package agent

import (
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestResumeSafetyHooksIgnoreProcessWithoutDriver(t *testing.T) {
	ap := &Process{SessionID: "thread-test-hook"}
	require.NotPanics(t, func() {
		ap.markResumeSafety(true)
		ap.SetOnResumeSafety(func(uint64, bool, string) error { return nil })
	})
}

func TestSetOnResumeSafetyReplaysTransitionObservedBeforeCallbackInstall(t *testing.T) {
	driver := &idleCompactionProcessStubDriver{}
	ap := &Process{
		SessionID: "thread-fast",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
		},
	}

	// This models SessionInit -> TurnEnd racing through parseStdout before
	// Dispatcher.launch has installed its durable-state callback.
	ap.markResumeSafety(false)
	ap.markResumeSafety(true)

	var got []bool
	var versions []uint64
	ap.SetOnResumeSafety(func(version uint64, safe bool, sessionID string) error {
		require.Equal(t, "thread-fast", sessionID)
		versions = append(versions, version)
		got = append(got, safe)
		return nil
	})

	require.Equal(t, []bool{true}, got)
	require.Len(t, versions, 1)
	require.NotZero(t, versions[0])
}

func TestResumeSafetyReplayCannotOverwriteNewerTransition(t *testing.T) {
	driver := &idleCompactionProcessStubDriver{}
	ap := &Process{
		SessionID: "thread-race",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
		},
	}
	ap.markResumeSafety(false)

	falseEntered := make(chan struct{})
	releaseFalse := make(chan struct{})
	setDone := make(chan struct{})
	markDone := make(chan struct{})
	var gotMu sync.Mutex
	var got []bool
	go func() {
		ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
			if !safe {
				close(falseEntered)
				<-releaseFalse
			}
			gotMu.Lock()
			got = append(got, safe)
			gotMu.Unlock()
			return nil
		})
		close(setDone)
	}()
	<-falseEntered
	go func() {
		ap.markResumeSafety(true)
		close(markDone)
	}()

	select {
	case <-markDone:
		t.Fatal("newer resume-safety callback bypassed serialized replay")
	case <-time.After(20 * time.Millisecond):
	}
	close(releaseFalse)
	<-setDone
	<-markDone

	gotMu.Lock()
	defer gotMu.Unlock()
	require.Equal(t, []bool{false, true}, got)
}

func TestIdleCompactionMakesResumeSafetyUnsafeAgain(t *testing.T) {
	driver := &idleCompactionProcessStubDriver{started: true}
	ap := &Process{
		SessionID: "thread-compacting",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
		},
	}

	var got []bool
	ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
		got = append(got, safe)
		return nil
	})
	ap.markResumeSafety(true)
	require.True(t, ap.startIdleCompactionAfterTurnEnd())

	require.Equal(t, []bool{true, false}, got)
	require.True(t, ap.resumeSafetyKnown)
	require.False(t, ap.resumeSafety)
}

