package tools

import (
	"maps"
	"pentagi/pkg/database"

	"github.com/invopop/jsonschema"
	"github.com/vxcontrol/langchaingo/llms"
)

const (
	FinalyToolName             = "done"
	AskUserToolName            = "ask"
	MaintenanceToolName        = "maintenance"
	MaintenanceResultToolName  = "maintenance_result"
	CoderToolName              = "coder"
	CodeResultToolName         = "code_result"
	PentesterToolName          = "pentester"
	HackResultToolName         = "hack_result"
	AdviceToolName             = "advice"
	MemoristToolName           = "memorist"
	MemoristResultToolName     = "memorist_result"
	BrowserToolName            = "browser"
	GoogleToolName             = "google"
	DuckDuckGoToolName         = "duckduckgo"
	TavilyToolName             = "tavily"
	TraversaalToolName         = "traversaal"
	PerplexityToolName         = "perplexity"
	SearxngToolName            = "searxng"
	SploitusToolName           = "sploitus"
	SearchToolName             = "search"
	SearchResultToolName       = "search_result"
	EnricherResultToolName     = "enricher_result"
	SearchInMemoryToolName     = "search_in_memory"
	SearchGuideToolName        = "search_guide"
	StoreGuideToolName         = "store_guide"
	SearchAnswerToolName       = "search_answer"
	StoreAnswerToolName        = "store_answer"
	SearchCodeToolName         = "search_code"
	StoreCodeToolName          = "store_code"
	GraphitiSearchToolName     = "graphiti_search"
	ReportResultToolName       = "report_result"
	SubtaskListToolName        = "subtask_list"
	SubtaskPatchToolName       = "subtask_patch"
	TerminalToolName           = "terminal"
	FileToolName               = "file"
	GetFlowStatusToolName      = "get_flow_status"
	StopFlowToolName           = "stop_flow"
	SubmitFlowInputToolName    = "submit_flow_input"
	PatchFlowSubtasksToolName  = "patch_flow_subtasks"
	WaitFlowCompletionToolName = "wait_flow_completion"
)

type ToolType int

const (
	NoneToolType ToolType = iota
	EnvironmentToolType
	SearchNetworkToolType
	SearchVectorDbToolType
	AgentToolType
	StoreAgentResultToolType
	StoreVectorDbToolType
	BarrierToolType
)

func (t ToolType) String() string {
	switch t {
	case EnvironmentToolType:
		return "environment"
	case SearchNetworkToolType:
		return "search_network"
	case SearchVectorDbToolType:
		return "search_vector_db"
	case AgentToolType:
		return "agent"
	case StoreAgentResultToolType:
		return "store_agent_result"
	case StoreVectorDbToolType:
		return "store_vector_db"
	case BarrierToolType:
		return "barrier"
	default:
		return "none"
	}
}

// GetToolType returns the tool type for a given tool name
func GetToolType(name string) ToolType {
	if toolType, ok := toolsTypeMapping[name]; ok {
		return toolType
	}
	return NoneToolType
}

