package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"

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

type memorySearchServiceFake struct {
	available bool
	agentID   string
	request   memorysearch.SearchRequest
	response  memorysearch.SearchResponse
	err       error
}

func (f *memorySearchServiceFake) MemorySearchAvailable(context.Context, string) bool {
	return f.available
}

func (f *memorySearchServiceFake) MemorySearch(
	_ context.Context,
	agentID string,
	request memorysearch.SearchRequest,
) (memorysearch.SearchResponse, error) {
	f.agentID = agentID
	f.request = request
	return f.response, f.err
}

func TestMemoryMCPToolsAlwaysAdvertisedWhenIndexUnavailable(t *testing.T) {
	fake := &memorySearchServiceFake{}
	handler := newMemoryHandler(fake, "route-agent")
	tools := handler.Tools()
	require.Len(t, tools, 1)
	require.Contains(t, tools[0].Description, "one focused topic")
	require.Contains(t, tools[0].InputSchema.Properties["query"].Description, "2-6 discriminative terms")
	require.Contains(t, tools[0].InputSchema.Properties["limit"].Description, "default 3")
	require.Contains(t, tools[0].InputSchema.Properties["maxChars"].Description, "default 1600")
	// The tool is always advertised; unavailable indexes fail at call time
	// instead of hiding the tool from tools/list.
	require.Len(t, handler.Tools(), 1)
}

func TestMemoryMCPCallUsesAuthenticatedRouteAgentAndCompactTextPayload(t *testing.T) {
	fake := &memorySearchServiceFake{
		available: true,
		response: memorysearch.SearchResponse{
			State: memorysearch.StateReady,
			SourceRoot: memorysearch.SourceRoot{
				Kind:         "cloudMemory",
				AbsolutePath: "/safe/memory",
			},
			Results: []memorysearch.SearchResult{{
				Path:        "knowledge/a.md",
				HeadingPath: []string{"A"},
				Kind:        "concept",
				MatchType:   memorysearch.MatchTypeLexical,
				MatchLines:  memorysearch.LineRange{Start: 2, End: 3},
				ReadLines:   memorysearch.LineRange{Start: 1, End: 5},
				Snippet:     "needle",
			}},
		},
	}
	handler := newMemoryHandler(fake, "route-agent")
	result := handler.CallTool(context.Background(), "ignored-call-agent", "memory_search",
		json.RawMessage(`{"query":"needle","limit":3,"pathPrefix":"knowledge/","kinds":["concept"],"maxChars":2000}`))

	require.False(t, result.IsError)
	require.Equal(t, "route-agent", fake.agentID)
	require.Equal(t, "needle", fake.request.Query)
	require.Equal(t, 3, fake.request.Limit)
	require.Equal(t, "knowledge/", fake.request.PathPrefix)
	require.Equal(t, []string{"concept"}, fake.request.Kinds)
	require.Nil(t, result.StructuredContent)
	require.Len(t, result.Content, 1)
	require.JSONEq(t, `{
		"results": [{
			"path": "knowledge/a.md",
			"lines": [1, 5],
			"snippet": "needle"
		}]
	}`, result.Content[0].Text)
	require.NotContains(t, result.Content[0].Text, "\n",
		"agent payload must stay compact instead of pretty-printed")
}

func TestMemoryMCPCallReturnsToolErrorWhenIndexUnavailable(t *testing.T) {
	fake := &memorySearchServiceFake{available: true}
	handler := newMemoryHandler(fake, "route-agent")
	require.Len(t, handler.Tools(), 1)
	fake.available = false

	result := handler.CallTool(context.Background(), "route-agent", "memory_search",
		json.RawMessage(`{"query":"needle"}`))
	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "index is unavailable")
	require.Nil(t, result.StructuredContent)
}

func TestMemoryMCPCallReturnsToolErrorForDegradedSearchResponse(t *testing.T) {
	fake := &memorySearchServiceFake{
		available: true,
		response: memorysearch.SearchResponse{
			State:  memorysearch.StateDegraded,
			Reason: memorysearch.ReasonIndexUnavailable,
		},
	}
	handler := newMemoryHandler(fake, "route-agent")

	result := handler.CallTool(context.Background(), "route-agent", "memory_search",
		json.RawMessage(`{"query":"needle"}`))

	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "index is unavailable")
	require.Nil(t, result.StructuredContent)
}

func TestMemoryMCPCallOmitsDefaultStateAndIncludesActionableFlags(t *testing.T) {
	fake := &memorySearchServiceFake{
		available: true,
		response: memorysearch.SearchResponse{
			State:     memorysearch.StateRefreshing,
			Stale:     true,
			Truncated: true,
		},
	}
	handler := newMemoryHandler(fake, "route-agent")

	result := handler.CallTool(context.Background(), "route-agent", "memory_search",
		json.RawMessage(`{"query":"needle"}`))

	require.False(t, result.IsError)
	require.Nil(t, result.StructuredContent)
	require.JSONEq(t, `{"results":[],"stale":true,"truncated":true}`, result.Content[0].Text)
	require.NotContains(t, result.Content[0].Text, "state")
	require.NotContains(t, result.Content[0].Text, "reason")
}

