package dispatcher

import (
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestPersistResumeSafetyRejectsDelayedCallbackFromOlderProcessEpoch(t *testing.T) {
	workDir := t.TempDir()
	oldEpoch, err := persistStartConfigState(workDir, "codex", "http://server.test", "thread-old", &types.AgentStartConfig{
		Runtime:   "codex",
		ServerURL: "http://server.test",
		SessionID: "thread-old",
	})
	require.NoError(t, err)
	require.NotZero(t, oldEpoch)

	releaseOldWrite := make(chan struct{})
	oldWriteDone := make(chan struct{})
	var oldApplied bool
	var oldErr error
	go func() {
		<-releaseOldWrite
		oldApplied, oldErr = PersistResumeSafety(workDir, "http://server.test", "thread-old", oldEpoch, 7, true)
		close(oldWriteDone)
	}()

	// The replacement Process persists its unsafe baseline first. The delayed
	// terminal callback from the previous Process then lands afterward and must
	// be rejected by the epoch CAS rather than resurrecting safe=true.
	newEpoch, err := persistStartConfigState(workDir, "codex", "http://server.test", "thread-new", &types.AgentStartConfig{
		Runtime:   "codex",
		ServerURL: "http://server.test",
		SessionID: "thread-new",
	})
	require.NoError(t, err)
	require.Greater(t, newEpoch, oldEpoch)
	close(releaseOldWrite)
	<-oldWriteDone
	require.NoError(t, oldErr)
	require.False(t, oldApplied)

	state, err := loadSessionResumeState(workDir)
	require.NoError(t, err)
	require.Equal(t, newEpoch, state.SessionResumeSafetyEpoch)
	require.Zero(t, state.SessionResumeSafetyVersion)
	require.False(t, state.SessionResumeSafe)
	require.Equal(t, "thread-new", state.LastSessionID)
}

func TestPersistResumeSafetyRejectsOutOfOrderVersionWithinProcessEpoch(t *testing.T) {
	workDir := t.TempDir()
	epoch, err := persistStartConfigState(workDir, "codex", "http://server.test", "thread-1", &types.AgentStartConfig{
		Runtime:   "codex",
		ServerURL: "http://server.test",
		SessionID: "thread-1",
	})
	require.NoError(t, err)

	applied, err := PersistResumeSafety(workDir, "http://server.test", "thread-1", epoch, 2, false)
	require.NoError(t, err)
	require.True(t, applied)

	var wg sync.WaitGroup
	wg.Add(1)
	var staleApplied bool
	go func() {
		defer wg.Done()
		staleApplied, err = PersistResumeSafety(workDir, "http://server.test", "thread-1", epoch, 1, true)
	}()
	wg.Wait()
	require.NoError(t, err)
	require.False(t, staleApplied)

	state, err := loadSessionResumeState(workDir)
	require.NoError(t, err)
	require.Equal(t, uint64(2), state.SessionResumeSafetyVersion)
	require.False(t, state.SessionResumeSafe)
}

func TestPersistResumeSafetyFailsClosedAroundLiveOwnershipInEitherOrder(t *testing.T) {
	for _, callbackFirst := range []bool{false, true} {
		name := "ownership_before_safe_callback"
		if callbackFirst {
			name = "safe_callback_before_ownership"
		}
		t.Run(name, func(t *testing.T) {
			workDir := t.TempDir()
			epoch, err := persistStartConfigState(workDir, "codex", "http://server.test", "thread-1", &types.AgentStartConfig{
				Runtime:   "codex",
				ServerURL: "http://server.test",
				SessionID: "thread-1",
			})
			require.NoError(t, err)

			store := newSessionDeliveryOwnershipStore(workDir)
			putOwner := func() {
				_, putErr := store.PutDeliveryOwnership(agent.DeliveryOwnershipRecord{
					Message:    types.InboxMessage{MessageID: "q1", Target: "#ops", Content: "accepted"},
					Input:      testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "q1"}, "accepted"),
					State:      agent.DeliveryOwnershipQueued,
					AcceptedAt: time.Now().UTC(),
				})
				require.NoError(t, putErr)
			}
			persistSafe := func() bool {
				applied, persistErr := PersistResumeSafety(workDir, "http://server.test", "thread-1", epoch, 1, true)
				require.NoError(t, persistErr)
				return applied
			}

			if callbackFirst {
				require.True(t, persistSafe())
				putOwner()
			} else {
				putOwner()
				require.False(t, persistSafe())
			}

			state, err := loadSessionResumeState(workDir)
			require.NoError(t, err)
			require.False(t, state.SessionResumeSafe)
			require.Equal(t, uint64(1), state.SessionResumeSafetyVersion)
			require.Len(t, state.RuntimeDeliveryOwnership, 1)

			if !callbackFirst {
				// The rejected snapshot's version is consumed. Removing the owner
				// cannot make an equal-version retry authoritative later.
				require.NoError(t, store.DeleteDeliveryOwnership("q1"))
				require.False(t, persistSafe())
				state, err = loadSessionResumeState(workDir)
				require.NoError(t, err)
				require.False(t, state.SessionResumeSafe)
			}
		})
	}
}