var toolsTypeMapping = map[string]ToolType{
	FinalyToolName:             BarrierToolType,
	AskUserToolName:            BarrierToolType,
	MaintenanceToolName:        AgentToolType,
	MaintenanceResultToolName:  StoreAgentResultToolType,
	CoderToolName:              AgentToolType,
	CodeResultToolName:         StoreAgentResultToolType,
	PentesterToolName:          AgentToolType,
	HackResultToolName:         StoreAgentResultToolType,
	AdviceToolName:             AgentToolType,
	MemoristToolName:           AgentToolType,
	MemoristResultToolName:     StoreAgentResultToolType,
	BrowserToolName:            SearchNetworkToolType,
	GoogleToolName:             SearchNetworkToolType,
	DuckDuckGoToolName:         SearchNetworkToolType,
	TavilyToolName:             SearchNetworkToolType,
	TraversaalToolName:         SearchNetworkToolType,
	PerplexityToolName:         SearchNetworkToolType,
	SearxngToolName:            SearchNetworkToolType,
	SploitusToolName:           SearchNetworkToolType,
	SearchToolName:             AgentToolType,
	SearchResultToolName:       StoreAgentResultToolType,
	EnricherResultToolName:     StoreAgentResultToolType,
	SearchInMemoryToolName:     SearchVectorDbToolType,
	SearchGuideToolName:        SearchVectorDbToolType,
	StoreGuideToolName:         StoreVectorDbToolType,
	SearchAnswerToolName:       SearchVectorDbToolType,
	StoreAnswerToolName:        StoreVectorDbToolType,
	SearchCodeToolName:         SearchVectorDbToolType,
	StoreCodeToolName:          StoreVectorDbToolType,
	GraphitiSearchToolName:     SearchVectorDbToolType,
	ReportResultToolName:       StoreAgentResultToolType,
	SubtaskListToolName:        StoreAgentResultToolType,
	SubtaskPatchToolName:       StoreAgentResultToolType,
	TerminalToolName:           EnvironmentToolType,
	FileToolName:               EnvironmentToolType,
	GetFlowStatusToolName:      EnvironmentToolType,
	StopFlowToolName:           EnvironmentToolType,
	SubmitFlowInputToolName:    EnvironmentToolType,
	PatchFlowSubtasksToolName:  EnvironmentToolType,
	WaitFlowCompletionToolName: EnvironmentToolType,
}

var reflector = &jsonschema.Reflector{
	DoNotReference: true,
	ExpandedStruct: true,
}

var allowedSummarizingToolsResult = []string{
	TerminalToolName,
	BrowserToolName,
}

var allowedStoringInMemoryTools = []string{
	TerminalToolName,
	FileToolName,
	SearchToolName,
	GoogleToolName,
	DuckDuckGoToolName,
	TavilyToolName,
	TraversaalToolName,
	PerplexityToolName,
	SearxngToolName,
	SploitusToolName,
	MaintenanceToolName,
	CoderToolName,
	PentesterToolName,
	AdviceToolName,
}

