//go:build !tinygo && !wasm

package starlark

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	codemcp "github.com/maximhq/bifrost/core/mcp"
	"github.com/maximhq/bifrost/core/schemas"
)

// createGetToolDocsTool creates the getToolDocs tool definition for code mode.
// This tool provides detailed documentation for a specific tool when the compact
// signatures from readToolFile are not sufficient to understand how to use it.
func (s *StarlarkCodeMode) createGetToolDocsTool() schemas.ChatTool {
	getToolDocsProps := schemas.NewOrderedMapFromPairs(
		schemas.KV("server", map[string]interface{}{
			"type":        "string",
			"description": "The server name (e.g., 'calculator'). Use listToolFiles to see available servers.",
		}),
		schemas.KV("tool", map[string]interface{}{
			"type":        "string",
			"description": "The tool name (e.g., 'add'). Use readToolFile to see available tools for a server.",
		}),
	)
	return schemas.ChatTool{
		Type: schemas.ChatToolTypeFunction,
		Function: &schemas.ChatToolFunction{
			Name: codemcp.ToolTypeGetToolDocs,
			Description: schemas.Ptr(
				"Get detailed documentation for a specific tool including full parameter descriptions, " +
					"types, and usage examples. Use this when the compact signature from readToolFile " +
					"is not sufficient to understand how to use a tool. " +
					"Requires both server name and tool name as parameters.",
			),
			Parameters: &schemas.ToolFunctionParameters{
				Type:       "object",
				Properties: getToolDocsProps,
				Required:   []string{"server", "tool"},
			},
		},
	}
}

// handleGetToolDocs handles the getToolDocs tool call.
func (s *StarlarkCodeMode) handleGetToolDocs(ctx context.Context, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, error) {
	// Parse tool arguments
	var arguments map[string]interface{}
	if err := json.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil {
		return nil, fmt.Errorf("failed to parse tool arguments: %v", err)
	}

	serverName, ok := arguments["server"].(string)
	if !ok || serverName == "" {
		return nil, fmt.Errorf("server parameter is required and must be a string")
	}

	toolName, ok := arguments["tool"].(string)
	if !ok || toolName == "" {
		return nil, fmt.Errorf("tool parameter is required and must be a string")
	}

	// Get available tools per client
	availableToolsPerClient := s.clientManager.GetToolPerClient(ctx)

	// Find matching client
	var matchedClientName string
	var matchedTool *schemas.ChatTool

	serverNameLower := strings.ToLower(serverName)
	for clientName, tools := range availableToolsPerClient {
		client := s.clientManager.GetClientByName(clientName)
		if client == nil {
			s.logger.Warn("%s Client %s not found, skipping", codemcp.CodeModeLogPrefix, clientName)
			continue
		}
		if !client.ExecutionConfig.IsCodeModeClient || len(tools) == 0 {
			continue
		}

		clientNameLower := strings.ToLower(clientName)
		if clientNameLower == serverNameLower {
			matchedClientName = clientName

			// Find the specific tool
			for i, tool := range tools {
				if tool.Function != nil {
					if matchesToolReference(toolName, clientName, tool.Function.Name) {
						matchedTool = &tools[i]
						break
					}
				}
			}
			break
		}
	}

	// Handle server not found
	if matchedClientName == "" {
		var availableServers []string
		for name := range availableToolsPerClient {
			client := s.clientManager.GetClientByName(name)
			if client != nil && client.ExecutionConfig.IsCodeModeClient {
				availableServers = append(availableServers, name)
			}
		}
		errorMsg := fmt.Sprintf("Server '%s' not found. Available servers are:\n", serverName)
		for _, sn := range availableServers {
			errorMsg += fmt.Sprintf("  - %s\n", sn)
		}
		return createToolResponseMessage(toolCall, errorMsg), nil
	}

	// Handle tool not found
	if matchedTool == nil {
		tools := availableToolsPerClient[matchedClientName]
		var availableTools []string
		for _, tool := range tools {
			if tool.Function != nil {
				availableTools = append(availableTools, getCanonicalToolName(matchedClientName, tool.Function.Name))
			}
		}
		errorMsg := fmt.Sprintf("Tool '%s' not found in server '%s'. Available tools are:\n", toolName, matchedClientName)
		for _, t := range availableTools {
			errorMsg += fmt.Sprintf("  - %s\n", t)
		}
		return createToolResponseMessage(toolCall, errorMsg), nil
	}

	// Generate detailed documentation using generateTypeDefinitions
	docContent := generateTypeDefinitions(matchedClientName, []schemas.ChatTool{*matchedTool}, true)

	return createToolResponseMessage(toolCall, docContent), nil
}

