package agent

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func TestPiRPCReconcileUsesLiveRuntimeStateAndEntries(t *testing.T) {
	t.Run("replacement idle without local correlation stays unknown", func(t *testing.T) {
		binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
			piFakeStateScript(false),
			piFakeEntriesScript(),
		})
		d := &PiRPCDriver{}
		result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
		require.NoError(t, err)
		t.Cleanup(result.Cleanup)

		got, err := d.ReconcileRuntimeInput(t.Context(), "owned-by-retired-process", result.SessionID)
		require.NoError(t, err)
		require.Equal(t, RuntimeInputReconciliationUnknown, got,
			"a replacement process cannot manufacture negative evidence for the retired process")
	})

	t.Run("idle without marker is authoritative absent", func(t *testing.T) {
		binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
			piFakePromptScript(),
			piFakeStateScript(false),
			piFakeEntriesScript(),
		})
		d := &PiRPCDriver{}
		result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
		require.NoError(t, err)
		t.Cleanup(result.Cleanup)

		require.NoError(t, d.DeliverRuntimeInput(t.Context(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "idle-absent"}, "accepted but never acked"), result.SessionID))
		got, err := d.ReconcileRuntimeInput(t.Context(), "idle-absent", result.SessionID)
		require.NoError(t, err)
		require.Equal(t, RuntimeInputReconciliationPending, got,
			"the first idle snapshot sequences RuntimeIdle before durable replay")
		got, err = d.ReconcileRuntimeInput(t.Context(), "idle-absent", result.SessionID)
		require.NoError(t, err)
		require.Equal(t, RuntimeInputReconciliationAbsent, got)
		d.mu.Lock()
		require.False(t, d.rootRunActive, "authoritative idle retires the phantom preflight generation")
		require.Empty(t, d.promptInputID)
		d.mu.Unlock()
	})

	t.Run("streaming without marker remains pending", func(t *testing.T) {
		binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
			piFakePromptScript(),
			piFakeStateScript(true),
			piFakeEntriesScript(),
		})
		d := &PiRPCDriver{}
		result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
		require.NoError(t, err)
		t.Cleanup(result.Cleanup)

		require.NoError(t, d.DeliverRuntimeInput(t.Context(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "active-pending"}, "still active"), result.SessionID))
		got, err := d.ReconcileRuntimeInput(t.Context(), "active-pending", result.SessionID)
		require.NoError(t, err)
		require.Equal(t, RuntimeInputReconciliationPending, got)
		d.mu.Lock()
		require.True(t, d.rootRunActive)
		d.mu.Unlock()
	})

	t.Run("idle in-memory marker is authoritative acked", func(t *testing.T) {
		marker := "delivered\n\n<!-- agentteam-runtime-input:idle-acked -->"
		entries := piFakeEntriesScript()
		entries.Data = map[string]any{
			"entries": []any{map[string]any{
				"type":    "message",
				"message": map[string]any{"role": "user", "content": marker},
			}},
			"leafId": "entry-user",
		}
		binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
			piFakePromptScript(),
			piFakeStateScript(false),
			entries,
		})
		d := &PiRPCDriver{}
		result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
		require.NoError(t, err)
		t.Cleanup(result.Cleanup)

		require.NoError(t, d.DeliverRuntimeInput(t.Context(), testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "idle-acked"}, "delivered"), result.SessionID))
		got, err := d.ReconcileRuntimeInput(t.Context(), "idle-acked", result.SessionID)
		require.NoError(t, err)
		require.Equal(t, RuntimeInputReconciliationPending, got,
			"the first idle snapshot sequences RuntimeIdle before ownership ack")
		got, err = d.ReconcileRuntimeInput(t.Context(), "idle-acked", result.SessionID)
		require.NoError(t, err)
		require.Equal(t, RuntimeInputReconciliationAcked, got)
		d.mu.Lock()
		require.False(t, d.rootRunActive, "authoritative idle also repairs a lost agent_settled event")
		d.mu.Unlock()
	})
}