func TestIdleCompactionTransportEOFAfterDispatchStaysUnsafeUntilRuntimeIdle(t *testing.T) {
	writer := &codexTransportSignalWriter{wrote: make(chan struct{})}
	transportFailure := &codexAppServerTransportFailure{}
	driver := &CodexAppServerDriver{
		launchGeneration: 7,
		codexRPCState: codexRPCState{
			pending:          make(map[int]pendingAppServerRequest),
			serverRequests:   make(map[string]inFlightServerRequest),
			rawStdin:         writer,
			transportFailure: transportFailure,
		},
		codexThreadState: codexThreadState{
			threadID:      "thread-compacting",
			threadStatus:  appServerThreadStatus{Type: "idle"},
			compactQueued: true,
		},
	}
	ap := &Process{
		AgentID:   "agent-compaction-ambiguous",
		SessionID: "thread-compacting",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
		},
		processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
	}

	var got []bool
	require.NoError(t, ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
		got = append(got, safe)
		return nil
	}))
	require.NoError(t, ap.markResumeSafety(true))

	started := make(chan bool, 1)
	go func() { started <- ap.startIdleCompactionAfterTurnEnd() }()
	<-writer.wrote // the thread/compact/start frame crossed the dispatch boundary
	driver.failTransport(7, transportFailure, "injected EOF after compaction dispatch")

	require.True(t, <-started, "ambiguous dispatch must conservatively hold the Process in compacting state")
	require.Equal(t, []bool{true, false}, got)
	require.True(t, ap.resumeSafetyKnown)
	require.False(t, ap.resumeSafety, "lost compaction response must not restore durable resume safety")
	require.False(t, ap.IsIdle())
	driver.mu.Lock()
	require.True(t, driver.compactionInProgress)
	require.False(t, driver.compactQueued)
	driver.mu.Unlock()

	// Only an authoritative runtime boundary may settle the ambiguous mutation.
	require.True(t, driver.updateThreadStatus("thread-compacting", appServerThreadStatus{Type: "idle"}))
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:      consts.ParsedEventKindRuntimeIdle,
		SessionID: "thread-compacting",
	})
	require.True(t, ap.IsIdle())
	require.True(t, ap.resumeSafety)
	require.Equal(t, []bool{true, false, true}, got)
}

func TestIdleCompactionKnownPreDispatchRejectionRestoresResumeSafety(t *testing.T) {
	driver := &idleCompactionProcessStubDriver{startErr: errCodexAppServerRequestDispatchRejected}
	ap := &Process{
		AgentID:   "agent-compaction-rejected",
		SessionID: "thread-compaction-rejected",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
		},
		processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
	}

	var got []bool
	require.NoError(t, ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
		got = append(got, safe)
		return nil
	}))
	require.NoError(t, ap.markResumeSafety(true))

	require.False(t, ap.startIdleCompactionAfterTurnEnd())
	require.Equal(t, []bool{true, false, true}, got)
	require.True(t, ap.IsIdle())
	require.True(t, ap.resumeSafety, "a known-unwritten request leaves the prior idle thread safe to resume")
}

func TestNoPromptSessionInitSafetyIsSupersededByOwnershipStoreRestore(t *testing.T) {
	newProcess := func() *Process {
		return &Process{
			AgentID: "agent-no-prompt",
			processRuntimeBinding: processRuntimeBinding{
				driver: &protocolRecoveryStubDriver{},
			},
			processLifecycleState: processLifecycleState{state: processStateRunning, isIdle: true},
		}
	}
	newHandoffStore := func() *memoryDeliveryOwnershipStore {
		return &memoryDeliveryOwnershipStore{records: map[string]DeliveryOwnershipRecord{
			"msg-1": {
				Message: types.InboxMessage{MessageID: "msg-1", Target: "#ops", Content: "queued"},
				Input:   drivercontract.RuntimeInput{ID: "msg-1"},
				State:   DeliveryOwnershipQueued,
			},
		}}
	}

	t.Run("callback before ownership store", func(t *testing.T) {
		ap := newProcess()
		var got []bool
		ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error { got = append(got, safe); return nil })
		ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindSessionInit, SessionID: "thread-1"})
		require.Equal(t, []bool{false, true}, got, "empty no-prompt initialization must establish an idle safe boundary")

		require.NoError(t, ap.SetDeliveryOwnershipStoreForHandoff(newHandoffStore(), "msg-1"))
		require.Equal(t, []bool{false, true, false}, got, "restored live ownership must supersede the provisional idle-safe snapshot")
	})

	t.Run("callback after ownership store", func(t *testing.T) {
		ap := newProcess()
		ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindSessionInit, SessionID: "thread-1"})
		require.True(t, ap.resumeSafety)
		require.NoError(t, ap.SetDeliveryOwnershipStoreForHandoff(newHandoffStore(), "msg-1"))
		require.False(t, ap.resumeSafety)

		var got []bool
		ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error { got = append(got, safe); return nil })
		require.Equal(t, []bool{false}, got, "callback replay must observe the post-restore unsafe version")
	})
}