package dispatcher

import (
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func isolateTrajectoryStorage(t *testing.T) {
	t.Helper()

	cfgDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, cfgDir)

	cfg := profile.DefaultConfig()
	cfg.Log.Dir = filepath.Join(cfgDir, "logs")
	require.NoError(t, profile.Save(cfg, filepath.Join(cfgDir, "config.yaml")))
}

func TestTrajectoryHubSubscribeCapsInitialReplay(t *testing.T) {
	isolateTrajectoryStorage(t)
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
	}
	items := make([]types.TrajectoryLogItem, trajectoryReplayLimit+50)
	now := time.Now().UnixMilli()
	for i := range items {
		items[i] = types.TrajectoryLogItem{
			AgentID:   "agent-a",
			Entry:     types.TrajectoryEntry{Kind: "prompt_input", Text: "item"},
			Timestamp: now + int64(i),
		}
	}
	h.hot["agent-a"] = &agentCache{items: items}

	ch := h.Subscribe("agent-a")
	defer h.Unsubscribe("agent-a", ch)

	got := drainTrajectoryChannel(ch)
	require.Len(t, got, trajectoryReplayLimit)
	require.Equal(t, items[len(items)-trajectoryReplayLimit].Timestamp, got[0].Timestamp)
	require.Equal(t, items[len(items)-1].Timestamp, got[len(got)-1].Timestamp)
}

func TestTrajectoryHubHistoryCompactsStreamedTextBeforeLimit(t *testing.T) {
	isolateTrajectoryStorage(t)
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
	}
	now := time.Now().UnixMilli()
	h.hot["agent-a"] = &agentCache{items: []types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "text", Text: "pi"}},
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: "text", Text: "ng"}},
		{AgentID: "agent-a", Timestamp: now + 2, Entry: types.TrajectoryEntry{Kind: "text", Text: "cap"}},
	}}

	got := h.History("agent-a", 1, 0)

	require.Len(t, got, 1)
	require.Equal(t, "pingcap", got[0].Entry.Text)
	require.Equal(t, now, got[0].Timestamp)
}

func TestTrajectoryHubSubscribeCompactsColdHotReplay(t *testing.T) {
	isolateTrajectoryStorage(t)
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
	}
	now := time.Now().UnixMilli()
	appendCold("agent-a", []types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "text", Text: "ping"}},
	})
	h.hot["agent-a"] = &agentCache{items: []types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: "text", Text: "cap"}},
	}}

	ch := h.Subscribe("agent-a")
	defer h.Unsubscribe("agent-a", ch)

	got := drainTrajectoryChannel(ch)
	require.Len(t, got, 1)
	require.Equal(t, "pingcap", got[0].Entry.Text)
}

func TestTrajectoryHubPreservesDistinctSubagentActivitiesAcrossHotColdAndHistory(t *testing.T) {
	isolateTrajectoryStorage(t)
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
	}
	now := time.Now().UnixMilli()
	activity := func(agentID, itemID string, timestamp int64, extra map[string]any) types.TrajectoryLogItem {
		metadata := map[string]any{
			"runtimeItemType": "subAgentActivity",
			"itemId":          itemID,
			"activityKind":    "interacted",
			"agentPath":       "/root/reviewer",
		}
		for key, value := range extra {
			metadata[key] = value
		}
		return types.TrajectoryLogItem{
			AgentID:   agentID,
			Timestamp: timestamp,
			Entry: types.TrajectoryEntry{
				Kind:     consts.TrajectoryKindStatus,
				Activity: "working",
				Detail:   "Interacted with /root/reviewer",
				Metadata: metadata,
			},
		}
	}
	itemID := func(item types.TrajectoryLogItem) string {
		value, _ := item.Entry.Metadata["itemId"].(string)
		return value
	}

	h.hot["agent-hot"] = &agentCache{items: []types.TrajectoryLogItem{
		activity("agent-hot", "activity-1", now, nil),
		activity("agent-hot", "activity-1", now+1, map[string]any{"replayed": true}),
		activity("agent-hot", "activity-2", now+2, nil),
	}}
	hot := h.Recent("agent-hot", 10)
	require.Len(t, hot, 2)
	require.Equal(t, []string{"activity-1", "activity-2"}, []string{itemID(hot[0]), itemID(hot[1])})
	require.Equal(t, true, hot[0].Entry.Metadata["replayed"], "same-item replay should merge metadata")

	appendCold("agent-cold-hot", []types.TrajectoryLogItem{
		activity("agent-cold-hot", "activity-1", now, nil),
	})
	h.hot["agent-cold-hot"] = &agentCache{items: []types.TrajectoryLogItem{
		activity("agent-cold-hot", "activity-1", now+1, map[string]any{"replayed": true}),
		activity("agent-cold-hot", "activity-2", now+2, nil),
	}}

	ch := h.Subscribe("agent-cold-hot")
	defer h.Unsubscribe("agent-cold-hot", ch)
	replay := drainTrajectoryChannel(ch)
	require.Len(t, replay, 2)
	require.Equal(t, []string{"activity-1", "activity-2"}, []string{itemID(replay[0]), itemID(replay[1])})

	history := h.History("agent-cold-hot", 10, 0)
	require.Len(t, history, 2)
	require.Equal(t, []string{"activity-2", "activity-1"}, []string{itemID(history[0]), itemID(history[1])})
}

