package mcp

import (
	"context"
	"encoding/json"
	"testing"

	"github.com/stretchr/testify/require"

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

func TestToolMuxRejectsDuplicateAdvertisedOwner(t *testing.T) {
	message := &toolMuxTestHandler{tools: toolMuxTestDefs("same")}
	execution := &toolMuxNoAckHandler{defs: toolMuxTestDefs("same")}
	_, err := newToolMux(
		toolChild{kind: toolChildMessage, inventory: message},
		toolChild{kind: toolChildMemory, inventory: execution},
	)
	require.EqualError(t, err, `tool mux: duplicate owner for "same"`)
}

func TestToolMuxRoutesEveryAdvertisedToolToExactlyOneOwner(t *testing.T) {
	message := &toolMuxTestHandler{tools: toolMuxTestDefs("speak")}
	execution := &toolMuxNoAckHandler{defs: toolMuxTestDefs("run_review")}
	kb := &toolMuxNoAckHandler{defs: toolMuxTestDefs("lookup_fact")}
	mux, err := newToolMux(
		toolChild{kind: toolChildMessage, inventory: message},
		toolChild{kind: toolChildMemory, inventory: execution},
		toolChild{kind: toolChildKB, inventory: kb},
	)
	require.NoError(t, err)

	for _, name := range []string{"speak", "run_review", "lookup_fact"} {
		result := mux.CallTool(context.Background(), "agent-a", name, nil)
		require.False(t, result.IsError)
		require.Equal(t, name, result.Content[0].Text)
	}
	require.Equal(t, []string{"speak"}, message.calls)
}

func TestToolMuxSnapshotsToolInventory(t *testing.T) {
	message := &toolMuxTestHandler{tools: toolMuxTestDefs("message_public")}
	mux, err := newToolMux(toolChild{kind: toolChildMessage, inventory: message})
	require.NoError(t, err)

	message.tools[0].Name = "mutated"
	require.Equal(t, []string{"message_public"}, toolMuxToolNames(mux.Tools()))
}

func TestToolMuxDoesNotClassifyToolNamesOrPrefixes(t *testing.T) {
	message := &toolMuxTestHandler{tools: toolMuxTestDefs("say")}
	execution := &toolMuxNoAckHandler{defs: toolMuxTestDefs("review_without_execution_prefix")}
	kb := &toolMuxNoAckHandler{defs: toolMuxTestDefs("knowledge_lookup_not_in_old_list")}
	mux, err := newToolMux(
		toolChild{kind: toolChildMessage, inventory: message},
		toolChild{kind: toolChildMemory, inventory: execution},
		toolChild{kind: toolChildKB, inventory: kb},
	)
	require.NoError(t, err)
	require.False(t, mux.CallTool(context.Background(), "agent", "review_without_execution_prefix", nil).IsError)
	require.False(t, mux.CallTool(context.Background(), "agent", "knowledge_lookup_not_in_old_list", nil).IsError)
}

func TestToolMuxForwardsFallbackDeliveryToMessageChild(t *testing.T) {
	message := &toolMuxTestHandler{tools: toolMuxTestDefs("say")}
	mux, err := newToolMux(toolChild{kind: toolChildMessage, inventory: message})
	require.NoError(t, err)

	mux.AcknowledgeMCPToolResultDelivery(
		"agent-a",
		"runtime-a",
		[]string{"one", "two"},
		types.MCPToolDeliveryUncertain,
	)
	require.Equal(t, "agent-a", message.ackAgentID)
	require.Equal(t, "runtime-a", message.ackRuntimeSessionID)
	require.Equal(t, []string{"one", "two"}, message.ackReceipts)
	require.Equal(t, types.MCPToolDeliveryUncertain, message.ackOutcome)
	require.Equal(t, 1, message.ackCalls)
}

func TestToolMuxRejectsMissingMessageChild(t *testing.T) {
	_, err := newToolMux(toolChild{
		kind:      toolChildMemory,
		inventory: &toolMuxTestHandler{tools: toolMuxTestDefs("run")},
	})
	require.EqualError(t, err, "tool mux: exactly one message child is required")
}

func TestToolMuxRejectsMessageChildWithoutAcknowledger(t *testing.T) {
	_, err := newToolMux(toolChild{
		kind:      toolChildMessage,
		inventory: toolMuxNoAckHandler{defs: toolMuxTestDefs("say")},
	})
	require.EqualError(t, err, "tool mux: message child must acknowledge fallback delivery")
}

func TestToolMuxRejectsSecondMessageChild(t *testing.T) {
	_, err := newToolMux(
		toolChild{kind: toolChildMessage, inventory: &toolMuxTestHandler{tools: toolMuxTestDefs("one")}},
		toolChild{kind: toolChildMessage, inventory: &toolMuxTestHandler{tools: toolMuxTestDefs("two")}},
	)
	require.EqualError(t, err, "tool mux: exactly one message child is required")
}

func TestToolMuxRejectsDuplicateFallbackAcknowledger(t *testing.T) {
	_, err := newToolMux(
		toolChild{kind: toolChildMessage, inventory: &toolMuxTestHandler{tools: toolMuxTestDefs("one")}},
		toolChild{kind: toolChildMemory, inventory: &toolMuxTestHandler{tools: toolMuxTestDefs("two")}},
	)
	require.EqualError(t, err, "tool mux: multiple fallback delivery acknowledgers")
}

type toolMuxTestHandler struct {
	tools               []types.MCPToolDef
	calls               []string
	ackAgentID          string
	ackRuntimeSessionID string
	ackReceipts         []string
	ackOutcome          types.MCPToolDeliveryOutcome
	ackCalls            int
}

func (h *toolMuxTestHandler) Tools() []types.MCPToolDef { return h.tools }
func (h *toolMuxTestHandler) CallTool(_ context.Context, _ string, name string, _ json.RawMessage) *types.MCPToolResult {
	h.calls = append(h.calls, name)
	return toolText(name)
}
func (*toolMuxTestHandler) SkillDoc(types.ToolNamer) string { return "" }
func (*toolMuxTestHandler) SkillDocSplit(types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{}
}
func (h *toolMuxTestHandler) AcknowledgeMCPToolResultDelivery(agentID string, runtimeSessionID string, receipts []string, outcome types.MCPToolDeliveryOutcome) {
	h.ackAgentID = agentID
	h.ackRuntimeSessionID = runtimeSessionID
	h.ackReceipts = append([]string(nil), receipts...)
	h.ackOutcome = outcome
	h.ackCalls++
}

type toolMuxNoAckHandler struct {
	defs []types.MCPToolDef
}

func (h toolMuxNoAckHandler) Tools() []types.MCPToolDef { return h.defs }
func (h toolMuxNoAckHandler) CallTool(_ context.Context, _ string, name string, _ json.RawMessage) *types.MCPToolResult {
	return toolText(name)
}
func (toolMuxNoAckHandler) SkillDoc(types.ToolNamer) string { return "" }
func (toolMuxNoAckHandler) SkillDocSplit(types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{}
}

func toolMuxTestDefs(names ...string) []types.MCPToolDef {
	defs := make([]types.MCPToolDef, 0, len(names))
	for _, name := range names {
		defs = append(defs, types.MCPToolDef{Name: name})
	}
	return defs
}

func toolMuxToolNames(defs []types.MCPToolDef) []string {
	names := make([]string, 0, len(defs))
	for _, def := range defs {
		names = append(names, def.Name)
	}
	return names
}