package mcp

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

// --- kbOpForTool tests ---

func TestKBAgentTeamOpMapping(t *testing.T) {
	tests := []struct {
		name    string
		tool    string
		wantOp  string
		wantErr bool
	}{
		{name: "kb_search", tool: "kb_search", wantOp: "search"},
		{name: "kb_fetch", tool: "kb_fetch", wantOp: "fetch"},
		{name: "old list removed", tool: "kb_list", wantErr: true},
		{name: "old read removed", tool: "kb_read", wantErr: true},
		{name: "unknown tool", tool: "kb_unknown", wantErr: true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			op, err := kbOpForTool(tt.tool)
			if tt.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.Equal(t, tt.wantOp, op)
		})
	}
}

// --- Tool def tests ---

func TestKBAgentTeamToolDefs(t *testing.T) {
	defs := (&kbAgentTeamHandler{}).Tools()
	require.Len(t, defs, 2)

	names := make(map[string]struct{}, len(defs))
	for _, def := range defs {
		names[def.Name] = struct{}{}
		require.NotNil(t, def.Annotations, "%s should have annotations", def.Name)
		require.True(t, def.Annotations.ReadOnlyHint, "%s should be read-only", def.Name)
		require.NotNil(t, def.Execution, "%s should have execution metadata", def.Name)
		require.Equal(t, "forbidden", def.Execution.TaskSupport, "%s should forbid background tasks", def.Name)
	}
	for _, name := range []string{"kb_search", "kb_fetch"} {
		require.Contains(t, names, name, "missing tool %s", name)
	}

	byName := map[string]types.MCPToolDef{}
	for _, def := range defs {
		byName[def.Name] = def
	}
	require.Equal(t, []string{"query"}, byName["kb_search"].InputSchema.Required)
	require.Equal(t, []string{"citation_ids"}, byName["kb_fetch"].InputSchema.Required)
	forbiddenFields := []string{"target", "kb_ids", "source_ids", "source_id", "repo_id", "repo_role", "path", "ref", "revision", "search_mode", "offset", "citation_uri"}
	for _, def := range defs {
		for _, field := range forbiddenFields {
			require.NotContains(t, def.InputSchema.Properties, field, "%s must not expose backend locator field %s", def.Name, field)
		}
	}
}

func TestKBAgentTeamHandlerRejectsUnknownTool(t *testing.T) {
	h := &kbAgentTeamHandler{}
	result := h.CallTool(context.Background(), "agent-1", "kb_unknown", json.RawMessage(`{}`))
	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "unknown tool: kb_unknown")
}

func TestNewKBHandlerRejectsMissingRequiredDependency(t *testing.T) {
	connection := testAgentTeamProtocol{}
	runtime := testKBDispatcher{}
	for _, config := range []KBHandlerConfig{
		{Runtime: runtime},
		{Connection: connection},
	} {
		_, err := NewKBHandler(config)
		require.Error(t, err)
	}
}

// --- Full call tests with mock HTTP server ---

type testAgentTeamProtocol struct {
	api *apiclient.Client
}

func (p testAgentTeamProtocol) ID() string                              { return "test" }
func (p testAgentTeamProtocol) APIClient(_ ...string) *apiclient.Client { return p.api }
func (p testAgentTeamProtocol) Send(context.Context, types.Metadata, types.OutboundMessage) error {
	return nil
}

type testKBDispatcher struct {
	messageID     string
	serverAgentID string
}

func (d testKBDispatcher) CurrentTurnMessageID(_ string) string { return d.messageID }
func (d testKBDispatcher) Session(_ string) (types.SessionSnapshot, bool) {
	return types.SessionSnapshot{ServerAgentID: d.serverAgentID}, true
}

func newKBAgentTeamTestHandler(t *testing.T, server *httptest.Server, messageID string) *kbAgentTeamHandler {
	t.Helper()
	api := &apiclient.Client{
		ServerURL:        server.URL,
		AgentID:          "srv-agent-1",
		RouteWorkspaceID: "ws_test",
		HTTPClient:       server.Client(),
	}
	return &kbAgentTeamHandler{
		c:       testAgentTeamProtocol{api: api},
		runtime: testKBDispatcher{messageID: messageID, serverAgentID: "srv-agent-1"},
	}
}

