package dispatcher

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestSessionKeyStableAndDelimiterSafe(t *testing.T) {
	key := NewSessionKey("agent:with:colon", "chn/main", "#agentteam:abc:def").WithKind("normal:kind")
	encoded := key.String()

	reparsed, ok := parseSessionKey(encoded)
	require.True(t, ok)
	require.Equal(t, key, reparsed)
	require.Equal(t, encoded, reparsed.String())
}

func TestConversationKeyForMetadataPrefersAgentTeamTarget(t *testing.T) {
	meta := &types.AgentTeamMeta{AgentID: "agent-a", Target: "#agentteam:thread-a"}
	require.Equal(t, "#agentteam:thread-a", conversationKeyForMetadata(meta))

	meta.Target = ""
	require.Equal(t, "agent-a", conversationKeyForMetadata(meta))
}

func TestWaitForSessionStartReturnsContextCancellation(t *testing.T) {
	waitCh := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	err := waitForSessionStart(ctx, waitCh)
	require.ErrorIs(t, err, context.Canceled)
}

func TestSessionStartReservationSerializesSameExpectedSession(t *testing.T) {
	manager := NewSessionManager()
	key := "agent-reserved"
	sess := &interactiveSession{sessionKey: key, agentID: key}
	manager.Set(key, sess)

	first, firstWait, firstOwns := manager.BeginStart(key, sess)
	require.True(t, firstOwns)
	require.NotNil(t, first)

	second, secondWait, secondOwns := manager.BeginStart(key, sess)
	require.False(t, secondOwns)
	require.Nil(t, second)
	require.Equal(t, firstWait, secondWait)
	select {
	case <-secondWait:
		t.Fatal("second starter was released before the first generation finished")
	default:
	}

	manager.FinishStart(first)
	select {
	case <-secondWait:
	case <-time.After(time.Second):
		t.Fatal("waiting starter was not released")
	}

	next, _, nextOwns := manager.BeginStart(key, sess)
	require.True(t, nextOwns)
	require.Greater(t, next.generation, first.generation)
	manager.FinishStart(next)
}

func TestSessionStartReservationRejectsLatePublishAfterRemoval(t *testing.T) {
	manager := NewSessionManager()
	key := "agent-removed"
	sess := &interactiveSession{sessionKey: key, agentID: key}
	manager.Set(key, sess)

	reservation, _, owns := manager.BeginStart(key, sess)
	require.True(t, owns)
	require.True(t, manager.RemoveIfCurrent(key, sess))
	require.True(t, sess.retired.Load())

	process := &agent.Process{AgentID: key}
	replaced, published := manager.PublishStart(reservation, sess, process, "old-generation", false)
	require.False(t, published)
	require.Nil(t, replaced)
	_, exists := manager.Get(key)
	require.False(t, exists)
	sess.mu.Lock()
	require.Nil(t, sess.process)
	sess.mu.Unlock()
	manager.FinishStart(reservation)
}

func TestSessionStartReservationRejectsFreshPublishAfterDelete(t *testing.T) {
	manager := NewSessionManager()
	key := "agent-fresh-deleted"

	reservation, waitCh, owns := manager.BeginStart(key, nil)
	require.True(t, owns)
	require.NotNil(t, reservation)
	manager.Delete(key)
	select {
	case <-waitCh:
	case <-time.After(time.Second):
		t.Fatal("Delete did not release waiters for the invalidated generation")
	}

	sess := &interactiveSession{sessionKey: key, agentID: key}
	process := &agent.Process{AgentID: key}
	replaced, published := manager.PublishStart(reservation, sess, process, "deleted-generation", false)
	require.False(t, published)
	require.Nil(t, replaced)
	_, exists := manager.Get(key)
	require.False(t, exists)
	sess.mu.Lock()
	require.Nil(t, sess.process)
	sess.mu.Unlock()

	// Finish after Delete invalidated and closed the reservation must be a
	// harmless no-op rather than a double close.
	manager.FinishStart(reservation)
}

func TestSessionManagerCloseFencesAllStartsAndRejectsNewGeneration(t *testing.T) {
	manager := NewSessionManager()
	live := &interactiveSession{sessionKey: "agent-live", agentID: "agent-live"}
	manager.Set(live.sessionKey, live)
	first, firstWait, firstOwns := manager.BeginStart("agent-first", nil)
	second, secondWait, secondOwns := manager.BeginStart("agent-second", nil)
	require.True(t, firstOwns)
	require.True(t, secondOwns)

	manager.Close()
	require.True(t, manager.Closed())
	require.True(t, live.retired.Load())
	require.Contains(t, manager.Snapshot(true), live, "Close must leave published sessions snapshot-able for Dispatcher teardown")
	for _, waitCh := range []<-chan struct{}{firstWait, secondWait} {
		select {
		case <-waitCh:
		case <-time.After(time.Second):
			t.Fatal("Close did not release an in-flight start waiter")
		}
	}

	for _, reservation := range []*sessionStartReservation{first, second} {
		_, published := manager.PublishStart(reservation, &interactiveSession{}, &agent.Process{}, "late", false)
		require.False(t, published)
		manager.FinishStart(reservation)
	}
	reservation, waitCh, owns := manager.BeginStart("agent-after-close", nil)
	require.False(t, owns)
	require.Nil(t, reservation)
	select {
	case <-waitCh:
	default:
		t.Fatal("closed manager must return an already-closed wait channel")
	}
}

