package daemon

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

// Presence-aware empty-snapshot watermark regressions. A modern Desktop Protocol
// ready always carries runningAgentTurnStatesWatermark (even when turn states are
// empty or no agents run), so an explicit empty current snapshot can retire prior
// per-agent working state; a legacy ready that omitted the field must not claim
// modern ordering authority. All assertions read the real Hub state/revisions/
// diagnostics and the real HubWorkingReader.

// (a) Prior active revision 5, then an explicit current watermark 6 with the
// agent still running but NO turn-state entry: the empty snapshot clears the
// stale working state; the responder enters the bounded grace (known=false) and
// becomes recoverable after the grace elapses.
func TestWatermarkEmptySnapshotClearsStaleActive(t *testing.T) {
	ctx := context.Background()
	now := time.Date(2026, 8, 16, 12, 0, 0, 0, time.UTC)
	client := newOrderingClient("mac_1")
	client.connectedAt = now.Add(-30 * time.Second)
	client.setAgentSession("agt_1", "sess_1")
	hub := &Hub{clients: map[string]*clientState{client.machineID: client}}
	reader := agentruntime.NewHubWorkingReaderWithOptions(
		fakeAgentMachineResolver{machineID: "mac_1"},
		hub,
		agentruntime.HubWorkingReaderOptions{Now: func() time.Time { return now }, Grace: 60 * time.Second},
	)

	// Prior active/queued working state at revision 5.
	client.setAgentTurnState("agt_1", "active", "sess_1", 5, "")
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working)

	// Explicit current empty snapshot (agent still running, no turn-state entry)
	// at watermark 6: newer than revision 5, so it clears the stale active.
	client.replaceReadyAgentTurnStates([]string{"agt_1"}, nil, wm(6))
	_, ok := client.agentTurnState("agt_1")
	require.False(t, ok, "an explicit current empty snapshot must retire the stale active state")
	require.Equal(t, int64(6), client.turnStateHighWater)
	require.True(t, client.turnStateSnapshotReceived)

	// Within the bounded grace the responder is deferred (known=false, no send).
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.False(t, known, "after the empty snapshot the responder is within the bounded grace")
	require.False(t, working)

	// After the grace elapses the responder becomes recoverable (idle).
	client.connectedAt = now.Add(-90 * time.Second)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known, "after the bounded grace an online responder with no snapshot becomes recoverable")
	require.False(t, working)
}

// (b) Explicit current watermark 6 with ZERO running agents: the empty snapshot
// retires the stale working state and the responder follows the bounded
// grace -> recoverable transition.
func TestWatermarkEmptySnapshotZeroRunningAgents(t *testing.T) {
	ctx := context.Background()
	now := time.Date(2026, 8, 16, 12, 0, 0, 0, time.UTC)
	client := newOrderingClient("mac_1")
	client.connectedAt = now.Add(-30 * time.Second)
	client.setAgentSession("agt_1", "sess_1")
	hub := &Hub{clients: map[string]*clientState{client.machineID: client}}
	reader := agentruntime.NewHubWorkingReaderWithOptions(
		fakeAgentMachineResolver{machineID: "mac_1"},
		hub,
		agentruntime.HubWorkingReaderOptions{Now: func() time.Time { return now }, Grace: 60 * time.Second},
	)

	client.setAgentTurnState("agt_1", "queued", "sess_1", 5, "")
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working)

	// Zero running agents + explicit watermark 6: empty snapshot retires stale.
	client.replaceReadyAgentTurnStates(nil, nil, wm(6))
	_, ok := client.agentTurnState("agt_1")
	require.False(t, ok, "an empty current snapshot with zero running agents must retire stale working state")
	require.Equal(t, int64(6), client.turnStateHighWater)
	require.True(t, client.turnStateSnapshotReceived)

	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.False(t, known, "within the bounded grace the responder defers")
	require.False(t, working)

	client.connectedAt = now.Add(-90 * time.Second)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known, "after grace the responder becomes recoverable")
	require.False(t, working)
}

// (c) A genuinely OLDER empty watermark 4 must NOT clear a revision-5
// active/queued increment.
func TestWatermarkOlderEmptySnapshotPreservesNewerActive(t *testing.T) {
	ctx := context.Background()
	client := newOrderingClient("mac_1")
	client.setAgentSession("agt_1", "sess_1")
	reader := orderingReader(client)

	// A newer active increment at revision 5.
	client.setAgentTurnState("agt_1", "active", "sess_1", 5, "")
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working)

	// An OLDER empty snapshot (watermark 4 < 5): must NOT clear the active.
	client.replaceReadyAgentTurnStates([]string{"agt_1"}, nil, wm(4))
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "an older empty snapshot must not clear a newer active/queued increment")
	require.Equal(t, "active", client.agentTurnStateRecords["agt_1"].State)
	require.Equal(t, int64(5), client.agentTurnStateRecords["agt_1"].Revision)
	// The older snapshot is not accepted: high-water and diagnostics must not
	// regress.
	require.Equal(t, int64(5), client.turnStateHighWater)
}

