package daemon

import (
	"context"
	"encoding/json"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
)

// sessionFenceStore records the UpdateAgentSessionFromMachine mutations so the
// test can prove a stale terminal clear is dropped BEFORE any store mutation.
type sessionFenceStore struct {
	recordingStartStore
	mu      sync.Mutex
	updates []string // agentID:sessionID per accepted mutation
}

func (store *sessionFenceStore) UpdateAgentSessionFromMachine(_ context.Context, _ string, agentID string, sessionID string, persistLegacySession bool) (bool, error) {
	store.mu.Lock()
	defer store.mu.Unlock()
	store.updates = append(store.updates, agentID+":"+sessionID)
	return true, nil
}

func (store *sessionFenceStore) sessionUpdates() []string {
	store.mu.Lock()
	defer store.mu.Unlock()
	return append([]string(nil), store.updates...)
}

// routeSession sends one agent:session frame through the real daemon protocol
// router into the given client.
func routeSession(t *testing.T, hub *Hub, client *clientState, sessionID, expectedSessionID string) {
	t.Helper()
	payload := map[string]any{
		"type":      "agent:session",
		"agentId":   "agt_1",
		"sessionId": sessionID,
	}
	if expectedSessionID != "" {
		payload["expectedSessionId"] = expectedSessionID
	}
	raw, err := json.Marshal(payload)
	require.NoError(t, err)
	require.NoError(t, hub.handleMessage(context.Background(), client, raw))
}

func newSessionFenceClient(currentSession string) *clientState {
	client := &clientState{
		machineID:             "mac_1",
		workspaceID:           "wsp_1",
		createdByUserID:       "usr_1",
		agentSessions:         map[string]string{},
		runningAgents:         map[string]bool{},
		startingAgents:        map[string]bool{},
		stoppedAgents:         map[string]bool{},
		agentStatuses:         map[string]string{},
		agentActivities:       map[string]agentActivityState{},
		agentTurnStateRecords: map[string]agentTurnStateRecord{},
	}
	if currentSession != "" {
		client.agentSessions["agt_1"] = currentSession
		client.runningAgents["agt_1"] = true
	}
	return client
}

// TestSessionFenceMatchingTerminalClearApplies proves a terminal clear carrying
// the CURRENT session identity still clears the current session (store mutation,
// running set cleared, session removed) — current-generation cleanup stays
// intact.
func TestSessionFenceMatchingTerminalClearApplies(t *testing.T) {
	store := &sessionFenceStore{}
	hub := &Hub{store: store}
	client := newSessionFenceClient("session-2")
	routeSession(t, hub, client, "", "session-2")
	require.Empty(t, client.agentSession("agt_1"), "a matching terminal clear must remove the current session")
	require.False(t, client.isAgentRunning("agt_1"), "a matching terminal clear must mark the agent not running")
	require.Equal(t, []string{"agt_1:"}, store.sessionUpdates(), "a matching terminal clear must reach the store exactly once")
}

// TestSessionFenceStaleTerminalClearDropped proves a delayed stop from a
// replaced generation — a clear carrying an expected session that is NOT the
// current session — is dropped ENTIRELY: no store mutation, no realtime/current
// session mutation, the replacement session and running state preserved.
func TestSessionFenceStaleTerminalClearDropped(t *testing.T) {
	store := &sessionFenceStore{}
	hub := &Hub{store: store}
	client := newSessionFenceClient("session-2")
	routeSession(t, hub, client, "", "session-1")
	require.Equal(t, "session-2", client.agentSession("agt_1"),
		"a stale terminal clear must never delete the replacement session")
	require.True(t, client.isAgentRunning("agt_1"), "a stale terminal clear must not clear the running state")
	require.Empty(t, store.sessionUpdates(), "a stale terminal clear must be dropped before any store mutation")
}

// TestSessionFenceLegacyClearUnconditional proves a legacy empty clear WITHOUT
// the optional expected-session fence keeps the pre-existing unconditional
// behavior: it clears whatever session is current.
func TestSessionFenceLegacyClearUnconditional(t *testing.T) {
	store := &sessionFenceStore{}
	hub := &Hub{store: store}
	client := newSessionFenceClient("session-2")
	routeSession(t, hub, client, "", "")
	require.Empty(t, client.agentSession("agt_1"), "a legacy unconditional clear must still clear the current session")
	require.False(t, client.isAgentRunning("agt_1"))
	require.Equal(t, []string{"agt_1:"}, store.sessionUpdates(), "a legacy clear must reach the store once")
}

// TestSessionFencePublicationUnchanged proves non-empty session publications are
// unaffected by the fence: a new session still replaces the old one.
func TestSessionFencePublicationUnchanged(t *testing.T) {
	store := &sessionFenceStore{}
	hub := &Hub{store: store}
	client := newSessionFenceClient("session-1")
	routeSession(t, hub, client, "session-2", "")
	require.Equal(t, "session-2", client.agentSession("agt_1"))
	require.True(t, client.isAgentRunning("agt_1"))
	require.Equal(t, []string{"agt_1:session-2"}, store.sessionUpdates())
}