func TestSessionManagerCloseReleasesWaitersBeforeReservationOwnerFinishes(t *testing.T) {
	manager := NewSessionManager()
	reservation, waiterDone, owns := manager.BeginStart("agent-owner-quiesce", nil)
	require.True(t, owns)
	require.NotNil(t, reservation)

	owners := manager.Close()
	require.Len(t, owners, 1)
	select {
	case <-waiterDone:
	case <-time.After(time.Second):
		t.Fatal("Close did not immediately release external start waiters")
	}
	select {
	case <-owners[0]:
		t.Fatal("owner completion closed before the reservation owner called FinishStart")
	default:
	}

	// Close itself never waits for ownerDone, so a reservation owner may safely
	// initiate manager shutdown and then unwind its own reservation.
	manager.FinishStart(reservation)
	select {
	case <-owners[0]:
	case <-time.After(time.Second):
		t.Fatal("FinishStart did not publish reservation-owner completion")
	}
}

func TestDispatcherCloseWaitsForInvalidatedStartOwnerAndRejectsLatePublish(t *testing.T) {
	d := mustNewDispatcher("http://server.test")
	reservation, waiterDone, owns := d.sessionManager.BeginStart("agent-close-owner", nil)
	require.True(t, owns)

	closeDone := make(chan struct{})
	go func() {
		d.Close()
		close(closeDone)
	}()
	select {
	case <-waiterDone:
	case <-time.After(time.Second):
		t.Fatal("dispatcher Close did not invalidate the start generation")
	}
	select {
	case <-closeDone:
		t.Fatal("dispatcher Close returned before the reservation owner quiesced")
	case <-time.After(5100 * time.Millisecond):
		// The historical bounded join returned at five seconds and permitted a
		// late owner write. Strong Close must remain blocked beyond that boundary.
	}

	lateStatePath := filepath.Join(t.TempDir(), "late-publish")
	if _, published := d.sessionManager.PublishStart(
		reservation,
		&interactiveSession{sessionKey: "agent-close-owner", agentID: "agent-close-owner"},
		&agent.Process{AgentID: "agent-close-owner"},
		"late",
		false,
	); published {
		require.NoError(t, os.WriteFile(lateStatePath, []byte("late"), 0o600))
	}
	d.sessionManager.FinishStart(reservation)
	select {
	case <-closeDone:
	case <-time.After(time.Second):
		t.Fatal("dispatcher Close did not return after the reservation owner finished")
	}
	require.NoFileExists(t, lateStatePath, "an invalidated reservation must not publish or perform its guarded late write")
}

func TestSessionManagerCloseRacingBeginStartNeverAllowsLatePublish(t *testing.T) {
	for i := 0; i < 100; i++ {
		manager := NewSessionManager()
		key := fmt.Sprintf("agent-close-race-%d", i)
		start := make(chan struct{})
		var (
			reservation *sessionStartReservation
			waitCh      <-chan struct{}
			owns        bool
			wg          sync.WaitGroup
		)
		wg.Add(2)
		go func() {
			defer wg.Done()
			<-start
			reservation, waitCh, owns = manager.BeginStart(key, nil)
		}()
		go func() {
			defer wg.Done()
			<-start
			manager.Close()
		}()
		close(start)
		wg.Wait()
		require.True(t, manager.Closed())
		select {
		case <-waitCh:
		default:
			t.Fatalf("iteration %d left a start waiter open after Close", i)
		}
		if owns {
			require.NotNil(t, reservation)
			_, published := manager.PublishStart(reservation, &interactiveSession{}, &agent.Process{}, "late", false)
			require.False(t, published)
			manager.FinishStart(reservation)
		} else {
			require.Nil(t, reservation)
		}
	}
}