func TestPiRPCIdleStartupSteerReconcilesAndReplayUsesPrompt(t *testing.T) {
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		piFakeSteerScript(),
		piFakeStateScript(false),
		piFakeEntriesScript(),
		piFakeStateScript(false),
		piFakeEntriesScript(),
		piFakePromptScript(),
	})
	d := &PiRPCDriver{}
	result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
	require.NoError(t, err)
	t.Cleanup(result.Cleanup)

	input := testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "startup-steer"}, "hello")
	require.NoError(t, d.DeliverAckedBusyRuntimeInput(t.Context(), input, result.SessionID))
	// A duplicate daemon admission must join the retained correlation instead of
	// resending the steer or exhausting the transient inbox retry budget.
	require.NoError(t, d.DeliverAckedBusyRuntimeInput(t.Context(), input, result.SessionID))

	got, err := d.ReconcileRuntimeInput(t.Context(), input.ID, result.SessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationPending, got,
		"the first idle snapshot starts the late-ack grace window")
	d.mu.Lock()
	require.True(t, d.steerInFlight[input.ID], "one racy idle snapshot must retain the steer correlation")
	d.startupSteerIdleObservedAt[input.ID] = time.Now().Add(-piStartupSteerIdleAbsentGrace)
	d.mu.Unlock()

	got, err = d.ReconcileRuntimeInput(t.Context(), input.ID, result.SessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationPending, got,
		"a stable delayed snapshot publishes RuntimeIdle before replay")
	d.mu.Lock()
	require.False(t, d.steerInFlight[input.ID], "confirmed idle retires stale steer correlation")
	d.mu.Unlock()

	got, err = d.ReconcileRuntimeInput(t.Context(), input.ID, result.SessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationAbsent, got)

	// The replay is now an idle root prompt. DeliverRuntimeInput clears the
	// negative fence before its one new mutation generation.
	require.NoError(t, d.DeliverRuntimeInput(t.Context(), input, result.SessionID))
	d.mu.Lock()
	require.True(t, d.rootRunActive)
	require.Equal(t, input.ID, d.promptInputID)
	d.mu.Unlock()
}

func TestPiRPCIdleStartupSteerLateMarkerPreventsReplay(t *testing.T) {
	marker := "hello\n\n<!-- agentteam-runtime-input:startup-steer-late -->"
	lateEntries := piFakeEntriesScript()
	lateEntries.Data = map[string]any{
		"entries": []any{map[string]any{
			"type":    "message",
			"message": map[string]any{"role": "user", "content": marker},
		}},
		"leafId": "entry-user-late",
	}
	binary, _ := writeFakePiRPCModeBinary(t, []piFakeScript{
		piFakeSteerScript(),
		piFakeStateScript(false),
		piFakeEntriesScript(),
		piFakeStateScript(false),
		lateEntries,
	})
	d := &PiRPCDriver{}
	result, err := d.Spawn(t.Context(), piTestConfig(t, binary))
	require.NoError(t, err)
	t.Cleanup(result.Cleanup)

	input := testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "startup-steer-late"}, "hello")
	require.NoError(t, d.DeliverAckedBusyRuntimeInput(t.Context(), input, result.SessionID))

	got, err := d.ReconcileRuntimeInput(t.Context(), input.ID, result.SessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationPending, got)
	d.mu.Lock()
	require.True(t, d.steerInFlight[input.ID])
	d.mu.Unlock()

	got, err = d.ReconcileRuntimeInput(t.Context(), input.ID, result.SessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationPending, got,
		"the late marker wins and RuntimeIdle is sequenced before acked ownership")

	got, err = d.ReconcileRuntimeInput(t.Context(), input.ID, result.SessionID)
	require.NoError(t, err)
	require.Equal(t, RuntimeInputReconciliationAcked, got)
	d.mu.Lock()
	require.Empty(t, d.startupSteerIdleObservedAt)
	d.mu.Unlock()
}