package daemon

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

	"github.com/tidbcloud/AgentTeam/internal/api/im"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
	"github.com/tidbcloud/AgentTeam/internal/store/rewind"
)

type recordingStartStore struct {
	includeLegacySessions            []bool
	persistLegacySessions            []bool
	markMachineReadyErr              error
	agentBelongsToMachine            *bool
	startVisibleReplyContractEnabled *bool
}

func (store *recordingStartStore) MarkMachineDisconnected(context.Context, string) error {
	return nil
}

func (store *recordingStartStore) PrepareAgentStart(_ context.Context, _ string, _ string, _ *message.ReceiveMessage, _ bool, includeLegacySession bool, _ bool) (agent.StartCommand, error) {
	store.includeLegacySessions = append(store.includeLegacySessions, includeLegacySession)
	visibleReplyContractEnabled := true
	if store.startVisibleReplyContractEnabled != nil {
		visibleReplyContractEnabled = *store.startVisibleReplyContractEnabled
	}
	command := agent.StartCommand{
		Config: agent.StartConfig{
			Runtime:                     "codex",
			PermissionMode:              agent.AgentPermissionModeFullAccess,
			RuntimePolicy:               agent.NewRuntimePolicySnapshot("codex", agent.AgentPermissionModeFullAccess, 2, visibleReplyContractEnabled),
			VisibleReplyContractEnabled: visibleReplyContractEnabled,
		},
	}
	if includeLegacySession {
		command.Config.SessionID = "db-session"
	}
	return command, nil
}

func (store *recordingStartStore) MarkMachineReady(context.Context, agent.MachineIdentity, agent.ReadyPayload) error {
	if store.markMachineReadyErr != nil {
		return store.markMachineReadyErr
	}
	return nil
}

func (store *recordingStartStore) UpdateAgentStatusFromMachine(context.Context, string, string, string) (bool, error) {
	return true, nil
}

func (store *recordingStartStore) UpdateAgentActivityFromMachine(context.Context, string, string, string, string) (bool, error) {
	return true, nil
}

func (store *recordingStartStore) UpdateAgentSessionFromMachine(_ context.Context, _ string, _ string, _ string, persistLegacySession bool) (bool, error) {
	store.persistLegacySessions = append(store.persistLegacySessions, persistLegacySession)
	return true, nil
}

func (store *recordingStartStore) AgentBelongsToMachine(context.Context, string, string) (bool, error) {
	if store.agentBelongsToMachine != nil {
		return *store.agentBelongsToMachine, nil
	}
	return true, nil
}

func (store *recordingStartStore) AgentNamesForMachine(_ context.Context, _ string, _ string) (string, string, bool, error) {
	if store.agentBelongsToMachine != nil && !*store.agentBelongsToMachine {
		return "", "", false, nil
	}
	return "agent-name", "Agent Name", true, nil
}

func (store *recordingStartStore) AgentPermissionModeForMachine(context.Context, string, string) (string, bool, error) {
	return agent.AgentPermissionModeFullAccess, true, nil
}

func (store *recordingStartStore) ClearWorkspaceWriteAgentAuthTokensForMachine(context.Context, string) error {
	return nil
}

func (store *recordingStartStore) RecordMessageRecoveryPointAndPruneOtherThreads(context.Context, string, string, string, string, string, string) ([]rewind.MessageRecoveryPoint, error) {
	return nil, nil
}

func (store *recordingStartStore) ListAgentsForMachine(context.Context, string) ([]agent.MachineAgent, error) {
	return nil, nil
}

func (store *recordingStartStore) GetTaskActiveTarget(context.Context, string, string, string) (message.ReceiveActiveTaskTarget, bool, error) {
	return message.ReceiveActiveTaskTarget{}, false, nil
}

func (store *recordingStartStore) MachineIDForAgent(_ context.Context, _ string, _ string) (string, error) {
	return "", nil
}

