package mcp

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

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

type mcpAgentToolHandler func(context.Context, string, json.RawMessage) *types.MCPToolResult

type mcpAPIToolHandler func(context.Context, *apiclient.Client, json.RawMessage) *types.MCPToolResult

type mcpToolRegistration struct {
	Def         types.MCPToolDef
	call        mcpAgentToolHandler
	callWithAPI func(context.Context, string, *apiclient.Client, json.RawMessage) *types.MCPToolResult
}

type mcpToolGroup struct {
	Name  string
	Tools []mcpToolRegistration
}

func mcpToolGroups(h *mcpHandler) []mcpToolGroup {
	return []mcpToolGroup{
		{Name: "message", Tools: messageMCPTools(h)},
	}
}

func mcpToolRegistrations(h *mcpHandler) []mcpToolRegistration {
	groups := mcpToolGroups(h)
	count := 0
	for _, group := range groups {
		count += len(group.Tools)
	}
	tools := make([]mcpToolRegistration, 0, count)
	for _, group := range groups {
		tools = append(tools, group.Tools...)
	}
	return tools
}

func mcpToolByName(h *mcpHandler, name string) (mcpToolRegistration, bool) {
	for _, tool := range mcpToolRegistrations(h) {
		if tool.Def.Name == name {
			return tool, true
		}
	}
	return mcpToolRegistration{}, false
}

func (h *mcpHandler) toolDefs() []types.MCPToolDef {
	registrations := mcpToolRegistrations(h)
	tools := make([]types.MCPToolDef, 0, len(registrations))
	for _, registration := range registrations {
		tools = append(tools, registration.Def)
	}

	// All AgentTeam MCP tools are session-bound communication tools and must
	// not be delegated as background tasks by the runtime.
	forbidden := &types.MCPToolExecution{TaskSupport: "forbidden"}
	for i := range tools {
		if tools[i].Execution == nil {
			tools[i].Execution = forbidden
		}
	}
	return tools
}

// MessageToolDefs returns the immutable MCP tool definitions for AgentTeam message
// tools without constructing a runtime-bound handler.
func MessageToolDefs() []types.MCPToolDef {
	return newMessageInventoryHandler().toolDefs()
}

func agentMCPTool(def types.MCPToolDef, fn mcpAgentToolHandler) mcpToolRegistration {
	return mcpToolRegistration{
		Def:  def,
		call: fn,
		callWithAPI: func(ctx context.Context, agentID string, _ *apiclient.Client, args json.RawMessage) *types.MCPToolResult {
			return fn(ctx, agentID, args)
		},
	}
}

func apiMCPTool(h *mcpHandler, def types.MCPToolDef, fn mcpAPIToolHandler) mcpToolRegistration {
	return mcpToolRegistration{
		Def: def,
		call: func(ctx context.Context, agentID string, args json.RawMessage) *types.MCPToolResult {
			return h.apiDispatch(ctx, agentID, args, fn)
		},
		callWithAPI: func(ctx context.Context, _ string, api *apiclient.Client, args json.RawMessage) *types.MCPToolResult {
			if api == nil {
				return toolError("agentteam API not available for this session")
			}
			return fn(ctx, api, args)
		},
	}
}

func (h *mcpHandler) CallTool(ctx context.Context, agentID, name string, args json.RawMessage) *types.MCPToolResult {
	tool, ok := mcpToolByName(h, name)
	if !ok {
		return toolError(fmt.Sprintf("unknown tool: %s", name))
	}
	return tool.call(ctx, agentID, args)
}