func TestMemoryMCPCallReturnsEmptyResultsForNoConfidentMatch(t *testing.T) {
	fake := &memorySearchServiceFake{
		available: true,
		response: memorysearch.SearchResponse{
			State:  memorysearch.StateReady,
			Reason: memorysearch.ReasonNoConfidentMatch,
		},
	}
	handler := newMemoryHandler(fake, "route-agent")

	result := handler.CallTool(context.Background(), "route-agent", "memory_search",
		json.RawMessage(`{"query":"needle"}`))

	require.False(t, result.IsError)
	require.JSONEq(t, `{"results":[]}`, result.Content[0].Text)
}

func TestMemoryMCPCallRejectsIdentityAndRootOverrides(t *testing.T) {
	fake := &memorySearchServiceFake{available: true}
	handler := newMemoryHandler(fake, "route-agent")

	for _, args := range []string{
		`{"query":"needle","agentID":"other-agent"}`,
		`{"query":"needle","root":"/tmp/foreign-memory"}`,
	} {
		result := handler.CallTool(context.Background(), "route-agent", "memory_search",
			json.RawMessage(args))
		require.True(t, result.IsError)
		require.Contains(t, result.Content[0].Text, "invalid tool arguments")
	}
	require.Empty(t, fake.agentID, "rejected overrides must never reach the search service")
}

func TestMemoryMCPCallDoesNotExposeInternalRuntimeError(t *testing.T) {
	const secret = "binding fp_old fp_new /Users/alice/.derived/memory-search/generations/gen-42/manifest.json"
	fake := &memorySearchServiceFake{
		available: true,
		err:       errors.New(secret),
	}
	handler := newMemoryHandler(fake, "route-agent")

	result := handler.CallTool(context.Background(), "route-agent", "memory_search",
		json.RawMessage(`{"query":"needle"}`))
	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "memory search is temporarily unavailable")
	require.NotContains(t, result.Content[0].Text, secret)
	require.NotContains(t, result.Content[0].Text, "manifest.json")
}

func TestMemoryMCPCallPreservesSafeInvalidRequestMessage(t *testing.T) {
	fake := &memorySearchServiceFake{
		available: true,
		err: fmt.Errorf("%w: query must contain at least 2 characters",
			memorysearch.ErrInvalidSearchRequest),
	}
	handler := newMemoryHandler(fake, "route-agent")

	result := handler.CallTool(context.Background(), "route-agent", "memory_search",
		json.RawMessage(`{"query":"x"}`))
	require.True(t, result.IsError)
	require.Contains(t, result.Content[0].Text, "query must contain at least 2 characters")
}

func TestToolMuxMemoryCallsUseAuthenticatedAgent(t *testing.T) {
	searcher := &recordingMemorySearcher{}
	mux := newMemoryToolMux(t, searcher)

	result := mux.CallTool(context.Background(), "agent-a", "memory_search", json.RawMessage(`{"query":"needle"}`))
	require.False(t, result.IsError)
	require.Equal(t, []string{"agent-a"}, searcher.agents())
}

func TestToolMuxMemoryCallsIsolateConcurrentAgents(t *testing.T) {
	searcher := &recordingMemorySearcher{}
	mux := newMemoryToolMux(t, searcher)

	var wg sync.WaitGroup
	for _, agentID := range []string{"agent-a", "agent-b"} {
		agentID := agentID
		wg.Add(1)
		go func() {
			defer wg.Done()
			result := mux.CallTool(context.Background(), agentID, "memory_search", json.RawMessage(`{"query":"needle"}`))
			require.False(t, result.IsError)
		}()
	}
	wg.Wait()
	require.ElementsMatch(t, []string{"agent-a", "agent-b"}, searcher.agents())
}

func TestToolMuxMemoryInventoryDoesNotCacheAgentIdentity(t *testing.T) {
	searcher := &recordingMemorySearcher{}
	mux := newMemoryToolMux(t, searcher)

	require.Equal(t, []string{"memory_search"}, toolMuxToolNames(mux.Tools()))
	require.Empty(t, searcher.agents())
	require.False(t, mux.CallTool(context.Background(), "agent-a", "memory_search", json.RawMessage(`{"query":"needle"}`)).IsError)
	require.False(t, mux.CallTool(context.Background(), "agent-b", "memory_search", json.RawMessage(`{"query":"needle"}`)).IsError)
	require.Equal(t, []string{"agent-a", "agent-b"}, searcher.agents())
}

func newMemoryToolMux(t *testing.T, searcher MemorySearcher) *toolMux {
	t.Helper()
	mux, err := newToolMux(
		toolChild{kind: toolChildMessage, inventory: &toolMuxTestHandler{}},
		toolChild{
			kind:      toolChildMemory,
			inventory: newMemoryHandler(searcher, ""),
			forAgent: func(agentID string) types.MCPHandler {
				return newMemoryHandler(searcher, agentID)
			},
		},
	)
	require.NoError(t, err)
	return mux
}

type recordingMemorySearcher struct {
	mu       sync.Mutex
	agentIDs []string
}

func (*recordingMemorySearcher) MemorySearchAvailable(context.Context, string) bool { return true }
func (s *recordingMemorySearcher) MemorySearch(_ context.Context, agentID string, _ memorysearch.SearchRequest) (memorysearch.SearchResponse, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.agentIDs = append(s.agentIDs, agentID)
	return memorysearch.SearchResponse{State: memorysearch.StateReady}, nil
}

func (s *recordingMemorySearcher) agents() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]string(nil), s.agentIDs...)
}