// (d) Presence-aware decoding: an explicit modern watermark 0 materializes a
// snapshot with modern ordering authority (agent revisions and high-water
// update), while a legacy ready that omitted the field does not claim modern
// ordering authority (agent revisions and high-water untouched).
func TestWatermarkPresenceAwareZeroVersusLegacyOmitted(t *testing.T) {
	// Modern explicit watermark 0: materializes idle at revision 0 and updates
	// the connection high-water + snapshot diagnostics.
	modern := newOrderingClient("mac_1")
	modern.setAgentSession("agt_1", "sess_1")
	modern.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "idle"}, wm(0))
	require.Equal(t, "idle", modern.agentTurnStateRecords["agt_1"].State)
	require.Equal(t, int64(0), modern.agentTurnStateRecords["agt_1"].Revision, "a modern watermark-0 snapshot claims ordering authority at revision 0")
	require.Equal(t, int64(0), modern.turnStateHighWater)
	require.True(t, modern.turnStateSnapshotReceived)

	// Legacy ready that omitted the watermark: materializes idle but does NOT
	// update agent revisions or the high-water (no modern ordering authority).
	legacy := newOrderingClient("mac_1")
	legacy.setAgentSession("agt_1", "sess_1")
	legacy.setAgentTurnState("agt_1", "active", "sess_1", 5, "")
	legacy.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "idle"}, nil)
	require.Equal(t, "idle", legacy.agentTurnStateRecords["agt_1"].State)
	require.Equal(t, int64(5), legacy.agentTurnStateRecords["agt_1"].Revision, "a legacy ready must not overwrite the stored per-agent revision")
	require.Equal(t, int64(5), legacy.turnStateHighWater, "a legacy ready must not claim modern high-water authority")

	// Consequence: a later modern snapshot at watermark 4 is older than the
	// legacy-tracked revision 5, so it cannot clear the legacy idle.
	legacy.replaceReadyAgentTurnStates([]string{"agt_1"}, nil, wm(4))
	require.Equal(t, "idle", legacy.agentTurnStateRecords["agt_1"].State, "a delayed older modern snapshot must not clear a legacy-tracked state")
}

// Watermark router regression through the production wire route: a ready frame
// carrying an explicit empty watermark decodes as a modern snapshot (pointer
// present), and one omitting the field decodes as legacy (nil), driving the
// correct merge.
func TestWatermarkRouterPresenceDecoding(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{store: &recordingStartStore{}, clients: map[string]*clientState{}}
	client := newOrderingClient("mac_1")
	client.sessionGeneration = "epoch_cur"
	hub.clients["mac_1"] = client
	client.setAgentSession("agt_1", "sess_1")
	reader := agentruntime.NewHubWorkingReader(fakeAgentMachineResolver{machineID: "mac_1"}, hub)

	route := func(frame string) {
		require.NoError(t, hub.handleMessage(ctx, client, []byte(frame)))
	}

	// Prior active at revision 5.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"active","sessionId":"sess_1","revision":5}`)
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working)

	// A modern ready frame with an explicit empty watermark 6 (agent running, no
	// turn states): the router decodes a present watermark and clears the stale
	// active via merge.
	route(`{"type":"ready","capabilities":["agent:permission-mode"],"runningAgents":["agt_1"],"runningAgentTurnStatesWatermark":6}`)
	_, ok := client.agentTurnState("agt_1")
	require.False(t, ok, "a modern empty watermark routed through the wire must retire stale working state")
	require.Equal(t, int64(6), client.turnStateHighWater)

	// A newer MODERN ready (watermark 7) carrying idle materializes it at
	// revision 7 (modern ordering authority).
	route(`{"type":"ready","capabilities":["agent:permission-mode"],"runningAgents":["agt_1"],"runningAgentTurnStates":{"agt_1":"idle"},"runningAgentTurnStatesWatermark":7}`)
	require.Equal(t, "idle", client.agentTurnStateRecords["agt_1"].State)
	require.Equal(t, int64(7), client.agentTurnStateRecords["agt_1"].Revision)
	require.Equal(t, int64(7), client.turnStateHighWater)

	// A legacy ready frame WITHOUT the watermark must not claim modern ordering
	// authority: it clear-and-replaces the state but leaves the per-agent
	// revision and high-water at 7.
	route(`{"type":"ready","capabilities":["agent:permission-mode"],"runningAgents":["agt_1"],"runningAgentTurnStates":{"agt_1":"active"}}`)
	require.Equal(t, "active", client.agentTurnStateRecords["agt_1"].State)
	require.Equal(t, int64(7), client.agentTurnStateRecords["agt_1"].Revision, "a legacy ready without a watermark must not reset the per-agent revision")
	require.Equal(t, int64(7), client.turnStateHighWater, "a legacy ready must not claim modern ordering authority")
}
