package wire

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/internal/api/daemon"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/rewind"
	"github.com/tidbcloud/AgentTeam/internal/workflow/agentruntime"
)

// turnStateHubStore is a minimal daemon.Store double that lets the real Hub
// connection handler/read agentteam run without persistence. It returns neutral
// defaults for every method the ready/telemetry/disconnect paths exercise.
type turnStateHubStore struct{}

func (turnStateHubStore) MarkMachineDisconnected(context.Context, string) error { return nil }
func (turnStateHubStore) PrepareAgentStart(context.Context, string, string, *message.ReceiveMessage, bool, bool, bool) (agent.StartCommand, error) {
	return agent.StartCommand{}, nil
}
func (turnStateHubStore) MarkMachineReady(context.Context, agent.MachineIdentity, agent.ReadyPayload) error {
	return nil
}
func (turnStateHubStore) UpdateAgentStatusFromMachine(context.Context, string, string, string) (bool, error) {
	return true, nil
}
func (turnStateHubStore) UpdateAgentActivityFromMachine(context.Context, string, string, string, string) (bool, error) {
	return true, nil
}
func (turnStateHubStore) UpdateAgentSessionFromMachine(context.Context, string, string, string, bool) (bool, error) {
	return true, nil
}
func (turnStateHubStore) AgentBelongsToMachine(context.Context, string, string) (bool, error) {
	return true, nil
}
func (turnStateHubStore) AgentNamesForMachine(context.Context, string, string) (string, string, bool, error) {
	return "agent-name", "Agent Name", true, nil
}
func (turnStateHubStore) AgentPermissionModeForMachine(context.Context, string, string) (string, bool, error) {
	return agent.AgentPermissionModeFullAccess, true, nil
}
func (turnStateHubStore) MachineIDForAgent(context.Context, string, string) (string, error) {
	return "mac_1", nil
}
func (turnStateHubStore) ClearWorkspaceWriteAgentAuthTokensForMachine(context.Context, string) error {
	return nil
}
func (turnStateHubStore) RecordMessageRecoveryPointAndPruneOtherThreads(context.Context, string, string, string, string, string, string) ([]rewind.MessageRecoveryPoint, error) {
	return nil, nil
}
func (turnStateHubStore) ListAgentsForMachine(context.Context, string) ([]agent.MachineAgent, error) {
	return nil, nil
}
func (turnStateHubStore) GetTaskActiveTarget(context.Context, string, string, string) (message.ReceiveActiveTaskTarget, bool, error) {
	return message.ReceiveActiveTaskTarget{}, false, nil
}

// turnStateResolver resolves every agent to the harness machine, mirroring the
// production ServiceAdapter the RecoveryRunner wiring uses.
type turnStateResolver struct{}

func (turnStateResolver) MachineIDForAgent(context.Context, string, string) (string, error) {
	return "mac_1", nil
}

// startTurnStateHubServer mounts the real daemon Hub connection handler over a
// real httptest WebSocket endpoint at the exact daemon connect path.
func startTurnStateHubServer(t *testing.T, hub *daemon.Hub) *httptest.Server {
	t.Helper()
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v1/workspaces/wsp_1/daemon/connect" {
			http.NotFound(w, r)
			return
		}
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		identity := agent.MachineIdentity{MachineID: "mac_1", WorkspaceID: "wsp_1", CreatedByUserID: "usr_1"}
		// ServeConnection is the real daemon Hub read agentteam: it registers the
		// client, decodes every ready/agent:session/agent:turn_state frame, and
		// unregisters on close.
		_ = hub.ServeConnection(r.Context(), identity, conn)
	}))
	t.Cleanup(server.Close)
	return server
}

func newTurnStateProtocol(serverURL string, disp *agentteamTestDispatcher) *Protocol {
	return &Protocol{
		id:          "agentteam-test",
		enabled:     true,
		serverURL:   serverURL,
		apiKey:      "machine-key",
		workspaceID: "wsp_1",
		machineID:   "mac_1",
		sendCh:      make(chan []byte, sendBufSize),
		agents:      make(map[string]*agentState),
		dedupRing:   newDedupRing(10),
		disp:        disp,
	}
}

func turnStateReader(hub *daemon.Hub, grace time.Duration) *agentruntime.HubWorkingReader {
	return agentruntime.NewHubWorkingReaderWithOptions(
		turnStateResolver{},
		hub,
		agentruntime.HubWorkingReaderOptions{Now: time.Now, Grace: grace},
	)
}