func TestPrepareAgentStartForClientGatesLegacySessionByCapability(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	hub := &Hub{store: store}

	legacyClient := &clientState{capabilities: map[string]struct{}{}}
	command, err := hub.prepareAgentStartForClient(context.Background(), legacyClient, "agt_legacy", nil)
	require.NoError(t, err)
	require.Equal(t, "db-session", command.Config.SessionID)

	capableClient := &clientState{capabilities: map[string]struct{}{agentLocalSessionResumeCapability: {}}}
	command, err = hub.prepareAgentStartForClient(context.Background(), capableClient, "agt_capable", nil)
	require.NoError(t, err)
	require.Empty(t, command.Config.SessionID)

	runtimeAuthorityClient := &clientState{capabilities: map[string]struct{}{agentRuntimeSessionAuthorityCapability: {}}}
	command, err = hub.prepareAgentStartForClient(context.Background(), runtimeAuthorityClient, "agt_runtime_authority", nil)
	require.NoError(t, err)
	require.Equal(t, "db-session", command.Config.SessionID)

	require.Equal(t, []bool{true, false, true}, store.includeLegacySessions)
}

func TestPrepareAgentStartForClientUsesConnectionSession(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	hub := &Hub{store: store}
	client := &clientState{capabilities: map[string]struct{}{agentLocalSessionResumeCapability: {}}}
	client.setAgentSession("agt_live", "live-session")

	command, err := hub.prepareAgentStartForClient(context.Background(), client, "agt_live", nil)
	require.NoError(t, err)
	require.Equal(t, "live-session", command.Config.SessionID)
	require.Equal(t, []bool{false}, store.includeLegacySessions)
}

func TestPrepareAgentStartForClientDoesNotUseRuntimeAuthorityAsProviderSession(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	hub := &Hub{store: store}
	client := &clientState{capabilities: map[string]struct{}{
		agentLocalSessionResumeCapability:      {},
		agentRuntimeSessionAuthorityCapability: {},
	}}
	client.setAgentSession("agt_live", "ars_runtime")

	command, err := hub.prepareAgentStartForClient(context.Background(), client, "agt_live", nil)
	require.NoError(t, err)
	require.Empty(t, command.Config.SessionID)
	require.Equal(t, []bool{false}, store.includeLegacySessions)
}

func TestPrepareAgentStartForClientGatesVisibleReplyContractByCapability(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	hub := &Hub{store: store}

	legacyClient := &clientState{capabilities: map[string]struct{}{}}
	command, err := hub.prepareAgentStartForClient(context.Background(), legacyClient, "agt_legacy", nil)
	require.NoError(t, err)
	require.False(t, command.Config.VisibleReplyContractEnabled)
	require.False(t, command.Config.RuntimePolicy.VisibleReplyContractEnabled)

	capableClient := &clientState{capabilities: map[string]struct{}{agentVisibleReplyContractCapability: {}}}
	command, err = hub.prepareAgentStartForClient(context.Background(), capableClient, "agt_capable", nil)
	require.NoError(t, err)
	require.True(t, command.Config.VisibleReplyContractEnabled)
	require.True(t, command.Config.RuntimePolicy.VisibleReplyContractEnabled)
}

func TestPrepareAgentStartForClientAlwaysOnDespiteFalseStoreValue(t *testing.T) {
	t.Parallel()

	// A historical DB row may still serve visibleReplyContractEnabled=false.
	// Capability-declaring clients must still receive enabled=true (always-on).
	falseValue := false
	store := &recordingStartStore{startVisibleReplyContractEnabled: &falseValue}
	hub := &Hub{store: store}

	capableClient := &clientState{capabilities: map[string]struct{}{agentVisibleReplyContractCapability: {}}}
	command, err := hub.prepareAgentStartForClient(context.Background(), capableClient, "agt_capable", nil)
	require.NoError(t, err)
	require.True(t, command.Config.VisibleReplyContractEnabled)
	require.True(t, command.Config.RuntimePolicy.VisibleReplyContractEnabled)
}