var registryDefinitions = map[string]llms.FunctionDefinition{
	TerminalToolName: {
		Name: TerminalToolName,
		Description: "Calls a terminal command in blocking mode. " +
			"Use timeout=0 or a negative value to apply the configured server default timeout. " +
			"Explicit positive values are accepted up to 10800 seconds (3 hours); values outside this range are replaced by the server default. " +
			"Only one command can be executed at a time",
		Parameters: reflector.Reflect(&TerminalAction{}),
	},
	FileToolName: {
		Name:        FileToolName,
		Description: "Modifies or reads local files",
		Parameters:  reflector.Reflect(&FileAction{}),
	},
	ReportResultToolName: {
		Name:        ReportResultToolName,
		Description: "Send the report result to the user with execution status and description",
		Parameters:  reflector.Reflect(&TaskResult{}),
	},
	SubtaskListToolName: {
		Name:        SubtaskListToolName,
		Description: "Send new generated subtask list to the user",
		Parameters:  reflector.Reflect(&SubtaskList{}),
	},
	SubtaskPatchToolName: {
		Name: SubtaskPatchToolName,
		Description: "Submit delta operations to modify the current subtask list instead of regenerating all subtasks. " +
			"Supports add (create new subtask at position), remove (delete by ID), modify (update title/description), " +
			"and reorder (move to different position) operations. Use empty operations array if no changes needed.",
		Parameters: reflector.Reflect(&SubtaskPatch{}),
	},
	SearchToolName: {
		Name: SearchToolName,
		Description: "Search in a different search engines in the internet and long-term memory " +
			"by your complex question to the researcher team member, also you can add some instructions to get result " +
			"in a specific format or structure or content type like " +
			"code or command samples, manuals, guides, exploits, vulnerability details, repositories, libraries, etc.",
		Parameters: reflector.Reflect(&ComplexSearch{}),
	},
	SearchResultToolName: {
		Name:        SearchResultToolName,
		Description: "Send the complex search result as a answer for the user question to the user",
		Parameters:  reflector.Reflect(&SearchResult{}),
	},
	BrowserToolName: {
		Name:        BrowserToolName,
		Description: "Opens a browser to look for additional information from the web site",
		Parameters:  reflector.Reflect(&Browser{}),
	},
	GoogleToolName: {
		Name: GoogleToolName,
		Description: "Search in the google search engine, it's a fast query and the shortest content " +
			"to check some information or collect public links by short query",
		Parameters: reflector.Reflect(&SearchAction{}),
	},
	DuckDuckGoToolName: {
		Name: DuckDuckGoToolName,
		Description: "Search in the duckduckgo search engine, it's a anonymous query and returns a small content " +
			"to check some information from different sources or collect public links by short query",
		Parameters: reflector.Reflect(&SearchAction{}),
	},
	TavilyToolName: {
		Name: TavilyToolName,
		Description: "Search in the tavily search engine, it's a more complex query and more detailed content " +
			"with answer by query and detailed information from the web sites",
		Parameters: reflector.Reflect(&SearchAction{}),
	},
	TraversaalToolName: {
		Name: TraversaalToolName,
		Description: "Search in the traversaal search engine, presents you answer and web-links " +
			"by your query according to relevant information from the web sites",
		Parameters: reflector.Reflect(&SearchAction{}),
	},
	PerplexityToolName: {
		Name: PerplexityToolName,
		Description: "Search in the perplexity search engine, it's a fully complex query and detailed research report " +
			"with answer by query and detailed information from the web sites and other sources augmented by the LLM",
		Parameters: reflector.Reflect(&SearchAction{}),
	},
	SearxngToolName: {
		Name: SearxngToolName,
		Description: "Search in the searxng meta search engine, it's a privacy-focused search engine " +
			"that aggregates results from multiple search engines with customizable categories, " +
			"language settings, and safety filters",
		Parameters: reflector.Reflect(&SearchAction{}),
	},
	SploitusToolName: {
		Name: SploitusToolName,
		Description: "Search the Sploitus exploit aggregator (https://sploitus.com) for public exploits, " +
			"proof-of-concept code, and offensive security tools. Sploitus indexes ExploitDB, Packet Storm, " +
			"GitHub Security Advisories, and many other sources. Use this tool to find exploit code and PoCs " +
			"for specific software, services, CVEs, or vulnerability classes (e.g. 'ssh', 'apache log4j', " +
			"'CVE-2021-44228'). Returns exploit URLs, CVSS scores, CVE references, and publication dates.",
		Parameters: reflector.Reflect(&SploitusAction{}),
	},
	EnricherResultToolName: {
		Name:        EnricherResultToolName,
		Description: "Send the enriched user's question with additional information to the user",
		Parameters:  reflector.Reflect(&EnricherResult{}),
	},
	SearchInMemoryToolName: {
		Name: SearchInMemoryToolName,
		Description: "Search in the vector database (long-term memory) for relevant information by providing one or more semantically rich, " +
			"context-aware natural language queries (1 to 5 queries). Formulate each query with sufficient context, intent, and detailed descriptions " +
			"to enhance semantic matching and retrieval accuracy. Multiple queries allow exploring different semantic angles and improve recall. " +
			"Results from all queries are merged, deduplicated, and ranked by relevance score. This function is ideal when you need to retrieve specific information " +
			"to assist in generating accurate and informative responses. If Task ID or Subtask ID are known, " +
			"they can be used as strict filters to further refine the search results and improve relevancy.",
		Parameters: reflector.Reflect(&SearchInMemoryAction{}),
	},
	SearchGuideToolName: {
		Name: SearchGuideToolName,
		Description: "Search in the vector database for relevant guides by providing one or more semantically rich, context-aware natural language queries (1 to 5 queries). " +
			"Formulate each query with sufficient context, intent, and detailed descriptions of the guides you need to enhance semantic matching and " +
			"retrieval accuracy. Multiple queries allow exploring different aspects of the guide topic and improve search coverage. " +
			"Specify the type of guide required to further refine the search. Results from all queries are merged, deduplicated, and ranked by relevance score. " +
			"This function is ideal when you need to retrieve specific guides to assist in accomplishing tasks or solving issues.",
		Parameters: reflector.Reflect(&SearchGuideAction{}),
	},
	StoreGuideToolName: {
		Name: StoreGuideToolName,
		Description: "Store the guide to the vector database for future use. " +
			"Anonymize all sensitive data (IPs, domains, credentials, paths) using descriptive placeholders",
		Parameters: reflector.Reflect(&StoreGuideAction{}),
	},
	SearchAnswerToolName: {
		Name: SearchAnswerToolName,
		Description: "Search in the vector database for relevant answers by providing one or more semantically rich, context-aware natural language queries (1 to 5 queries). " +
			"Formulate each query with sufficient context, intent, and detailed descriptions of what you want to find and why you need it " +
			"to enhance semantic matching and retrieval accuracy. Multiple queries allow exploring different formulations and improve search coverage. " +
			"Specify the type of answer required to further refine the search. Results from all queries are merged, deduplicated, and ranked by relevance score. " +
			"This function is ideal when you need to retrieve specific answers to assist in tasks, solve issues, or answer questions.",
		Parameters: reflector.Reflect(&SearchAnswerAction{}),
	},
	StoreAnswerToolName: {
		Name: StoreAnswerToolName,
		Description: "Store the question answer to the vector database for future use. " +
			"Anonymize all sensitive data (IPs, domains, credentials) using descriptive placeholders",
		Parameters: reflector.Reflect(&StoreAnswerAction{}),
	},
	SearchCodeToolName: {
		Name: SearchCodeToolName,
		Description: "Search in the vector database for relevant code samples by providing one or more semantically rich, context-aware natural language queries (1 to 5 queries). " +
			"Formulate each query with sufficient context, intent, and detailed descriptions of what you want to achieve with the code and what should be included, " +
			"to enhance semantic matching and retrieval accuracy. Multiple queries allow exploring different code patterns and use cases. " +
			"Specify the programming language to further refine the search. Results from all queries are merged, deduplicated, and ranked by relevance score. " +
			"This function is ideal when you need to retrieve specific code examples to assist in development tasks or solve programming issues.",
		Parameters: reflector.Reflect(&SearchCodeAction{}),
	},
	StoreCodeToolName: {
		Name: StoreCodeToolName,
		Description: "Store the code sample to the vector database for future use. It's should be a sample like a one source code file for some question. " +
			"Anonymize all sensitive data (IPs, domains, credentials, API keys) using descriptive placeholders",
		Parameters: reflector.Reflect(&StoreCodeAction{}),
	},
	GraphitiSearchToolName: {
		Name: GraphitiSearchToolName,
		Description: "Search the Graphiti temporal knowledge graph for historical penetration testing context, " +
			"including previous agent responses, tool execution records, discovered entities, and their relationships. " +
			"Supports 7 search types: temporal_window (time-bounded search), entity_relationships (graph traversal from an entity), " +
			"diverse_results (anti-redundancy search), episode_context (full agent reasoning and tool outputs), " +
			"successful_tools (proven techniques), recent_context (latest findings), and entity_by_label (type-specific entity search). " +
			"Use this to avoid repeating failed approaches, reuse successful exploitation techniques, understand entity relationships, " +
			"and build on previous findings within the same penetration testing engagement.",
		Parameters: reflector.Reflect(&GraphitiSearchAction{}),
	},
	MemoristToolName: {
		Name:        MemoristToolName,
		Description: "Call to Archivist team member who remember all the information about the past work and made tasks and can answer your question about it",
		Parameters:  reflector.Reflect(&MemoristAction{}),
	},
	MemoristResultToolName: {
		Name:        MemoristResultToolName,
		Description: "Send the search in long-term memory result as a answer for the user question to the user",
		Parameters:  reflector.Reflect(&MemoristResult{}),
	},
	MaintenanceToolName: {
		Name:        MaintenanceToolName,
		Description: "Call to DevOps team member to maintain local environment and tools inside the docker container",
		Parameters:  reflector.Reflect(&MaintenanceAction{}),
	},
	MaintenanceResultToolName: {
		Name:        MaintenanceResultToolName,
		Description: "Send the maintenance result to the user with task status and fully detailed report about using the result",
		Parameters:  reflector.Reflect(&TaskResult{}),
	},
	CoderToolName: {
		Name:        CoderToolName,
		Description: "Call to developer team member to write a code for the specific task",
		Parameters:  reflector.Reflect(&CoderAction{}),
	},
	CodeResultToolName: {
		Name:        CodeResultToolName,
		Description: "Send the code result to the user with execution status and fully detailed report about using the result",
		Parameters:  reflector.Reflect(&CodeResult{}),
	},
	PentesterToolName: {
		Name:        PentesterToolName,
		Description: "Call to pentester team member to perform a penetration test or looking for vulnerabilities and weaknesses",
		Parameters:  reflector.Reflect(&PentesterAction{}),
	},
	HackResultToolName: {
		Name:        HackResultToolName,
		Description: "Send the penetration test result to the user with detailed report",
		Parameters:  reflector.Reflect(&HackResult{}),
	},
	AdviceToolName: {
		Name:        AdviceToolName,
		Description: "Get more complex answer from the mentor about some issue or difficult situation",
		Parameters:  reflector.Reflect(&AskAdvice{}),
	},
	AskUserToolName: {
		Name:        AskUserToolName,
		Description: "If you need to ask user for input, use this tool",
		Parameters:  reflector.Reflect(&AskUser{}),
	},
	FinalyToolName: {
		Name:        FinalyToolName,
		Description: "If you need to finish the task with success or failure, use this tool",
		Parameters:  reflector.Reflect(&Done{}),
	},
	GetFlowStatusToolName: {
		Name: GetFlowStatusToolName,
		Description: "Return structured information about the automation flow. " +
			"Choose the detail level based on what you need to know:\n" +
			"  'summary'  — flow health snapshot: status, task counts, subtask counts, active task/subtask IDs. Best first call.\n" +
			"  'tasks'    — all tasks with ID, status, title. Add verbose=true for inputs and results.\n" +
			"  'subtasks' — all subtasks, optionally filtered by task_id. Add verbose=true for descriptions and results.\n" +
			"  'running'  — full Task→Subtask execution chain: task input, subtask description, partial result, and recent agent messages (10 default, 50 with verbose=true). Use when you need to understand what is happening right now.\n" +
			"  'planned'  — subtasks with status 'created' (not yet started), optionally filtered by task_id. Add verbose=true for full descriptions.",
		Parameters: reflector.Reflect(&GetFlowStatusAction{}),
	},
	StopFlowToolName: {
		Name: StopFlowToolName,
		Description: "Cancel the currently executing automation task. " +
			"If a task is actively running its execution context is cancelled; the flow transitions to 'waiting' state and can accept new input or subtask patches. " +
			"If no task is running the tool returns an informational message and takes no further action. " +
			"After calling this tool, verify the flow reached 'waiting' state before making further changes.",
		Parameters: reflector.Reflect(&StopFlowAction{}),
	},
	SubmitFlowInputToolName: {
		Name: SubmitFlowInputToolName,
		Description: "Deliver text to the automation flow. The behaviour depends on the current flow state:\n" +
			"  Subtask at 'ask' checkpoint — the text is delivered as the user's answer and the subtask resumes immediately.\n" +
			"  Flow waiting with no active subtask — the text becomes the goal for a new task; the generator and refiner agents decompose it into a subtask list. Write a complete, detailed description with all context the automation needs.\n" +
			"Returns an error if a task is currently running — cancel the running task first.",
		Parameters: reflector.Reflect(&SubmitFlowInputAction{}),
	},
	PatchFlowSubtasksToolName: {
		Name: PatchFlowSubtasksToolName,
		Description: "Replace the planned (not-yet-started) subtask list for a specific task using delta operations. " +
			"Supported operations: add (insert a new subtask at a position), remove (delete a planned subtask by ID), modify (update title or description), reorder (move to a different position). " +
			"The currently executing or waiting subtask can also be targeted by including its ID in an operation; doing so resets it to 'created' status. " +
			"After patching, submit new input to trigger execution of the updated plan. " +
			"Returns an error if a task is currently running — cancel it first.",
		Parameters: reflector.Reflect(&PatchFlowSubtasksAction{}),
	},
	WaitFlowCompletionToolName: {
		Name: WaitFlowCompletionToolName,
		Description: "Block until the currently running automation task completes (or the timeout expires). " +
			"Use this when you need to wait for the automation to finish before inspecting its results or taking further action. " +
			"Returns immediately with an informational message if no tasks exist or if no task is currently running — " +
			"use " + GetFlowStatusToolName + " first to confirm the flow state before calling this tool. " +
			"On return, always call " + GetFlowStatusToolName + " with detail='summary' to inspect the final status.",
		Parameters: reflector.Reflect(&WaitFlowCompletionAction{}),
	},
}

