package daemon

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/workflow/agentruntime"
)

// newOrderingClient builds a clientState with the maps the ordering tests need.
func newOrderingClient(machineID string) *clientState {
	return &clientState{
		machineID:             machineID,
		runningAgents:         map[string]bool{},
		startingAgents:        map[string]bool{},
		stoppedAgents:         map[string]bool{},
		agentStatuses:         map[string]string{},
		agentActivities:       map[string]agentActivityState{},
		agentSessions:         map[string]string{},
		agentTurnStateRecords: map[string]agentTurnStateRecord{},
		connectedAt:           time.Now().UTC(),
	}
}

// wm returns a pointer to an int64 watermark, matching the presence-aware
// replaceReadyAgentTurnStates signature (non-nil = modern ordering authority).
func wm(v int64) *int64 { return &v }

func orderingReader(client *clientState) *agentruntime.HubWorkingReader {
	hub := &Hub{clients: map[string]*clientState{client.machineID: client}}
	return agentruntime.NewHubWorkingReader(fakeAgentMachineResolver{machineID: client.machineID}, hub)
}

// Reorder direction 1: a NEWER active/queued increment arriving BEFORE an OLDER
// idle snapshot must survive the snapshot merge. The per-agent revision keeps the
// newer increment; active/queued can never become recoverable.
func TestTurnStateOrderingNewerIncrementSurvivesOlderSnapshot(t *testing.T) {
	client := newOrderingClient("mac_1")
	client.setAgentSession("agt_1", "sess_1")
	reader := orderingReader(client)

	// A newer active increment (revision 5) arrives first.
	client.setAgentTurnState("agt_1", "active", "sess_1", 5, "")
	working, known := reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "active increment must be working")

	// An OLDER idle snapshot (watermark 3) arrives late. It must NOT overwrite the
	// newer increment: agt_1 stays active (working), never recoverable.
	client.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "idle"}, wm(3))
	working, known = reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "an older idle snapshot must not overwrite a newer active increment")
	require.Equal(t, "active", client.agentTurnStateRecords["agt_1"].State)
	require.Equal(t, int64(5), client.agentTurnStateRecords["agt_1"].Revision)
}

// Reorder direction 2: a NEWER active/queued snapshot arriving BEFORE a delayed
// OLDER idle increment must win. The delayed increment (revision <= high-water)
// is dropped, so active/queued can never become recoverable.
func TestTurnStateOrderingNewerSnapshotBeatsDelayedOlderIncrement(t *testing.T) {
	client := newOrderingClient("mac_1")
	client.setAgentSession("agt_1", "sess_1")
	reader := orderingReader(client)

	// A newer queued snapshot (watermark 5) arrives first.
	client.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "queued"}, wm(5))
	working, known := reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "queued snapshot must be working")

	// A delayed OLDER idle increment (revision 2, below high-water 5) arrives. It
	// must be dropped.
	client.setAgentTurnState("agt_1", "idle", "sess_1", 2, "")
	working, known = reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "a delayed older idle increment must not overwrite a newer queued snapshot")
	require.Equal(t, "queued", client.agentTurnStateRecords["agt_1"].State)
}

// Duplicate revision: an equal-revision increment is dropped (never overwrites
// the applied state), so the applied idle (recoverable) stays recoverable and an
// applied active stays working.
func TestTurnStateOrderingDuplicateRevisionDropped(t *testing.T) {
	client := newOrderingClient("mac_1")
	client.setAgentSession("agt_1", "sess_1")
	reader := orderingReader(client)

	client.setAgentTurnState("agt_1", "idle", "sess_1", 5, "")
	working, known := reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working, "idle is recoverable")

	// A duplicate revision 5 with a different state is dropped.
	client.setAgentTurnState("agt_1", "active", "sess_1", 5, "")
	working, known = reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working, "an equal-revision increment must be dropped; idle stays recoverable")
	require.Equal(t, "idle", client.agentTurnStateRecords["agt_1"].State)
}

