package daemon

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

// L5: full-chain integration from the daemon-protocol `agent:turn_state`
// command through the server Hub overlay to the HubWorkingReader that the
// RecoveryRunner consumes. This closes the missing link between
// emitAgentTeamSessionTelemetrySnapshot (covered by the wire L1 test) and the
// recovery working judgment. A matching sessionId is accepted; a stale sessionId
// is dropped and the current known state is preserved (never cleared to unknown).
func TestTurnStateFullChainHandlerToRecoveryReader(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{clients: map[string]*clientState{}}
	client := &clientState{
		machineID:             "mac_1",
		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{},
	}
	hub.clients["mac_1"] = client

	// The daemon reported the current session (agent:session) then the turn_state
	// carrying the same RuntimeSessionID.
	client.setAgentSession("agt_1", "sess_1")

	reader := agentruntime.NewHubWorkingReader(fakeAgentMachineResolver{machineID: "mac_1"}, hub)

	// Route every frame through hub.handleMessage so the real wire router
	// (registry -> handler) is exercised, not the handler directly.
	route := func(frame string) {
		require.NoError(t, hub.handleMessage(ctx, client, []byte(frame)))
	}

	// idle live session -> recoverable.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1"}`)
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working, "idle live session must be recoverable")

	// active turn -> working, never recovered.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"active","sessionId":"sess_1"}`)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "an active turn must never be recovered")

	// real queued turn -> working.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"queued","sessionId":"sess_1"}`)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working, "a real queued turn must be working")

	// queued -> active -> idle transitions flow through the reader.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"active","sessionId":"sess_1"}`)
	working, _ = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, working)
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1"}`)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working, "turn must return to idle and be recoverable")
}

// TestTurnStateFullChainStaleSessionFence proves a stale turn_state from a
// replaced session is dropped and the current known state is preserved (never
// cleared to unknown); only disconnect/session-clear/ready-replacement leave
// unknown.
func TestTurnStateFullChainStaleSessionFence(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{clients: map[string]*clientState{}}
	client := &clientState{
		machineID:             "mac_1",
		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{},
	}
	hub.clients["mac_1"] = client
	reader := agentruntime.NewHubWorkingReader(fakeAgentMachineResolver{machineID: "mac_1"}, hub)

	// New session replaces the old one.
	client.setAgentSession("agt_1", "sess_new")
	route := func(frame string) {
		require.NoError(t, hub.handleMessage(ctx, client, []byte(frame)))
	}
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"queued","sessionId":"sess_new"}`)
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working)

	// A stale turn_state from the old session arrives late: dropped, current
	// known state preserved.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_old"}`)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known, "a stale turn_state must not clear the current known state to unknown")
	require.True(t, working, "the current session's queued state must be preserved after a stale event is dropped")

	// The new session's matching turn_state still wins.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_new"}`)
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working)

	// Disconnect (client removed) -> unknown (fail-safe).
	delete(hub.clients, "mac_1")
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.False(t, known, "disconnect must leave unknown")
	require.False(t, working)
}

type fakeAgentMachineResolver struct {
	machineID string
}

func (resolver fakeAgentMachineResolver) MachineIDForAgent(_ context.Context, _ string, _ string) (string, error) {
	return resolver.machineID, nil
}

// TestTurnStateFullChainEpochFenceDropsStaleEpoch proves a turn_state carrying a
// connection epoch that does not match the current connection's session
// generation is dropped, so an old-connection event can never overwrite the
// current-epoch snapshot.
func TestTurnStateFullChainEpochFenceDropsStaleEpoch(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{clients: map[string]*clientState{}}
	client := &clientState{
		machineID:                 "mac_1",
		sessionGeneration:         "epoch_current",
		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{},
		turnStateSnapshotReceived: true,
	}
	hub.clients["mac_1"] = client
	client.setAgentSession("agt_1", "sess_1")
	client.setAgentTurnState("agt_1", "queued", "sess_1", 0, "")

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

	// A stale-epoch event (from an old connection) is dropped: the current known
	// state is preserved.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1","epoch":"epoch_old"}`)
	state, ok := client.agentTurnState("agt_1")
	require.True(t, ok)
	require.Equal(t, "queued", state, "a stale-epoch turn_state must not overwrite the current snapshot")

	// A current-epoch event is accepted.
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"idle","sessionId":"sess_1","epoch":"epoch_current"}`)
	state, ok = client.agentTurnState("agt_1")
	require.True(t, ok)
	require.Equal(t, "idle", state)

	// A turn_state without an epoch is accepted (legacy daemon compatibility).
	route(`{"type":"agent:turn_state","agentId":"agt_1","turnState":"active","sessionId":"sess_1"}`)
	state, ok = client.agentTurnState("agt_1")
	require.True(t, ok)
	require.Equal(t, "active", state)
}

// TestTurnStateFullChainReconnectMissingSnapshotBoundedGrace proves the recovery
// chain across a real daemon reconnect: a fresh connection with no snapshot
// defers within the bounded grace, and once the ready snapshot reports an idle
// state the responder becomes recoverable.
func TestTurnStateFullChainReconnectMissingSnapshotBoundedGrace(t *testing.T) {
	ctx := context.Background()
	now := time.Date(2026, 8, 16, 12, 0, 0, 0, time.UTC)
	hub := &Hub{clients: map[string]*clientState{}}
	client := &clientState{
		machineID:             "mac_1",
		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:           now.Add(-30 * time.Second),
	}
	hub.clients["mac_1"] = client

	reader := agentruntime.NewHubWorkingReaderWithOptions(
		fakeAgentMachineResolver{machineID: "mac_1"},
		hub,
		agentruntime.HubWorkingReaderOptions{Now: func() time.Time { return now }, Grace: 60 * time.Second},
	)

	// Just reconnected, no snapshot yet: within grace, deferred (no send, no timing).
	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.False(t, known, "reconnected machine without a snapshot defers within grace")
	require.False(t, working)

	// The ready snapshot arrives and reports an idle turn: recoverable.
	client.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "idle"}, wm(0))
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.False(t, working, "idle snapshot after reconnect must be recoverable")

	// The same snapshot reports an active turn: working, never recovered.
	client.replaceReadyAgentTurnStates([]string{"agt_1"}, map[string]string{"agt_1": "active"}, wm(0))
	working, known = reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known)
	require.True(t, working)
}

// TestTurnStateFullChainReconnectGraceExpiryRecoverable proves that even if the
// daemon never reports a snapshot after reconnect, the bounded grace elapses and
// the responder becomes recoverable instead of being stuck unknown forever.
func TestTurnStateFullChainReconnectGraceExpiryRecoverable(t *testing.T) {
	ctx := context.Background()
	now := time.Date(2026, 8, 16, 12, 0, 0, 0, time.UTC)
	hub := &Hub{clients: map[string]*clientState{}}
	client := &clientState{
		machineID:             "mac_1",
		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:           now.Add(-2 * time.Minute), // past the 60s grace
	}
	hub.clients["mac_1"] = client

	reader := agentruntime.NewHubWorkingReaderWithOptions(
		fakeAgentMachineResolver{machineID: "mac_1"},
		hub,
		agentruntime.HubWorkingReaderOptions{Now: func() time.Time { return now }, Grace: 60 * time.Second},
	)

	working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
	require.True(t, known, "an online responder with no snapshot must become recoverable after the bounded grace")
	require.False(t, working)
}