func TestCompactTrajectoryLogItemsKeepsToolOutputBoundaries(t *testing.T) {
	now := time.Now().UnixMilli()
	got := compactTrajectoryLogItems([]types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "tool_output", ToolName: "shell", ToolCallID: "call-1", Text: "line 1"}},
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: "tool_output", ToolName: "shell", ToolCallID: "call-1", Text: "\nline 2"}},
		{AgentID: "agent-a", Timestamp: now + 2, Entry: types.TrajectoryEntry{Kind: "tool_output", ToolName: "shell", ToolCallID: "call-2", Text: "separate"}},
	})

	require.Len(t, got, 2)
	require.Equal(t, "shell", got[0].Entry.ToolName)
	require.Equal(t, "call-1", got[0].Entry.ToolCallID)
	require.Equal(t, "line 1\nline 2", got[0].Entry.Text)
	require.Equal(t, "shell", got[1].Entry.ToolName)
	require.Equal(t, "call-2", got[1].Entry.ToolCallID)
	require.Equal(t, "separate", got[1].Entry.Text)
}

func TestCompactTrajectoryLogItemsReplacesCumulativeToolOutput(t *testing.T) {
	now := time.Now().UnixMilli()
	got := compactTrajectoryLogItems([]types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "tool_output", ToolName: "shell", ToolCallID: "call-1", Text: "line 1\n"}},
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: "tool_output", ToolName: "shell", ToolCallID: "call-1", Text: "line 1\nline 2\n"}},
		{AgentID: "agent-a", Timestamp: now + 2, Entry: types.TrajectoryEntry{Kind: "tool_output", ToolName: "shell", ToolCallID: "call-1", Text: "line 1\nline 2\n"}},
	})

	require.Len(t, got, 1)
	require.Equal(t, "line 1\nline 2\n", got[0].Entry.Text)
}

func TestCompactTrajectoryLogItemsKeepsLatestContextSnapshotPerSession(t *testing.T) {
	now := time.Now().UnixMilli()
	got := compactTrajectoryLogItems([]types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: consts.TrajectoryKindContextUsage, SessionID: "sess-1", ActiveContextTokens: 40_000, ContextWindowTokens: 258_400}},
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: consts.TrajectoryKindContextUsage, SessionID: "sess-1", ActiveContextTokens: 48_000, ContextWindowTokens: 258_400}},
		{AgentID: "agent-a", Timestamp: now + 2, Entry: types.TrajectoryEntry{Kind: consts.TrajectoryKindContextUsage, SessionID: "sess-2", ActiveContextTokens: 20_000, ContextWindowTokens: 258_400}},
	})

	require.Len(t, got, 2)
	require.Equal(t, uint64(48_000), got[0].Entry.ActiveContextTokens)
	require.Equal(t, now+1, got[0].Timestamp)
	require.Equal(t, "sess-2", got[1].Entry.SessionID)
}

func TestCompactTrajectoryLogItemsKeepsThinkingSectionBoundaries(t *testing.T) {
	now := time.Now().UnixMilli()
	got := compactTrajectoryLogItems([]types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "thinking", Text: "**Plan", SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: "thinking", Text: "**\nInspect", Title: "Plan", SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 2, Entry: types.TrajectoryEntry{Kind: "thinking", Text: "**Act**\nEdit", Title: "Act", SectionID: "reasoning-1:summary:1"}},
	})

	require.Len(t, got, 2)
	require.Equal(t, "**Plan**\nInspect", got[0].Entry.Text)
	require.Equal(t, "Plan", got[0].Entry.Title)
	require.Equal(t, "reasoning-1:summary:0", got[0].Entry.SectionID)
	require.Equal(t, "**Act**\nEdit", got[1].Entry.Text)
	require.Equal(t, "Act", got[1].Entry.Title)
	require.Equal(t, "reasoning-1:summary:1", got[1].Entry.SectionID)
}

