package mcp

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

const (
	defaultKBSearchMaxResults = 10
	maxKBSearchMaxResults     = 50
)

// kbAgentTeamHandler implements ordinary-agent KB retrieval MCP tools. The model
// sees only search/fetch. AgentTeam injects the current turn message_id for
// attested scope resolution; source/path/ref/revision stay hidden behind
// server-issued citation IDs.
type kbAgentTeamHandler struct {
	c       AgentTeamConnection
	runtime KBRouteContext
}

// KBHandlerConfig contains the required remote connection and route context
// for ordinary-agent KB tools.
type KBHandlerConfig struct {
	Connection AgentTeamConnection
	Runtime    KBRouteContext
}

// KBRouteContext is the exact runtime surface consumed by ordinary-agent KB
// tools.
type KBRouteContext interface {
	Session(string) (types.SessionSnapshot, bool)
	CurrentTurnMessageID(string) string
}

// NewKBHandler constructs an ordinary-agent KB retrieval handler.
func NewKBHandler(config KBHandlerConfig) (types.MCPHandler, error) {
	if config.Connection == nil {
		return nil, errors.New("KB connection is required")
	}
	if config.Runtime == nil {
		return nil, errors.New("KB route context is required")
	}
	return &kbAgentTeamHandler{c: config.Connection, runtime: config.Runtime}, nil
}

func (h *kbAgentTeamHandler) Tools() []types.MCPToolDef { return kbAgentTeamToolDefs() }

func (h *kbAgentTeamHandler) CallTool(ctx context.Context, agentID, name string, args json.RawMessage) *types.MCPToolResult {
	name = strings.TrimSpace(name)
	switch name {
	case "kb_search", "kb_fetch":
		return h.callKBTool(ctx, agentID, name, args)
	default:
		return toolError("unknown tool: " + name)
	}
}

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

func (h *kbAgentTeamHandler) SkillDocSplit(toolName types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{
		Core: "Use `" + toolName("kb_search") + "` to search all knowledge bases available to the current turn, then `" + toolName("kb_fetch") + "` to fetch details for returned citation_ids. Do not browse KB structure; refine the query if search/fetch results are truncated or stale.",
	}
}

func (h *kbAgentTeamHandler) apiClient(agentID string) *apiclient.Client {
	if h.c == nil || h.runtime == nil {
		return nil
	}
	sess, ok := h.runtime.Session(agentID)
	if !ok {
		return nil
	}
	return h.c.APIClient(sess.ServerAgentID)
}

func (h *kbAgentTeamHandler) currentTurnMessageID(agentID string) string {
	if h.runtime == nil {
		return ""
	}
	return strings.TrimSpace(h.runtime.CurrentTurnMessageID(agentID))
}

// callKBTool resolves the op, builds the /agent-thread/upstream payload, and
// POSTs it via the agent's API client.
func (h *kbAgentTeamHandler) callKBTool(ctx context.Context, agentID, toolName string, args json.RawMessage) *types.MCPToolResult {
	api := h.apiClient(agentID)
	if api == nil {
		return toolError("agentteam API not available for this session")
	}
	api.LocalAgentID = strings.TrimSpace(agentID)

	op, err := kbOpForTool(toolName)
	if err != nil {
		return toolError(err.Error())
	}
	body, err := decodeKBAgentTeamInput(toolName, args)
	if err != nil {
		return toolError(err.Error())
	}

	messageID := h.currentTurnMessageID(agentID)
	if messageID == "" {
		return toolError("current turn message_id is required for KB retrieval")
	}

	payload := kbUpstreamEnvelope{
		MessageID: messageID,
		Method:    "POST",
		Path:      "/api/v1/kb/retrieval/knowledge/" + op,
		Body:      body,
	}

	path := api.AgentRuntimePath("/agent-thread/upstream")
	var response json.RawMessage
	if err := api.Post(ctx, path, payload, &response); err != nil {
		return knowledgeMCPToolError(err)
	}

	text, structured, err := kbModelResult(response)
	if err != nil {
		return toolError(err.Error())
	}
	return types.NewToolTextStructured(text, structured)
}

func kbOpForTool(toolName string) (string, error) {
	switch strings.TrimSpace(toolName) {
	case "kb_search":
		return "search", nil
	case "kb_fetch":
		return "fetch", nil
	default:
		return "", errUnknownKBTool
	}
}

func decodeKBAgentTeamInput(toolName string, args json.RawMessage) (json.RawMessage, error) {
	switch strings.TrimSpace(toolName) {
	case "kb_search":
		var input kbSearchInput
		if err := json.Unmarshal(normalizeKBInputArgs(args), &input); err != nil {
			return nil, stringError("invalid tool arguments")
		}
		input, err := input.normalized()
		if err != nil {
			return nil, err
		}
		return json.Marshal(input)
	case "kb_fetch":
		var input kbFetchInput
		if err := json.Unmarshal(normalizeKBInputArgs(args), &input); err != nil {
			return nil, stringError("invalid tool arguments")
		}
		input, err := input.normalized()
		if err != nil {
			return nil, err
		}
		return json.Marshal(input)
	default:
		return nil, errUnknownKBTool
	}
}

func normalizeKBInputArgs(args json.RawMessage) json.RawMessage {
	if len(args) == 0 || string(args) == "null" {
		return json.RawMessage(`{}`)
	}
	return args
}

func normalizeKBAgentTeamMaxResults(n int) int {
	if n <= 0 {
		return defaultKBSearchMaxResults
	}
	if n > maxKBSearchMaxResults {
		return maxKBSearchMaxResults
	}
	return n
}

var errUnknownKBTool = stringError("unknown KB tool")

type stringError string

func (e stringError) Error() string { return string(e) }

// ── Tool definitions ─────────────────────────────────────────────────────────

func kbAgentTeamToolDefs() []types.MCPToolDef {
	readOnly := &types.MCPToolAnnotation{ReadOnlyHint: true}
	forbidden := &types.MCPToolExecution{TaskSupport: "forbidden"}
	return []types.MCPToolDef{
		{
			Name:        "kb_search",
			Description: "Search accessible KBs for a natural-language query; returns citation_ids. max_results defaults to 10 and is capped at 50.",
			Annotations: readOnly,
			Execution:   forbidden,
			InputSchema: types.MCPInputSchema{
				Type: "object",
				Properties: map[string]types.MCPProperty{
					"query":       {Type: "string", Description: "Natural-language search query."},
					"max_results": {Type: "number", Description: "Maximum citations to return; default 10, hard max 50."},
				},
				Required: []string{"query"},
			},
		},
		{
			Name:        "kb_fetch",
			Description: "Fetch details for citation_ids returned by a knowledge-base search. Re-search/refine if a citation is stale or content is truncated.",
			Annotations: readOnly,
			Execution:   forbidden,
			InputSchema: types.MCPInputSchema{
				Type: "object",
				Properties: map[string]types.MCPProperty{
					"citation_ids":       {Type: "array", Description: "Citation IDs returned by a knowledge-base search.", Items: &types.MCPProperty{Type: "string"}},
					"max_chars_per_item": {Type: "number", Description: "Optional per-citation content budget."},
				},
				Required: []string{"citation_ids"},
			},
		},
	}
}