func TestAgentSessionTelemetryPersistsLegacyMirrorByCapability(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	hub := &Hub{store: store}
	payload, err := json.Marshal(map[string]any{
		"type":      "agent:session",
		"agentId":   "agt_live",
		"sessionId": "sess-live",
	})
	require.NoError(t, err)

	legacyClient := &clientState{
		machineID:     "mac_legacy",
		workspaceID:   "wsp_legacy",
		capabilities:  map[string]struct{}{},
		agentSessions: map[string]string{},
	}
	require.NoError(t, hub.handleMessage(context.Background(), legacyClient, payload))
	require.Equal(t, "sess-live", legacyClient.agentSession("agt_live"))

	capableClient := &clientState{
		machineID:     "mac_capable",
		workspaceID:   "wsp_capable",
		capabilities:  map[string]struct{}{agentLocalSessionResumeCapability: {}},
		agentSessions: map[string]string{},
	}
	require.NoError(t, hub.handleMessage(context.Background(), capableClient, payload))
	require.Equal(t, "sess-live", capableClient.agentSession("agt_live"))

	runtimeAuthorityClient := &clientState{
		machineID:     "mac_runtime_authority",
		workspaceID:   "wsp_runtime_authority",
		capabilities:  map[string]struct{}{agentRuntimeSessionAuthorityCapability: {}},
		agentSessions: map[string]string{},
	}
	runtimeAuthorityPayload, err := json.Marshal(map[string]any{
		"type":      "agent:session",
		"agentId":   "agt_live",
		"sessionId": "ars_live",
	})
	require.NoError(t, err)
	require.NoError(t, hub.handleMessage(context.Background(), runtimeAuthorityClient, runtimeAuthorityPayload))
	require.Equal(t, "ars_live", runtimeAuthorityClient.agentSession("agt_live"))

	require.Equal(t, []bool{true, false, false}, store.persistLegacySessions)
}

func TestReadySeedsRunningAgentRuntimeSessionsForReconnect(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	client := &clientState{
		machineID:     "mac_reconnect",
		workspaceID:   "wsp_reconnect",
		runningAgents: map[string]bool{},
		agentSessions: map[string]string{},
	}
	hub := &Hub{
		store: store,
		clients: map[string]*clientState{
			client.machineID: client,
		},
	}
	payload, err := json.Marshal(map[string]any{
		"type":         "ready",
		"capabilities": []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
		"runtimes":     []string{"claude"},
		"runningAgents": []string{
			"agt_live",
		},
		"runningAgentSessions": map[string]string{
			"agt_live":  " runtime-session-1 ",
			"agt_other": "runtime-session-ignored",
		},
	})
	if err != nil {
		t.Fatalf("marshal ready payload: %v", err)
	}

	if err := hub.handleMessage(context.Background(), client, payload); err != nil {
		t.Fatalf("handle ready payload: %v", err)
	}

	if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_live"); !ok || got != "runtime-session-1" {
		t.Fatalf("AgentRuntimeSessionID(live) = %q, %v; want runtime-session-1, true", got, ok)
	}
	if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_other"); ok || got != "" {
		t.Fatalf("AgentRuntimeSessionID(other) = %q, %v; want empty, false", got, ok)
	}
}

func TestRegisterPreservesRuntimeSessionUntilReadyReplay(t *testing.T) {
	t.Parallel()

	serverConn := newProtocolTestWebSocketPair(t)
	oldClient := &clientState{
		machineID:     "mac_reconnect_transition",
		workspaceID:   "wsp_reconnect_transition",
		conn:          serverConn,
		agentSessions: map[string]string{"agt_live": "runtime-session-1"},
	}
	store := &recordingStartStore{}
	hub := &Hub{
		store: store,
		clients: map[string]*clientState{
			oldClient.machineID: oldClient,
		},
	}

	newClient := &clientState{
		machineID:     oldClient.machineID,
		workspaceID:   oldClient.workspaceID,
		agentSessions: map[string]string{},
	}
	hub.register(newClient)

	if got, ok := hub.AgentRuntimeSessionID(newClient.machineID, "agt_live"); !ok || got != "runtime-session-1" {
		t.Fatalf("AgentRuntimeSessionID(live) during reconnect = %q, %v; want runtime-session-1, true", got, ok)
	}

	readyPayload, err := json.Marshal(map[string]any{
		"type":          "ready",
		"capabilities":  []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
		"runtimes":      []string{"claude"},
		"runningAgents": []string{"agt_live"},
	})
	require.NoError(t, err)
	require.NoError(t, hub.handleMessage(context.Background(), newClient, readyPayload))

	if got, ok := hub.AgentRuntimeSessionID(newClient.machineID, "agt_live"); ok || got != "" {
		t.Fatalf("AgentRuntimeSessionID(live) after ready replay = %q, %v; want empty, false", got, ok)
	}
}