// TestTurnStateTransportRestartEmptySnapshotClearsStaleActive is the real
// production transport restart harness. A real Desktop wire.Protocol connects
// over a real httptest WebSocket to the real daemon Hub ServeConnection read
// agentteam. It observes real register/ready decoding/HubWorkingReader active
// behavior, then an explicit current empty snapshot (re-ready) retires the stale
// active state and the responder becomes recoverable after the bounded grace.
func TestTurnStateTransportRestartEmptySnapshotClearsStaleActive(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	hub := daemon.New(slog.Default(), turnStateHubStore{}, "", nil)
	server := startTurnStateHubServer(t, hub)
	reader := turnStateReader(hub, 2*time.Second)

	// Connection 1: the dispatcher reports a live active turn.
	disp := &agentteamTestDispatcher{sessions: []types.SessionSnapshot{{
		ChannelID:        "agentteam-test",
		AgentID:          "agt_1",
		Status:           "active",
		RuntimeSessionID: "sess_1",
		TurnState:        types.AgentTurnStateActive,
	}}}
	proto := newTurnStateProtocol(server.URL, disp)
	require.NoError(t, proto.Start(ctx))
	t.Cleanup(func() { _ = proto.Stop(context.Background()) })
	require.NoError(t, proto.WaitConnected(ctx))

	// Real ready + telemetry propagate to the Hub: the reader sees working.
	require.Eventually(t, func() bool {
		working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
		return known && working
	}, 10*time.Second, 20*time.Millisecond, "the real ready/turn-state path must make an active responder working")

	// Mutate the dispatcher to an empty snapshot and re-send ready over the real
	// transport: the current-epoch empty watermark must retire the stale active.
	disp.sessions = nil
	proto.ResendMachineCapabilities()

	// The stale active is cleared: the responder is no longer working.
	require.Eventually(t, func() bool {
		working, _ := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
		return !working
	}, 10*time.Second, 20*time.Millisecond, "an explicit current empty snapshot must retire the stale active state")

	// After the bounded grace the responder becomes recoverable (known=true,
	// not working). The real reader uses the injected short grace for speed.
	require.Eventually(t, func() bool {
		working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
		return known && !working
	}, 10*time.Second, 20*time.Millisecond, "after the bounded grace the responder becomes recoverable")
}

// TestTurnStateTransportRestartOfflineThenReconnectEmptySnapshot proves the
// cross-connection restart: stopping the real connection drives the real Hub
// unregister (offline defer), and a fresh Protocol connection whose ready
// snapshot omits turn-state entries restores the responder through the bounded
// grace to recoverable.
func TestTurnStateTransportRestartOfflineThenReconnectEmptySnapshot(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	hub := daemon.New(slog.Default(), turnStateHubStore{}, "", nil)
	server := startTurnStateHubServer(t, hub)
	reader := turnStateReader(hub, 2*time.Second)

	// Connection 1: active turn -> working.
	disp1 := &agentteamTestDispatcher{sessions: []types.SessionSnapshot{{
		ChannelID:        "agentteam-test",
		AgentID:          "agt_1",
		Status:           "active",
		RuntimeSessionID: "sess_1",
		TurnState:        types.AgentTurnStateActive,
	}}}
	proto1 := newTurnStateProtocol(server.URL, disp1)
	require.NoError(t, proto1.Start(ctx))
	require.NoError(t, proto1.WaitConnected(ctx))
	require.Eventually(t, func() bool {
		working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
		return known && working
	}, 10*time.Second, 20*time.Millisecond, "connection 1 active must be working")

	// Stop the real connection: real unregister -> machine offline -> defer.
	require.NoError(t, proto1.Stop(context.Background()))
	require.Eventually(t, func() bool {
		_, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
		return !known
	}, 10*time.Second, 20*time.Millisecond, "stopping the real connection must leave the responder offline/deferred")

	// Reconnect a FRESH Protocol whose ready snapshot omits turn-state entries
	// (empty dispatcher, no sessions): the current-epoch empty watermark is
	// authoritative and the responder follows grace -> recoverable.
	disp2 := &agentteamTestDispatcher{} // no sessions, zero running agents
	proto2 := newTurnStateProtocol(server.URL, disp2)
	require.NoError(t, proto2.Start(ctx))
	t.Cleanup(func() { _ = proto2.Stop(context.Background()) })
	require.NoError(t, proto2.WaitConnected(ctx))

	require.Eventually(t, func() bool {
		working, known := reader.RuntimeWorking(ctx, "wsp_1", "agt_1")
		return known && !working
	}, 10*time.Second, 20*time.Millisecond, "after reconnect with an empty snapshot the responder must become recoverable")
}

// TestEmitTurnStateEventWireFrameCarriesSessionID proves the Desktop wire
// encoder preserves the SessionID stamped by the AgentTeamReporterSink on a terminal
// idle TurnStateEvent. This is the exact frame the current-generation exit
// publishes (handleProcessExit -> reportExitStatus -> sink -> wire), so the
// server Hub can apply its session fence and drop a stale-epoch idle from a
// replaced session while accepting the current generation's idle.
func TestEmitTurnStateEventWireFrameCarriesSessionID(t *testing.T) {
	proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
	proto.Emit(types.TurnStateEvent{
		AgentID:   "agt_1",
		TurnState: types.AgentTurnStateIdle,
		SessionID: "sess_1",
	})
	frames := drainSendCh(t, proto.sendCh)
	require.Len(t, frames, 1)
	var decoded map[string]any
	require.NoError(t, json.Unmarshal(frames[0], &decoded))
	require.Equal(t, "agent:turn_state", decoded["type"])
	require.Equal(t, "agt_1", decoded["agentId"])
	require.Equal(t, types.AgentTurnStateIdle, decoded["turnState"])
	require.Equal(t, "sess_1", decoded["sessionId"])
}