func getMessageType(name string) database.MsglogType {
	switch name {
	case TerminalToolName:
		return database.MsglogTypeTerminal
	case FileToolName:
		return database.MsglogTypeFile
	case BrowserToolName:
		return database.MsglogTypeBrowser
	case MemoristToolName, SearchToolName, GoogleToolName, DuckDuckGoToolName, TavilyToolName, TraversaalToolName,
		PerplexityToolName, SearxngToolName, SploitusToolName,
		SearchGuideToolName, SearchAnswerToolName, SearchCodeToolName, SearchInMemoryToolName, GraphitiSearchToolName:
		return database.MsglogTypeSearch
	case AdviceToolName:
		return database.MsglogTypeAdvice
	case AskUserToolName:
		return database.MsglogTypeAsk
	case FinalyToolName:
		return database.MsglogTypeDone
	default:
		return database.MsglogTypeThoughts
	}
}

func getMessageResultFormat(name string) database.MsglogResultFormat {
	switch name {
	case TerminalToolName:
		return database.MsglogResultFormatTerminal
	case FileToolName, BrowserToolName:
		return database.MsglogResultFormatPlain
	default:
		return database.MsglogResultFormatMarkdown
	}
}

// GetRegistryDefinitions returns tool definitions from the tools package
func GetRegistryDefinitions() map[string]llms.FunctionDefinition {
	registry := make(map[string]llms.FunctionDefinition, len(registryDefinitions))
	maps.Copy(registry, registryDefinitions)
	return registry
}

// GetToolTypeMapping returns a mapping from tool names to tool types
func GetToolTypeMapping() map[string]ToolType {
	mapping := make(map[string]ToolType, len(toolsTypeMapping))
	maps.Copy(mapping, toolsTypeMapping)
	return mapping
}

// GetToolsByType returns a mapping from tool types to a list of tool names
func GetToolsByType() map[ToolType][]string {
	result := make(map[ToolType][]string)

	for toolName, toolType := range toolsTypeMapping {
		result[toolType] = append(result[toolType], toolName)
	}

	return result
}