func TestReadyDoesNotSeedRuntimeSessionBeforeSessionPublish(t *testing.T) {
	t.Parallel()

	store := &recordingStartStore{}
	client := &clientState{
		machineID:     "mac_resume",
		workspaceID:   "wsp_resume",
		runningAgents: map[string]bool{},
		agentSessions: map[string]string{},
	}
	hub := &Hub{
		store: store,
		clients: map[string]*clientState{
			client.machineID: client,
		},
	}
	payload, err := json.Marshal(map[string]any{
		"type":         "ready",
		"capabilities": []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
		"runtimes":     []string{"claude"},
		"runningAgents": []string{
			"agt_live",
		},
	})
	if err != nil {
		t.Fatalf("marshal ready payload: %v", err)
	}

	if err := hub.handleMessage(context.Background(), client, payload); err != nil {
		t.Fatalf("handle ready payload: %v", err)
	}

	if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_live"); ok || got != "" {
		t.Fatalf("AgentRuntimeSessionID(live) before session publish = %q, %v; want empty, false", got, ok)
	}

	sessionPayload, err := json.Marshal(map[string]any{
		"type":      "agent:session",
		"agentId":   "agt_live",
		"sessionId": "runtime-session-1",
	})
	if err != nil {
		t.Fatalf("marshal session payload: %v", err)
	}
	if err := hub.handleMessage(context.Background(), client, sessionPayload); err != nil {
		t.Fatalf("handle session payload: %v", err)
	}

	if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_live"); !ok || got != "runtime-session-1" {
		t.Fatalf("AgentRuntimeSessionID(live) after session publish = %q, %v; want runtime-session-1, true", got, ok)
	}
}

func TestWorkspaceRuntimeSnapshotReturnsLiveRuntimeEvidence(t *testing.T) {
	t.Parallel()

	hub := &Hub{
		clients: map[string]*clientState{
			"mac_2": {
				machineID:           " mac_2 ",
				workspaceID:         " wsp_live ",
				createdByUserID:     " usr_2 ",
				runtimeCatalogKnown: true,
				runningAgents:       map[string]bool{},
				startingAgents:      map[string]bool{},
				agentSessions:       map[string]string{},
			},
			"mac_1": {
				machineID:           " mac_1 ",
				workspaceID:         " wsp_live ",
				createdByUserID:     " usr_1 ",
				runtimeCatalogKnown: true,
				runningAgents: map[string]bool{
					"agt_live":     true,
					"agt_starting": true,
				},
				startingAgents: map[string]bool{
					"agt_starting": true,
				},
				agentSessions: map[string]string{
					"agt_live":         " sess-live ",
					"agt_session_only": " sess-only ",
					"agt_blank":        " ",
				},
			},
			"mac_unknown": {
				machineID:           "mac_unknown",
				workspaceID:         "wsp_live",
				runtimeCatalogKnown: false,
				runningAgents: map[string]bool{
					"agt_hidden": true,
				},
				agentSessions: map[string]string{
					"agt_hidden": "sess-hidden",
				},
			},
			"mac_other_workspace": {
				machineID:           "mac_other_workspace",
				workspaceID:         "wsp_other",
				runtimeCatalogKnown: true,
				runningAgents: map[string]bool{
					"agt_other": true,
				},
				agentSessions: map[string]string{
					"agt_other": "sess-other",
				},
			},
		},
	}

	require.Equal(t, im.WorkspaceRuntimeSnapshot{
		Machines: []im.WorkspaceRuntimeMachineSnapshot{
			{
				MachineID:   "mac_1",
				OwnerUserID: "usr_1",
				Agents: []im.WorkspaceRuntimeAgentSnapshot{
					{
						AgentID:   "agt_live",
						Status:    "active",
						SessionID: "sess-live",
					},
					{
						AgentID:   "agt_session_only",
						Status:    "active",
						SessionID: "sess-only",
					},
				},
			},
			{
				MachineID:   "mac_2",
				OwnerUserID: "usr_2",
				Agents:      []im.WorkspaceRuntimeAgentSnapshot{},
			},
		},
	}, hub.WorkspaceRuntimeSnapshot(" wsp_live "))
}

func newProtocolTestWebSocketPair(t *testing.T) *websocket.Conn {
	t.Helper()

	serverConnCh := make(chan *websocket.Conn, 1)
	upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Errorf("upgrade websocket: %v", err)
			return
		}
		serverConnCh <- conn
	}))
	t.Cleanup(server.Close)

	clientConn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
	require.NoError(t, err)
	t.Cleanup(func() { _ = clientConn.Close() })

	select {
	case serverConn := <-serverConnCh:
		t.Cleanup(func() { _ = serverConn.Close() })
		return serverConn
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for server websocket")
	}
	return nil
}