func TestSessionStartPublicationRetiresExpectedReplacement(t *testing.T) {
	manager := NewSessionManager()
	key := "agent-replaced"
	oldSession := &interactiveSession{sessionKey: key, agentID: key}
	manager.Set(key, oldSession)

	reservation, waitCh, owns := manager.BeginStart(key, oldSession)
	require.True(t, owns)
	newSession := &interactiveSession{sessionKey: key, agentID: key}
	process := &agent.Process{AgentID: key}
	replaced, published := manager.PublishStart(reservation, newSession, process, "new-generation", true)
	require.True(t, published)
	require.Same(t, oldSession, replaced)
	require.True(t, oldSession.retired.Load())
	current, exists := manager.Get(key)
	require.True(t, exists)
	require.Same(t, newSession, current)
	require.True(t, newSession.processUsesWorkspaceRootMemory)
	newSession.mu.Lock()
	require.Same(t, process, newSession.process)
	require.Equal(t, "new-generation", newSession.runtimeStartFingerprint)
	newSession.mu.Unlock()

	manager.FinishStart(reservation)
	select {
	case <-waitCh:
	case <-time.After(time.Second):
		t.Fatal("published start generation did not finish")
	}
}

func TestSessionManagerLookupUsesAgentScopedSessionKey(t *testing.T) {
	d := mustNewDispatcher("")
	bindingA := conversationBinding{channelID: "chn", conversationKey: "thread-a"}
	bindingB := conversationBinding{channelID: "chn", conversationKey: "thread-b"}
	keyA := sessionKeyForBinding("agent-a", bindingA).String()
	keyB := sessionKeyForBinding("agent-a", bindingB).String()
	require.Equal(t, "agent-a", keyA)
	require.Equal(t, keyA, keyB)
	d.sessionManager.Set(keyA, &interactiveSession{sessionKey: keyA, agentID: "agent-a", binding: bindingB, startedAt: time.Now()})

	snap, ok := d.Session(keyA)
	require.True(t, ok)
	require.Equal(t, keyA, snap.SessionKey)
	require.Equal(t, "thread-b", snap.ConversationKey)

	latest, ok := d.Session("agent-a")
	require.True(t, ok)
	require.Equal(t, keyA, latest.SessionKey)
}

func TestAgentScopedSessionWorkDirUsesAgentRoot(t *testing.T) {
	configDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, configDir)

	key := sessionKeyForBinding("agent-a", conversationBinding{
		channelID:       "chn",
		conversationKey: "#agentteam:thread-a",
	}).String()
	require.Equal(t, "agent-a", key)
	require.Equal(t, filepath.Join(configDir, "agents", "agent-a"), sessionWorkDirForKey("agent-a", key, false))
	require.False(t, strings.Contains(sessionWorkDirForKey("agent-a", key, false), string(os.PathSeparator)+"sessions"+string(os.PathSeparator)))
}

func TestAgentRootMemorySeedsFromLatestLegacySession(t *testing.T) {
	configDir := t.TempDir()
	t.Setenv(profile.ConfigDirEnvKey, configDir)

	agentRoot := filepath.Join(configDir, "agents", "agent-a")
	oldSession := filepath.Join(agentRoot, "sessions", "old")
	newSession := filepath.Join(agentRoot, "sessions", "new")
	require.NoError(t, os.MkdirAll(oldSession, 0o755))
	require.NoError(t, os.MkdirAll(newSession, 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(oldSession, "MEMORY.md"), []byte("# Old\n"), 0o644))
	require.NoError(t, os.WriteFile(filepath.Join(newSession, "MEMORY.md"), []byte("# New\n"), 0o644))
	require.NoError(t, os.WriteFile(filepath.Join(newSession, sessionResumeStateFilename), []byte(`{"lastSessionId":"sess-new"}`), 0o644))
	require.NoError(t, os.MkdirAll(filepath.Join(newSession, "state"), 0o755))
	require.NoError(t, os.WriteFile(filepath.Join(newSession, "state", "active-task.json"), []byte(`{"currentObjective":"new"}`), 0o644))
	oldTime := time.Now().Add(-time.Hour)
	newTime := time.Now()
	require.NoError(t, os.Chtimes(filepath.Join(oldSession, "MEMORY.md"), oldTime, oldTime))
	require.NoError(t, os.Chtimes(filepath.Join(newSession, "MEMORY.md"), newTime, newTime))

	seedSessionMemoryFromAgentRoot("agent-a", agentRoot)

	content, err := os.ReadFile(filepath.Join(agentRoot, "MEMORY.md"))
	require.NoError(t, err)
	require.Equal(t, "# New\n", string(content))
	resumeBytes, err := os.ReadFile(filepath.Join(agentRoot, sessionResumeStateFilename))
	require.NoError(t, err)
	require.JSONEq(t, `{"lastSessionId":"sess-new"}`, string(resumeBytes))
	stateBytes, err := os.ReadFile(filepath.Join(agentRoot, "state", "active-task.json"))
	require.NoError(t, err)
	require.JSONEq(t, `{"currentObjective":"new"}`, string(stateBytes))
}