func TestKBAgentTeamCallInjectsMessageIDAndStripsTarget(t *testing.T) {
	var capturedPayload map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "/api/v1/workspaces/ws_test/agents/srv-agent-1/runtime/agent-thread/upstream", r.URL.Path)
		body, _ := io.ReadAll(r.Body)
		require.NoError(t, json.Unmarshal(body, &capturedPayload))
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"ok":true}`))
	}))
	defer server.Close()

	h := newKBAgentTeamTestHandler(t, server, "msg-trigger-123")
	result := h.CallTool(context.Background(), "agent-1", "kb_search",
		json.RawMessage(`{"query":"hello","max_results":75,"message_id":"should-be-stripped","target":"should-be-stripped"}`))

	require.False(t, result.IsError, "unexpected error: %s", result.Content[0].Text)

	// message_id must come from the dispatcher, not from args
	require.Equal(t, "msg-trigger-123", capturedPayload["message_id"])
	// model-controlled target/locator fields must be stripped from body
	body := capturedPayload["body"].(map[string]any)
	_, hasTarget := body["target"]
	require.False(t, hasTarget, "target must be stripped from body")
	_, hasMessageID := body["message_id"]
	require.False(t, hasMessageID, "message_id must be stripped from body")
	// query should survive and max_results should be clamped before upstream.
	require.Equal(t, "hello", body["query"])
	require.Equal(t, float64(50), body["max_results"])
	require.Equal(t, float64(50), body["limit"])
	// target is no longer model-supplied; server derives scope from message_id.
	require.NotContains(t, capturedPayload, "thread_channel_id")
	require.NotContains(t, capturedPayload, "target_key")
	// path must use the RESTful retrieval endpoint
	require.Equal(t, "/api/v1/kb/retrieval/knowledge/search", capturedPayload["path"])
	require.Equal(t, "POST", capturedPayload["method"])
}

func TestKBAgentTeamCallRequiresCurrentTurnMessageID(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Fatal("server should not be called when message_id is missing")
	}))
	defer server.Close()

	h := newKBAgentTeamTestHandler(t, server, "")
	result := h.CallTool(context.Background(), "agent-1", "kb_search",
		json.RawMessage(`{"query":"hello"}`))
	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "current turn message_id is required")
}

func TestKBAgentTeamCallUsesCorrectOp(t *testing.T) {
	tests := []struct {
		tool   string
		args   string
		wantOp string
	}{
		{"kb_search", `{"query":"test"}`, "search"},
		{"kb_fetch", `{"citation_ids":["cit_1"]}`, "fetch"},
	}
	for _, tt := range tests {
		t.Run(tt.tool+"_"+tt.wantOp, func(t *testing.T) {
			var capturedPath string
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				body, _ := io.ReadAll(r.Body)
				var payload map[string]any
				require.NoError(t, json.Unmarshal(body, &payload))
				capturedPath = payload["path"].(string)
				w.Header().Set("Content-Type", "application/json")
				_, _ = w.Write([]byte(`{"ok":true}`))
			}))
			defer server.Close()

			h := newKBAgentTeamTestHandler(t, server, "msg-1")
			result := h.CallTool(context.Background(), "agent-1", tt.tool, json.RawMessage(tt.args))
			require.False(t, result.IsError, "unexpected error: %s", result.Content[0].Text)
			require.Equal(t, "/api/v1/kb/retrieval/knowledge/"+tt.wantOp, capturedPath)
		})
	}
}

func TestKBAgentTeamCallReturnsErrorWhenAPIUnavailable(t *testing.T) {
	h := &kbAgentTeamHandler{c: nil, runtime: nil}
	result := h.CallTool(context.Background(), "agent-1", "kb_search",
		json.RawMessage(`{"query":"hello"}`))
	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "agentteam API not available")
}

func TestKBInputPreservesSafeUnknownRawFields(t *testing.T) {
	body, err := decodeKBAgentTeamInput("kb_search", json.RawMessage(`{
		"query":" q ",
		"max_results":3,
		"future":{"nested":[1.20,{"x":true}]}
	}`))
	require.NoError(t, err)
	require.JSONEq(t, `{"query":"q","max_results":3,"limit":3,"future":{"nested":[1.20,{"x":true}]}}`, string(body))
}

func TestKBInputNeverForwardsModelControlledScopeOrLocator(t *testing.T) {
	body, err := decodeKBAgentTeamInput("kb_search", json.RawMessage(`{
		"query":"q","target":"x","message_id":"x","kb_ids":["kb"],"sourceId":"src",
		"path":"internal/a.go","ref":"main","revision":"sha","offset":7,"future":"safe"
	}`))
	require.NoError(t, err)
	for _, forbidden := range []string{"target", "message_id", "kb_ids", "sourceId", "path", "ref", "revision", "offset"} {
		require.NotContains(t, string(body), `"`+forbidden+`"`)
	}
	require.Contains(t, string(body), `"future":"safe"`)
}

func TestKBInputPreservesKnownAliasWireFields(t *testing.T) {
	body, err := decodeKBAgentTeamInput("kb_fetch", json.RawMessage(`{
		"citation_ids":[" cit-1 ","cit-2"],"max_chars_per_item":1200
	}`))
	require.NoError(t, err)
	var decoded map[string]any
	require.NoError(t, json.Unmarshal(body, &decoded))
	require.Equal(t, decoded["citation_ids"], decoded["citationIds"])
	require.Equal(t, float64(1200), decoded["max_chars_per_item"])
	require.Equal(t, float64(1200), decoded["maxBytes"])
}

func TestKBSearchMaxResultsTruncatesFractionalNumbers(t *testing.T) {
	body, err := decodeKBAgentTeamInput("kb_search", json.RawMessage(`{"query":"q","max_results":7.9}`))
	require.NoError(t, err)
	var decoded map[string]any
	require.NoError(t, json.Unmarshal(body, &decoded))
	require.Equal(t, float64(7), decoded["max_results"])
	require.Equal(t, float64(7), decoded["limit"])
}

func TestKBSearchMaxResultsFallsBackForWrongTypes(t *testing.T) {
	for _, value := range []string{`"7"`, `true`, `null`, `{}`, `[]`} {
		body, err := decodeKBAgentTeamInput("kb_search", json.RawMessage(`{"query":"q","max_results":`+value+`}`))
		require.NoError(t, err)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(body, &decoded))
		require.Equal(t, float64(defaultKBSearchMaxResults), decoded["max_results"], value)
	}
}

func TestKBFetchCitationIDsIgnoreNonStringsTrimAndDeduplicate(t *testing.T) {
	body, err := decodeKBAgentTeamInput("kb_fetch", json.RawMessage(`{
		"citation_ids":[" cit-1 ",2,null,true,"","cit-1","cit-2"]
	}`))
	require.NoError(t, err)
	var decoded map[string]any
	require.NoError(t, json.Unmarshal(body, &decoded))
	require.Equal(t, []any{"cit-1", "cit-2"}, decoded["citation_ids"])
}

func TestKBInputDoesNotExposeDynamicValuesToPolicy(t *testing.T) {
	input := kbSearchInput{}
	require.NoError(t, json.Unmarshal([]byte(`{"query":"q","future":{"unsafe":true}}`), &input))
	require.Equal(t, "q", input.Query)
	require.Equal(t, map[string]json.RawMessage{"future": json.RawMessage(`{"unsafe":true}`)}, input.extra)
}

func TestKBResultPreservesExactUpstreamRawJSON(t *testing.T) {
	fixture := "{\n  \"unknown\": {\"n\": 1.2300, \"values\": [true, null]},\n  \"items\": [1,2]\n}\n"
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(fixture))
	}))
	defer server.Close()
	h := newKBAgentTeamTestHandler(t, server, "msg-raw")
	result := h.CallTool(context.Background(), "agent-1", "kb_search", json.RawMessage(`{"query":"q"}`))
	require.False(t, result.IsError)
	require.Equal(t, strings.TrimSpace(fixture), result.Content[0].Text)
	structured, ok := result.StructuredContent.(json.RawMessage)
	require.True(t, ok, "structuredContent type = %T, want json.RawMessage", result.StructuredContent)
	require.Equal(t, strings.TrimSpace(fixture), string(structured))
	encoded, err := json.Marshal(result.StructuredContent)
	require.NoError(t, err)
	require.Equal(t, `{"unknown":{"n":1.2300,"values":[true,null]},"items":[1,2]}`, string(encoded))
}