package mcp

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

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

type toolChildKind uint8

const (
	toolChildMessage toolChildKind = iota + 1
	toolChildMemory
	toolChildKB
)

type toolChild struct {
	kind      toolChildKind
	inventory types.MCPHandler
	forAgent  func(string) types.MCPHandler
}

type toolOwner struct {
	handler  types.MCPHandler
	forAgent func(string) types.MCPHandler
}

type toolMux struct {
	tools         []types.MCPToolDef
	owners        map[string]toolOwner
	skillHandlers []types.MCPHandler
	acknowledger  types.MCPToolResultDeliveryAcknowledger
}

func newToolMux(children ...toolChild) (*toolMux, error) {
	mux := &toolMux{owners: make(map[string]toolOwner)}
	messageCount := 0
	fallbackCount := 0
	for _, child := range children {
		if child.inventory == nil {
			continue
		}
		acknowledger, canAcknowledge := child.inventory.(types.MCPToolResultDeliveryAcknowledger)
		if canAcknowledge {
			fallbackCount++
		}
		if child.kind == toolChildMessage {
			messageCount++
			if !canAcknowledge {
				return nil, fmt.Errorf("tool mux: message child must acknowledge fallback delivery")
			}
			mux.acknowledger = acknowledger
		}

		for _, def := range cloneMCPToolDefs(child.inventory.Tools()) {
			name := strings.TrimSpace(def.Name)
			if name == "" {
				continue
			}
			if _, duplicate := mux.owners[name]; duplicate {
				return nil, fmt.Errorf("tool mux: duplicate owner for %q", name)
			}
			mux.owners[name] = toolOwner{handler: child.inventory, forAgent: child.forAgent}
			mux.tools = append(mux.tools, def)
		}
		mux.skillHandlers = append(mux.skillHandlers, child.inventory)
	}
	if messageCount != 1 {
		return nil, fmt.Errorf("tool mux: exactly one message child is required")
	}
	if fallbackCount > 1 {
		return nil, fmt.Errorf("tool mux: multiple fallback delivery acknowledgers")
	}
	return mux, nil
}

func (h *toolMux) Tools() []types.MCPToolDef {
	return cloneMCPToolDefs(h.tools)
}

func (h *toolMux) CallTool(ctx context.Context, agentID string, name string, args json.RawMessage) *types.MCPToolResult {
	name = strings.TrimSpace(name)
	owner, ok := h.owners[name]
	if !ok {
		return toolError("unknown tool: " + name)
	}
	handler := owner.handler
	if owner.forAgent != nil {
		handler = owner.forAgent(agentID)
	}
	if handler == nil {
		return toolError("unknown tool: " + name)
	}
	return handler.CallTool(ctx, agentID, name, args)
}

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

func (h *toolMux) SkillDocSplit(toolName types.ToolNamer) types.SkillDocParts {
	var result types.SkillDocParts
	for _, handler := range h.skillHandlers {
		parts := handler.SkillDocSplit(toolName)
		result.Core = joinSkillDocSection(result.Core, parts.Core)
		result.Reference = joinSkillDocSection(result.Reference, parts.Reference)
	}
	return result
}

func (h *toolMux) AcknowledgeMCPToolResultDelivery(agentID string, runtimeSessionID string, receipts []string, outcome types.MCPToolDeliveryOutcome) {
	h.acknowledger.AcknowledgeMCPToolResultDelivery(agentID, runtimeSessionID, receipts, outcome)
}

func joinSkillDocSection(left string, right string) string {
	return strings.TrimSpace(strings.Join([]string{strings.TrimSpace(left), strings.TrimSpace(right)}, "\n\n"))
}

func cloneMCPToolDefs(defs []types.MCPToolDef) []types.MCPToolDef {
	cloned := make([]types.MCPToolDef, len(defs))
	for index, def := range defs {
		cloned[index] = cloneMCPToolDef(def)
	}
	return cloned
}

func cloneMCPToolDef(def types.MCPToolDef) types.MCPToolDef {
	encoded, err := json.Marshal(def)
	if err != nil {
		return def
	}
	var cloned types.MCPToolDef
	if err := json.Unmarshal(encoded, &cloned); err != nil {
		return def
	}
	return cloned
}

// NewAgentTeamHandler composes the exact message, Memory, and KB tool owners into
// one immutable inventory and dispatch mux.
func NewAgentTeamHandler(
	message types.MCPHandler,
	memory MemorySearcher,
	kb types.MCPHandler,
) (types.MCPHandler, error) {
	children := []toolChild{{kind: toolChildMessage, inventory: message}}
	if memory != nil {
		children = append(children, toolChild{
			kind:      toolChildMemory,
			inventory: newMemoryHandler(memory, ""),
			forAgent: func(agentID string) types.MCPHandler {
				return newMemoryHandler(memory, agentID)
			},
		})
	}
	if kb != nil {
		children = append(children, toolChild{kind: toolChildKB, inventory: kb})
	}
	return newToolMux(children...)
}