// Per-agent merge: an older snapshot must not overwrite a NEWER per-agent
// increment for that agent, but its still-current entries for OTHER agents remain
// materializable.
func TestTurnStateOrderingPerAgentMerge(t *testing.T) {
	client := newOrderingClient("mac_1")
	client.setAgentSession("agt_a", "sess_a")
	client.setAgentSession("agt_b", "sess_b")

	// Newer increments for A (active, rev 5) and B (idle, rev 6).
	client.setAgentTurnState("agt_a", "active", "sess_a", 5, "")
	client.setAgentTurnState("agt_b", "idle", "sess_b", 6, "")

	// An OLDER snapshot (watermark 3) with A=idle, B=queued arrives late.
	client.replaceReadyAgentTurnStates([]string{"agt_a", "agt_b"}, map[string]string{
		"agt_a": "idle",
		"agt_b": "queued",
	}, wm(3))

	// A: rev 5 > watermark 3 -> older snapshot must not overwrite; stays active.
	require.Equal(t, "active", client.agentTurnStateRecords["agt_a"].State, "older snapshot must not overwrite a newer per-agent increment")
	// B: rev 6 > watermark 3 -> older snapshot must not overwrite; stays idle.
	require.Equal(t, "idle", client.agentTurnStateRecords["agt_b"].State, "older snapshot must not overwrite a newer per-agent increment")

	// A NEWER snapshot (watermark 7) with A=idle, B=active, C=queued wins for all.
	client.replaceReadyAgentTurnStates([]string{"agt_a", "agt_b", "agt_c"}, map[string]string{
		"agt_a": "idle",
		"agt_b": "active",
		"agt_c": "queued",
	}, wm(7))
	require.Equal(t, "idle", client.agentTurnStateRecords["agt_a"].State)
	require.Equal(t, "active", client.agentTurnStateRecords["agt_b"].State)
	require.Equal(t, "queued", client.agentTurnStateRecords["agt_c"].State)

	// A delayed OLDER snapshot (watermark 2) must not regress A/B/C.
	client.replaceReadyAgentTurnStates([]string{"agt_a", "agt_b", "agt_c"}, map[string]string{
		"agt_a": "active",
		"agt_b": "idle",
		"agt_c": "idle",
	}, wm(2))
	require.Equal(t, "idle", client.agentTurnStateRecords["agt_a"].State, "delayed older snapshot must not regress a newer state")
	require.Equal(t, "active", client.agentTurnStateRecords["agt_b"].State)
	require.Equal(t, "queued", client.agentTurnStateRecords["agt_c"].State)
}

// New connection reset: a fresh clientState has no per-agent revisions, so the
// first snapshot (even at base watermark 0) materializes and any carry-over
// ordering from the previous connection is gone.
func TestTurnStateOrderingNewConnectionResetsOrdering(t *testing.T) {
	// Old connection had a high watermark.
	oldClient := newOrderingClient("mac_1")
	oldClient.setAgentSession("agt_1", "sess_1")
	oldClient.setAgentTurnState("agt_1", "active", "sess_1", 100, "")

	// New connection (fresh clientState) resets ordering: base watermark snapshot
	// materializes.
	newClient := newOrderingClient("mac_1")
	newClient.setAgentSession("agt_1", "sess_1")
	reader := orderingReader(newClient)
	newClient.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "idle"}, wm(0))
	working, known := reader.RuntimeWorking(context.Background(), "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working, "a fresh connection must materialize the current snapshot, not inherit the old connection's ordering")
	require.Equal(t, int64(0), newClient.agentTurnStateRecords["agt_1"].Revision)
	_ = oldClient
}