func TestReadyRejectedByStoreClearsRuntimeSessions(t *testing.T) {
	t.Parallel()

	readyErr := errors.New("ready rejected")
	store := &recordingStartStore{markMachineReadyErr: readyErr}
	client := &clientState{
		machineID:     "mac_rejected",
		workspaceID:   "wsp_rejected",
		runningAgents: map[string]bool{},
		agentSessions: map[string]string{},
	}
	hub := &Hub{
		store: store,
		clients: map[string]*clientState{
			client.machineID: client,
		},
	}
	payload, err := json.Marshal(map[string]any{
		"type":         "ready",
		"capabilities": []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
		"runtimes":     []string{"claude"},
		"runningAgents": []string{
			"agt_live",
		},
		"runningAgentSessions": map[string]string{
			"agt_live": "runtime-session-rejected",
		},
	})
	require.NoError(t, err)

	require.ErrorIs(t, hub.handleMessage(context.Background(), client, payload), readyErr)

	if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_live"); ok || got != "" {
		t.Fatalf("AgentRuntimeSessionID(live) after rejected ready = %q, %v; want empty, false", got, ok)
	}
	if hub.AgentRuntimeRunning(client.machineID, "agt_live") {
		t.Fatalf("AgentRuntimeRunning(live) after rejected ready = true, want false")
	}
}

func TestReadySnapshotsReplaceRuntimeSessions(t *testing.T) {
	t.Parallel()

	for _, tc := range []struct {
		name             string
		updatePayload    map[string]any
		wantStillRunning bool
	}{
		{
			name: "omitted_running_agent_sessions",
			updatePayload: map[string]any{
				"type":          "ready",
				"capabilities":  []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
				"runtimes":      []string{"claude"},
				"runningAgents": []string{"agt_live"},
			},
			wantStillRunning: true,
		},
		{
			name: "empty_running_agent_sessions",
			updatePayload: map[string]any{
				"type":                 "ready",
				"capabilities":         []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
				"runtimes":             []string{"claude"},
				"runningAgents":        []string{"agt_live"},
				"runningAgentSessions": map[string]string{},
			},
			wantStillRunning: true,
		},
		{
			name: "agent_no_longer_running",
			updatePayload: map[string]any{
				"type":          "ready",
				"capabilities":  []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
				"runtimes":      []string{"claude"},
				"runningAgents": []string{},
			},
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			store := &recordingStartStore{}
			client := &clientState{
				machineID:     "mac_" + tc.name,
				workspaceID:   "wsp_" + tc.name,
				runningAgents: map[string]bool{},
				agentSessions: map[string]string{},
			}
			hub := &Hub{
				store: store,
				clients: map[string]*clientState{
					client.machineID: client,
				},
			}
			initialPayload, err := json.Marshal(map[string]any{
				"type":         "ready",
				"capabilities": []string{agentPermissionModeCapability, agentMaaSDirectCapability, agentLocalSessionResumeCapability},
				"runtimes":     []string{"claude"},
				"runningAgents": []string{
					"agt_live",
				},
				"runningAgentSessions": map[string]string{
					"agt_live": "runtime-session-1",
				},
			})
			require.NoError(t, err)
			require.NoError(t, hub.handleMessage(context.Background(), client, initialPayload))
			if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_live"); !ok || got != "runtime-session-1" {
				t.Fatalf("AgentRuntimeSessionID(live) after initial ready = %q, %v; want runtime-session-1, true", got, ok)
			}

			updatePayload, err := json.Marshal(tc.updatePayload)
			require.NoError(t, err)
			require.NoError(t, hub.handleMessage(context.Background(), client, updatePayload))

			if got, ok := hub.AgentRuntimeSessionID(client.machineID, "agt_live"); ok || got != "" {
				t.Fatalf("AgentRuntimeSessionID(live) after replacement ready = %q, %v; want empty, false", got, ok)
			}
			if got := hub.AgentRuntimeRunning(client.machineID, "agt_live"); got != tc.wantStillRunning {
				t.Fatalf("AgentRuntimeRunning(live) after replacement ready = %v, want %v", got, tc.wantStillRunning)
			}
		})
	}
}