package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"log/slog"

	"github.com/tidbcloud/AgentTeam/desktop/internal/memorysearch"
	"github.com/tidbcloud/AgentTeam/desktop/internal/promptcompiler"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// MemorySearcher is the read-only Memory search port consumed by the AgentTeam
// handler.
type MemorySearcher interface {
	MemorySearchAvailable(ctx context.Context, agentID string) bool
	MemorySearch(ctx context.Context, agentID string, request memorysearch.SearchRequest) (memorysearch.SearchResponse, error)
}

type memoryMCPHandler struct {
	service MemorySearcher
	agentID string
}

func newMemoryHandler(service MemorySearcher, agentID string) types.MCPHandler {
	return &memoryMCPHandler{service: service, agentID: agentID}
}

func (h *memoryMCPHandler) Tools() []types.MCPToolDef {
	return []types.MCPToolDef{memorySearchToolDef()}
}

func (h *memoryMCPHandler) CallTool(
	ctx context.Context,
	_ string,
	name string,
	args json.RawMessage,
) *types.MCPToolResult {
	if name != "memory_search" {
		return toolError("unknown tool: " + name)
	}
	var request struct {
		Query           string   `json:"query"`
		Limit           int      `json:"limit"`
		PathPrefix      string   `json:"pathPrefix"`
		Kinds           []string `json:"kinds"`
		IncludeArchived bool     `json:"includeArchived"`
		MaxChars        int      `json:"maxChars"`
	}
	if len(args) == 0 || string(args) == "null" {
		args = []byte(`{}`)
	}
	decoder := json.NewDecoder(bytes.NewReader(args))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&request); err != nil {
		return toolError("invalid tool arguments")
	}
	if h.service == nil || !h.service.MemorySearchAvailable(ctx, h.agentID) {
		return toolError("memory search index is unavailable; use bounded grep")
	}
	response, err := h.service.MemorySearch(ctx, h.agentID, memorysearch.SearchRequest{
		Query:           request.Query,
		Limit:           request.Limit,
		PathPrefix:      request.PathPrefix,
		Kinds:           request.Kinds,
		IncludeArchived: request.IncludeArchived,
		MaxChars:        request.MaxChars,
	})
	if err != nil {
		if errors.Is(err, memorysearch.ErrInvalidSearchRequest) {
			return toolError(err.Error())
		}
		slog.WarnContext(ctx, "memory MCP search failed",
			"agent_id", h.agentID, "error", err)
		return toolError("memory search is temporarily unavailable")
	}
	if response.State == memorysearch.StateDegraded {
		return toolError("memory search index is unavailable; use bounded grep")
	}
	return memorySearchResult(response)
}

type memorySearchAgentPayload struct {
	Results   []memorySearchAgentResult `json:"results"`
	Stale     bool                      `json:"stale,omitempty"`
	Truncated bool                      `json:"truncated,omitempty"`
}

type memorySearchAgentResult struct {
	Path    string `json:"path"`
	Lines   [2]int `json:"lines"`
	Snippet string `json:"snippet"`
}

// memorySearchResult renders the minimum agent-actionable response as one
// compact text payload. It deliberately omits structuredContent so runtimes do
// not receive a duplicate copy of the same result.
func memorySearchResult(response memorysearch.SearchResponse) *types.MCPToolResult {
	payload := memorySearchAgentPayload{
		Results:   make([]memorySearchAgentResult, 0, len(response.Results)),
		Stale:     response.Stale,
		Truncated: response.Truncated,
	}
	for _, result := range response.Results {
		payload.Results = append(payload.Results, memorySearchAgentResult{
			Path:    result.Path,
			Lines:   [2]int{result.ReadLines.Start, result.ReadLines.End},
			Snippet: result.Snippet,
		})
	}
	text, err := json.Marshal(payload)
	if err != nil {
		return toolError("memory search result could not be encoded")
	}
	return types.NewToolText(string(text))
}

func (h *memoryMCPHandler) SkillDoc(toolName types.ToolNamer) string {
	return types.JoinSkillDocParts(h.SkillDocSplit(toolName))
}

func (h *memoryMCPHandler) SkillDocSplit(toolName types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{
		Core: promptcompiler.MemorySearchGuidance(toolName("memory_search")),
	}
}

func memorySearchToolDef() types.MCPToolDef {
	return types.MCPToolDef{
		Name:        "memory_search",
		Description: "Find durable Memory by semantic terms and return compact file locators for bounded reads. Search one focused topic; use exact identifiers or stable phrases when known.",
		Annotations: &types.MCPToolAnnotation{ReadOnlyHint: true},
		InputSchema: types.MCPInputSchema{
			Type: "object",
			Properties: map[string]types.MCPProperty{
				"query":           {Type: "string", Description: "One focused topic using 2-6 discriminative terms; include exact IDs, symbols, paths, or stable phrases when known. Required, 2-512 Unicode characters."},
				"limit":           {Type: "integer", Description: "Maximum results; omit for the default 3, max 8."},
				"pathPrefix":      {Type: "string", Description: "Optional canonical Memory-relative path prefix."},
				"kinds":           {Type: "array", Description: "Optional source kind filters: system, profile, concept, workflow, event, reflection, root, or other.", Items: &types.MCPProperty{Type: "string"}},
				"includeArchived": {Type: "boolean", Description: "Whether archived Memory is included; defaults false."},
				"maxChars":        {Type: "integer", Description: "Whole response snippet budget; omit for the default 1600, max 8000."},
			},
			Required: []string{"query"},
		},
	}
}