func TestCompactTrajectoryLogItemsPreservesThinkingDeltaWhitespace(t *testing.T) {
	now := time.Now().UnixMilli()
	got := compactTrajectoryLogItems([]types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "thinking", Text: "**Constructing valid JSON**\n\nI", SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 1, Entry: types.TrajectoryEntry{Kind: "thinking", Text: " need", SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 2, Entry: types.TrajectoryEntry{Kind: "thinking", Text: " to", SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 3, Entry: types.TrajectoryEntry{Kind: "thinking", Text: " include ", SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 4, Entry: types.TrajectoryEntry{Kind: "thinking", Text: `"schema_version"`, SectionID: "reasoning-1:summary:0"}},
		{AgentID: "agent-a", Timestamp: now + 5, Entry: types.TrajectoryEntry{Kind: "thinking", Text: ".", SectionID: "reasoning-1:summary:0"}},
	})

	require.Len(t, got, 1)
	require.Equal(t, `**Constructing valid JSON**

I need to include "schema_version".`, got[0].Entry.Text)
}

func TestTrajectoryHubHotBytesAccounting(t *testing.T) {
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
	}
	h.Broadcast("agent-a", []types.TrajectoryEntry{{Kind: "text", Text: "hello"}})
	require.Greater(t, h.hotBytes, 0)
	before := h.hotBytes

	h.ClearAgent("agent-a")
	require.Zero(t, h.hotBytes)
	require.Less(t, h.hotBytes, before)
}

func TestTrajectoryHubBroadcastPreservesRetainedLargeFields(t *testing.T) {
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
	}
	text := strings.Repeat("x", 2400)
	toolInput := `{"target":"dm:test","content":"` + strings.Repeat("y", 2400) + `"}`
	h.Broadcast("agent-a", []types.TrajectoryEntry{{
		Kind:      "prompt_input",
		Text:      text,
		ToolInput: toolInput,
	}})

	recent := h.Recent("agent-a", 1)
	require.Len(t, recent, 1)
	require.Equal(t, text, recent[0].Entry.Text)
	require.Equal(t, toolInput, recent[0].Entry.ToolInput)
}

func TestTrajectoryHubSubscribeLiveFansOutNewEventsOnly(t *testing.T) {
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
		allSubs:     make(map[chan types.TrajectoryLogItem]struct{}),
	}
	h.Broadcast("agent-a", []types.TrajectoryEntry{{Kind: "text", Text: "old"}})

	ready := make(chan (<-chan types.TrajectoryLogItem), 1)
	go func() {
		ready <- h.SubscribeLive("agent-a")
	}()

	var ch <-chan types.TrajectoryLogItem
	select {
	case ch = <-ready:
	case <-time.After(time.Second):
		t.Fatal("SubscribeLive did not return")
	}
	require.Empty(t, drainTrajectoryChannel(ch), "per-agent live subscribers should not replay history")

	h.Broadcast("agent-a", []types.TrajectoryEntry{{Kind: "status", Activity: "working", Detail: "new-a"}})
	h.Broadcast("agent-b", []types.TrajectoryEntry{{Kind: "status", Activity: "working", Detail: "new-b"}})

	got := drainTrajectoryChannel(ch)
	require.Len(t, got, 1)
	require.Equal(t, "agent-a", got[0].AgentID)
	require.Equal(t, "new-a", got[0].Entry.Detail)

	h.Unsubscribe("agent-a", ch)
	_, ok := <-ch
	require.False(t, ok)
}

func TestTrajectoryHubSubscribeAllLiveFansOutNewEventsOnly(t *testing.T) {
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
		allSubs:     make(map[chan types.TrajectoryLogItem]struct{}),
	}
	h.Broadcast("agent-a", []types.TrajectoryEntry{{Kind: "text", Text: "old"}})

	ch := h.SubscribeAllLive()
	require.Empty(t, drainTrajectoryChannel(ch), "all-agent live subscribers should not replay history")

	h.Broadcast("agent-a", []types.TrajectoryEntry{{Kind: "text", Text: "new-a"}})
	h.Broadcast("agent-b", []types.TrajectoryEntry{{Kind: "status", Activity: "working", Detail: "new-b"}})

	got := drainTrajectoryChannel(ch)
	require.Len(t, got, 2)
	require.Equal(t, "agent-a", got[0].AgentID)
	require.Equal(t, "new-a", got[0].Entry.Text)
	require.Equal(t, "agent-b", got[1].AgentID)
	require.Equal(t, "new-b", got[1].Entry.Detail)

	h.UnsubscribeAll(ch)
	_, ok := <-ch
	require.False(t, ok)
}