// TestEmitSessionEventWireFrameCarriesExpectedSessionID proves the wire encoder
// carries the generation-fenced terminal clear: a SessionEvent with an empty
// SessionID and a non-empty ExpectedSessionID serializes expectedSessionId so
// the daemon can compare-and-clear against the dying generation. Legacy clears
// and non-empty publications omit the fence entirely.
func TestEmitSessionEventWireFrameCarriesExpectedSessionID(t *testing.T) {
	t.Run("terminal clear carries expected session", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.SessionEvent{AgentID: "agt_1", SessionID: "", ExpectedSessionID: "sess_1"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "agent:session", decoded["type"])
		require.Equal(t, "agt_1", decoded["agentId"])
		require.Equal(t, "", decoded["sessionId"])
		require.Equal(t, "sess_1", decoded["expectedSessionId"])
	})

	t.Run("legacy clear omits expected session", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.SessionEvent{AgentID: "agt_1", SessionID: ""})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "agent:session", decoded["type"])
		require.Equal(t, "", decoded["sessionId"])
		_, hasFence := decoded["expectedSessionId"]
		require.False(t, hasFence, "a legacy clear must not carry the optional fence")
	})

	t.Run("non-empty publication unchanged", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.SessionEvent{AgentID: "agt_1", SessionID: "sess_2"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "sess_2", decoded["sessionId"])
		_, hasFence := decoded["expectedSessionId"]
		require.False(t, hasFence, "a non-empty session publication must not carry the fence")
	})
}

// TestEmitStatusEventWireFrameCarriesExpectedSessionID proves the wire encoder
// carries the generation-fenced terminal status: a StatusEvent with Status
// "inactive" and a non-empty ExpectedSessionID serializes expectedSessionId so
// the daemon drops a delayed stop from a replaced generation. Legacy and
// non-terminal frames omit the fence entirely.
func TestEmitStatusEventWireFrameCarriesExpectedSessionID(t *testing.T) {
	t.Run("terminal inactive carries expected session", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.StatusEvent{AgentID: "agt_1", Status: "inactive", ExpectedSessionID: "sess_1"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "agent:status", decoded["type"])
		require.Equal(t, "agt_1", decoded["agentId"])
		require.Equal(t, "inactive", decoded["status"])
		require.Equal(t, "sess_1", decoded["expectedSessionId"])
	})

	t.Run("legacy inactive omits expected session", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.StatusEvent{AgentID: "agt_1", Status: "inactive"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "inactive", decoded["status"])
		_, hasFence := decoded["expectedSessionId"]
		require.False(t, hasFence, "a legacy status must not carry the optional fence")
	})

	t.Run("non-terminal status unchanged", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.StatusEvent{AgentID: "agt_1", Status: "active"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "active", decoded["status"])
		_, hasFence := decoded["expectedSessionId"]
		require.False(t, hasFence, "a non-terminal status must not carry the fence")
	})
}

// TestEmitActivityEventWireFrameCarriesExpectedSessionID proves the wire encoder
// carries the generation-fenced terminal activity: an ActivityEvent with
// Activity "offline" and a non-empty ExpectedSessionID serializes
// expectedSessionId so the daemon drops a delayed stop from a replaced
// generation. Legacy and non-terminal frames omit the fence entirely.
func TestEmitActivityEventWireFrameCarriesExpectedSessionID(t *testing.T) {
	t.Run("terminal offline carries expected session", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.ActivityEvent{AgentID: "agt_1", Activity: "offline", Detail: "process exited normally", ExpectedSessionID: "sess_1"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "agent:activity", decoded["type"])
		require.Equal(t, "agt_1", decoded["agentId"])
		require.Equal(t, "offline", decoded["activity"])
		require.Equal(t, "sess_1", decoded["expectedSessionId"])
	})

	t.Run("legacy offline omits expected session", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.ActivityEvent{AgentID: "agt_1", Activity: "offline", Detail: "process exited normally"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "offline", decoded["activity"])
		_, hasFence := decoded["expectedSessionId"]
		require.False(t, hasFence, "a legacy activity must not carry the optional fence")
	})

	t.Run("non-terminal activity unchanged", func(t *testing.T) {
		proto := &Protocol{id: "agentteam-test", sendCh: make(chan []byte, 8)}
		proto.Emit(types.ActivityEvent{AgentID: "agt_1", Activity: "thinking", Detail: "Working on it"})
		frames := drainSendCh(t, proto.sendCh)
		require.Len(t, frames, 1)
		var decoded map[string]any
		require.NoError(t, json.Unmarshal(frames[0], &decoded))
		require.Equal(t, "thinking", decoded["activity"])
		_, hasFence := decoded["expectedSessionId"]
		require.False(t, hasFence, "a non-terminal activity must not carry the fence")
	})
}