// Wire path: the production router applies the same-connection ordering. A frame
// with a stale revision is dropped; a frame with a newer revision applies.
func TestTurnStateOrderingWirePathRevisionFence(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{clients: map[string]*clientState{}}
	client := newOrderingClient("mac_1")
	client.sessionGeneration = "epoch_cur"
	hub.clients["mac_1"] = client
	client.setAgentSession("agt_1", "sess_1")

	route := func(frame string) error {
		return hub.handleMessage(ctx, client, []byte(frame))
	}

	// Newer queued increment (revision 5, current epoch/session).
	require.NoError(t, route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"queued","sessionId":"sess_1","epoch":"epoch_cur","revision":5}`))
	state, ok := client.agentTurnState("agt_1")
	require.True(t, ok)
	require.Equal(t, "queued", state)

	// A stale-epoch frame (even with a newer revision) is dropped by the epoch fence.
	require.NoError(t, route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1","epoch":"epoch_old","revision":6}`))
	state, _ = client.agentTurnState("agt_1")
	require.Equal(t, "queued", state, "a stale-epoch frame must be dropped regardless of revision")

	// A stale-session frame (matching epoch but old session) is dropped.
	require.NoError(t, route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_old","epoch":"epoch_cur","revision":6}`))
	state, _ = client.agentTurnState("agt_1")
	require.Equal(t, "queued", state, "a stale-session frame must be dropped")

	// An older-revision frame (revision 4 < high-water 5) is dropped by the
	// ordering fence even though epoch/session match.
	require.NoError(t, route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1","epoch":"epoch_cur","revision":4}`))
	state, _ = client.agentTurnState("agt_1")
	require.Equal(t, "queued", state, "an older-revision increment must not overwrite a newer one")

	// A newer revision applies.
	require.NoError(t, route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1","epoch":"epoch_cur","revision":7}`))
	state, _ = client.agentTurnState("agt_1")
	require.Equal(t, "idle", state)
}

// Concurrent ready/event publication under race: many goroutines concurrently
// publish snapshots and increments; the reader must always observe a coherent
// state and never race (run under go test -race). A concurrent diagnostics query
// must also never hang (self-deadlock regression).
func TestTurnStateOrderingConcurrentPublicationUnderRace(t *testing.T) {
	client := newOrderingClient("mac_1")
	client.setAgentSession("agt_1", "sess_1")
	hub := &Hub{clients: map[string]*clientState{client.machineID: client}}
	reader := agentruntime.NewHubWorkingReader(fakeAgentMachineResolver{machineID: "mac_1"}, hub)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func(worker int) {
			defer wg.Done()
			for j := 0; j < 100; j++ {
				revision := int64(worker*1000 + j + 1)
				switch j % 3 {
				case 0:
					client.setAgentTurnState("agt_1", "active", "sess_1", revision, "")
				case 1:
					client.setAgentTurnState("agt_1", "idle", "sess_1", revision, "")
				default:
					client.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "queued"}, wm(revision))
				}
			}
		}(i)
	}
	// Concurrent diagnostics reads must complete without hanging (deadlock).
	for i := 0; i < 4; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 100; j++ {
				_ = hub.AgentRuntimeDiagnostics("mac_1")
				_, _ = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
			}
		}()
	}
	wg.Wait()

	// The final observed state is one of the applied states; the reader must not
	// panic and the high-water must be monotonic.
	_, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.GreaterOrEqual(t, client.turnStateHighWater, int64(1))
	require.NotEmpty(t, client.agentTurnStateRecords["agt_1"].State)
}

// Diagnostics deadlock regression: AgentRuntimeDiagnostics must not self-deadlock
// (it reads every clientState field under client.mu and must not call a method
// that re-locks client.mu). A bounded call completes.
func TestAgentRuntimeDiagnosticsDoesNotSelfDeadlock(t *testing.T) {
	client := newOrderingClient("mac_1")
	client.sessionGeneration = "epoch_cur"
	client.setAgentSession("agt_1", "sess_1")
	client.setAgentTurnState("agt_1", "idle", "sess_1", 5, "")
	hub := &Hub{clients: map[string]*clientState{client.machineID: client}}

	done := make(chan map[string]any, 1)
	go func() {
		done <- hub.AgentRuntimeDiagnostics("mac_1")
	}()
	select {
	case diagnostics := <-done:
		require.Equal(t, true, diagnostics["connected"])
		require.Equal(t, "epoch_cur", diagnostics["sessionGeneration"])
		require.Equal(t, "idle", diagnostics["turnStates"].(map[string]string)["agt_1"])
		require.Equal(t, int64(5), diagnostics["turnStateRevisions"].(map[string]int64)["agt_1"])
	case <-time.After(5 * time.Second):
		require.Fail(t, "AgentRuntimeDiagnostics self-deadlocked")
	}
}