package dispatcher

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestMemorySearchEligibilityRechecksAuthoritativeAccessAndRetention(t *testing.T) {
	tmp := t.TempDir()
	agentID := "agt-memory-mcp"
	workDir := filepath.Join(tmp, agentID)
	memoryDir := filepath.Join(workDir, "memory")
	require.NoError(t, os.MkdirAll(filepath.Join(memoryDir, "system"), 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(memoryDir, "system", "identity.md"),
		[]byte("# Identity\n\nmemory search needle\n"), 0o644))
	access := MemoryAccess{MemoryID: "mem_1", Branch: "main"}
	stateBytes, err := json.Marshal(vault.RetentionState{
		SchemaVersion: 1,
		MemoryID:      access.MemoryID,
		AgentID:       agentID,
		Branch:        access.Branch,
	})
	require.NoError(t, err)
	require.NoError(t, os.WriteFile(vault.RetentionStatePath(memoryDir), stateBytes, 0o600))

	provider := &mutableMemoryAccessProvider{access: access, state: MemoryAccessBound}
	d := mustNewDispatcher("")
	d.SetMemoryAccessProvider(provider)
	d.EnableMemorySearch(func(id string) string {
		if id == agentID {
			return workDir
		}
		return ""
	}, memorysearch.Scan, nil)
	t.Cleanup(d.Close)

	reg, err := d.snapshotMemorySearch(context.Background())
	require.NoError(t, err)
	rec, _, err := reg.getOrCreate(context.Background(), agentID, access)
	require.NoError(t, err)
	require.Eventually(t, func() bool {
		gen, err := rec.Current(context.Background())
		if err != nil {
			return false
		}
		basename := filepath.Base(gen.Dir())
		_ = gen.Close()
		return rec.SearchState(basename) == memorysearch.StateReady
	}, 5*time.Second, 20*time.Millisecond)

	require.True(t, d.MemorySearchAvailable(context.Background(), agentID))
	response, err := d.MemorySearch(context.Background(), agentID,
		memorysearch.SearchRequest{Query: "needle"})
	require.NoError(t, err)
	require.Equal(t, memorysearch.StateReady, response.State)
	response, err = d.MemorySearch(context.Background(), agentID,
		memorysearch.SearchRequest{Query: "needle"})
	require.NoError(t, err)
	require.Equal(t, memorysearch.StateReady, response.State,
		"tool-time no-op freshness checks must not mark a healthy index stale")

	provider.state = MemoryAccessUnbound
	// The MCP hot path intentionally does not perform provider HTTP. Until an
	// authoritative refresh lands, the last Bound snapshot remains eligible.
	require.True(t, d.MemorySearchAvailable(context.Background(), agentID))
	d.RefreshMemoryAccess(context.Background(), []string{agentID})
	require.False(t, d.MemorySearchAvailable(context.Background(), agentID))
}

func TestMemorySearchGetOrCreateFailsClosedWhenRetentionMarkerMissing(t *testing.T) {
	tmp := t.TempDir()
	agentID := "agt-memory-mcp-missing"
	workDir := filepath.Join(tmp, agentID)
	require.NoError(t, os.MkdirAll(filepath.Join(workDir, "memory"), 0o755))
	access := MemoryAccess{MemoryID: "mem_1", Branch: "main"}
	d := mustNewDispatcher("")
	d.EnableMemorySearch(func(id string) string {
		if id == agentID {
			return workDir
		}
		return ""
	}, memorysearch.Scan, nil)
	t.Cleanup(d.Close)
	reg, err := d.snapshotMemorySearch(context.Background())
	require.NoError(t, err)
	_, _, err = reg.getOrCreate(context.Background(), agentID, access)
	require.Error(t, err)
	require.False(t, d.MemorySearchAvailable(context.Background(), agentID))
}

type mutableMemoryAccessProvider struct {
	access MemoryAccess
	state  MemoryAccessState
	err    error
}

func (p *mutableMemoryAccessProvider) Lookup(
	context.Context,
	string,
) (MemoryAccess, MemoryAccessState, error) {
	return p.access, p.state, p.err
}