func TestTrajectoryHubLatestByAgentReturnsNewestRequestedItems(t *testing.T) {
	isolateTrajectoryStorage(t)
	h := &trajectoryHub{
		hot:         make(map[string]*agentCache),
		subscribers: make(map[string]map[chan types.TrajectoryLogItem]struct{}),
		allSubs:     make(map[chan types.TrajectoryLogItem]struct{}),
	}
	now := time.Now().UnixMilli()
	h.hot["agent-a"] = &agentCache{items: []types.TrajectoryLogItem{
		{AgentID: "agent-a", Timestamp: now - 2000, Entry: types.TrajectoryEntry{Kind: "prompt_input", Text: "old-a"}},
		{AgentID: "agent-a", Timestamp: now, Entry: types.TrajectoryEntry{Kind: "text", Text: "new-a"}},
	}}
	h.hot["agent-b"] = &agentCache{items: []types.TrajectoryLogItem{
		{AgentID: "agent-b", Timestamp: now - 1000, Entry: types.TrajectoryEntry{Kind: "status", Activity: "working"}},
	}}

	got := h.LatestByAgent([]string{"agent-a", "missing"})
	require.Len(t, got, 1)
	require.Equal(t, "new-a", got["agent-a"].Entry.Text)
	require.NotContains(t, got, "agent-b")
	require.NotContains(t, got, "missing")
}

func drainTrajectoryChannel(ch <-chan types.TrajectoryLogItem) []types.TrajectoryLogItem {
	var out []types.TrajectoryLogItem
	for {
		select {
		case item := <-ch:
			out = append(out, item)
		default:
			return out
		}
	}
}

// TestTrajectoryHubSubscribeSeesConcurrentBroadcasts verifies that no
// trajectory events are lost when Broadcast and Subscribe run concurrently.
// This is a regression test for a race where Subscribe released the lock
// before registering the subscriber, creating a window where broadcasts
// would be invisible to the new subscriber.
func TestTrajectoryHubSubscribeSeesConcurrentBroadcasts(t *testing.T) {
	isolateTrajectoryStorage(t)
	h := newTrajectoryHub()
	defer h.Stop()

	const broadcasters = 4
	const broadcastsEach = 50

	var wg sync.WaitGroup
	wg.Add(broadcasters + 1)

	// Subscriber goroutine: subscribe and collect all events.
	var subItems []types.TrajectoryLogItem
	var subDone bool
	subscribed := make(chan struct{})
	go func() {
		defer wg.Done()
		ch := h.Subscribe("agent-a")
		close(subscribed)
		defer h.Unsubscribe("agent-a", ch)
		for {
			item, ok := <-ch
			if !ok {
				return
			}
			subItems = append(subItems, item)
			// Stop after we've seen enough events (broadcasts + possible initial replay).
			if len(subItems) >= broadcasters*broadcastsEach {
				subDone = true
				return
			}
		}
	}()

	select {
	case <-subscribed:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for trajectory subscriber registration")
	}

	// Broadcaster goroutines.
	for i := 0; i < broadcasters; i++ {
		go func(idx int) {
			defer wg.Done()
			for j := 0; j < broadcastsEach; j++ {
				h.Broadcast("agent-a", []types.TrajectoryEntry{
					{Kind: "text", Text: strings.Repeat("x", 200)},
				})
			}
		}(i)
	}

	wg.Wait()
	require.True(t, subDone, "subscriber should have received all broadcast events")
	require.Equal(t, broadcasters*broadcastsEach, len(subItems),
		"subscriber must not miss any events despite concurrent broadcasts")
}

func TestTrajectoryHubConcurrentStopIsIdempotentAndJoinsGC(t *testing.T) {
	h := newTrajectoryHub()

	const callers = 64
	start := make(chan struct{})
	var wg sync.WaitGroup
	wg.Add(callers)
	for range callers {
		go func() {
			defer wg.Done()
			<-start
			h.Stop()
		}()
	}
	close(start)

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("concurrent trajectory Stop callers did not join GC completion")
	}
	select {
	case <-h.gcDone:
	default:
		t.Fatal("Stop returned before the trajectory GC owner completed")
	}
}

func TestTrajectoryHubStopWaitsForAlreadyRunningGCOwner(t *testing.T) {
	h := &trajectoryHub{
		stopGC: make(chan struct{}),
		gcDone: make(chan struct{}),
	}
	gcObservedStop := make(chan struct{})
	releaseGC := make(chan struct{})
	go func() {
		<-h.stopGC
		close(gcObservedStop)
		<-releaseGC
		close(h.gcDone)
	}()

	stopReturned := make(chan struct{})
	go func() {
		h.Stop()
		close(stopReturned)
	}()
	select {
	case <-gcObservedStop:
	case <-time.After(time.Second):
		t.Fatal("Stop did not signal the GC owner")
	}
	select {
	case <-stopReturned:
		t.Fatal("Stop returned while the GC owner was still running")
	case <-time.After(30 * time.Millisecond):
	}
	close(releaseGC)
	select {
	case <-stopReturned:
	case <-time.After(time.Second):
		t.Fatal("Stop did not return after the GC owner completed")
	}
}