// generateTypeDefinitions generates Python documentation with docstrings from ChatTool schemas.
func generateTypeDefinitions(clientName string, tools []schemas.ChatTool, isToolLevel bool) string {
	var sb strings.Builder

	// Write comprehensive header
	sb.WriteString("# ============================================================================\n")
	if isToolLevel && len(tools) == 1 && tools[0].Function != nil {
		sb.WriteString(fmt.Sprintf("# Documentation for %s.%s tool\n", clientName, getCanonicalToolName(clientName, tools[0].Function.Name)))
	} else {
		sb.WriteString(fmt.Sprintf("# Documentation for %s MCP server\n", clientName))
	}
	sb.WriteString("# ============================================================================\n")
	sb.WriteString("#\n")
	if isToolLevel && len(tools) == 1 {
		sb.WriteString("# This file contains Python documentation for a specific tool on this MCP server.\n")
	} else {
		sb.WriteString("# This file contains Python documentation for all tools available on this MCP server.\n")
	}
	sb.WriteString("#\n")
	sb.WriteString("# USAGE INSTRUCTIONS:\n")
	sb.WriteString(fmt.Sprintf("# Call tools using: result = %s.tool_name(param=value)\n", clientName))
	sb.WriteString("# No async/await needed - calls are synchronous.\n")
	sb.WriteString("#\n")
	sb.WriteString("# STARLARK DIFFERENCE FROM PYTHON:\n")
	sb.WriteString("# for/if/while at top level MUST be inside a function.\n")
	sb.WriteString("# Wrap loops: def main(): for x in items: ... then result = main()\n")
	sb.WriteString("#\n")
	sb.WriteString("# CRITICAL - HANDLING RESPONSES:\n")
	sb.WriteString("# Tool responses are dicts. To avoid runtime errors:\n")
	sb.WriteString("# 1. Use print(result) to inspect the response structure first\n")
	sb.WriteString("# 2. Access dict values with brackets: result[\"key\"] NOT result.key\n")
	sb.WriteString("# 3. Use .get() for safe access: result.get(\"key\", default)\n")
	sb.WriteString("#\n")
	sb.WriteString("# Common error: \"key not found\" or \"has no attribute\"\n")
	sb.WriteString("# Fix: Use print() to see actual structure, then use result[\"key\"] or .get()\n")
	sb.WriteString("# ============================================================================\n\n")

	// Generate function definitions for each tool
	for _, tool := range tools {
		if tool.Function == nil || tool.Function.Name == "" {
			continue
		}

		originalToolName := tool.Function.Name
		toolName := getCanonicalToolName(clientName, originalToolName)
		description := ""
		if tool.Function.Description != nil {
			description = *tool.Function.Description
		}

		// Generate function signature
		params := formatPythonParams(tool.Function.Parameters)
		sb.WriteString(fmt.Sprintf("def %s(%s) -> dict:\n", toolName, params))

		// Generate docstring
		sb.WriteString("    \"\"\"\n")
		if description != "" {
			sb.WriteString(fmt.Sprintf("    %s\n", description))
			sb.WriteString("\n")
		}

		// Args section
		if tool.Function.Parameters != nil && tool.Function.Parameters.Properties != nil {
			props := tool.Function.Parameters.Properties
			required := make(map[string]bool)
			if tool.Function.Parameters.Required != nil {
				for _, req := range tool.Function.Parameters.Required {
					required[req] = true
				}
			}

			if props.Len() > 0 {
				sb.WriteString("    Args:\n")

				// Sort properties for consistent output
				propNames := make([]string, 0, props.Len())
				props.Range(func(name string, _ interface{}) bool {
					propNames = append(propNames, name)
					return true
				})
				for i := 0; i < len(propNames)-1; i++ {
					for j := i + 1; j < len(propNames); j++ {
						if propNames[i] > propNames[j] {
							propNames[i], propNames[j] = propNames[j], propNames[i]
						}
					}
				}

				for _, propName := range propNames {
					prop, _ := props.Get(propName)
					propMap, ok := prop.(map[string]interface{})
					if !ok {
						continue
					}

					pyType := jsonSchemaToPython(propMap)
					propDesc := ""
					if desc, ok := propMap["description"].(string); ok && desc != "" {
						propDesc = desc
					} else {
						propDesc = fmt.Sprintf("%s parameter", propName)
					}

					requiredNote := ""
					if required[propName] {
						requiredNote = " (required)"
					} else {
						requiredNote = " (optional)"
					}

					sb.WriteString(fmt.Sprintf("        %s (%s): %s%s\n", propName, pyType, propDesc, requiredNote))
				}
				sb.WriteString("\n")
			}
		}

		// Returns section
		sb.WriteString("    Returns:\n")
		sb.WriteString("        dict: Response from the tool. Structure varies by tool.\n")
		sb.WriteString("              Use print(result) to inspect the actual structure.\n")
		sb.WriteString("\n")

		// Example section
		sb.WriteString("    Example:\n")
		sb.WriteString(fmt.Sprintf("        result = %s.%s(%s)\n", clientName, toolName, getExampleParams(tool.Function.Parameters)))
		sb.WriteString("        print(result)  # Always inspect response first!\n")
		sb.WriteString("        value = result.get(\"key\", default)  # Safe access\n")
		sb.WriteString("    \"\"\"\n")
		sb.WriteString("    ...\n\n")
	}

	return sb.String()
}

// getExampleParams generates example parameter usage for a function.
func getExampleParams(params *schemas.ToolFunctionParameters) string {
	if params == nil || params.Properties == nil || params.Properties.Len() == 0 {
		return ""
	}

	required := make(map[string]bool)
	if params.Required != nil {
		for _, req := range params.Required {
			required[req] = true
		}
	}

	keys := params.Properties.Keys()

	// Get first required param as example
	for _, name := range keys {
		if required[name] {
			return fmt.Sprintf("%s=\"...\"", name)
		}
	}

	// If no required, get first param
	if len(keys) > 0 {
		return fmt.Sprintf("%s=\"...\"", keys[0])
	}

	return ""
}