package daemon

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api/dto"
	"github.com/tidbcloud/AgentTeam/internal/archive"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
	"github.com/tidbcloud/AgentTeam/internal/service/agentruntime"
	"github.com/tidbcloud/AgentTeam/internal/store"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
	"github.com/tidbcloud/AgentTeam/internal/store/app/runtime"
)

func newTestAgentRuntime(db coredb.DB, graph *store.Graph) *agentruntime.Service {
	return agentruntime.New(agentruntime.Dependencies{
		Agents:   graph.Agent,
		Messages: graph.Message,
		Recovery: graph.Rewind,
	})
}

type recordingEmitter struct {
	machineCapabilities []struct {
		runtimes []string
		defaults map[string]runtime.RuntimeDefaults
	}
	sessions             []struct{ agentID, sessionID string }
	trajectories         []json.RawMessage
	trajectoryTimestamps []int64
	trajectoryOwners     []string
	sessionIDs           []string
	wakeStatuses         []struct {
		agentID        string
		agentName      string
		displayName    string
		messageID      string
		wakeToken      string
		target         string
		status         string
		replyMessageID string
		detail         string
		timestamp      string
	}
}

func (emitter *recordingEmitter) EmitMachineStatus(string, string, string) {}

func (emitter *recordingEmitter) EmitMachineCapabilities(_ string, _ string, runtimes []string, defaults map[string]runtime.RuntimeDefaults, _ string, _ string, _ string) {
	emitter.machineCapabilities = append(emitter.machineCapabilities, struct {
		runtimes []string
		defaults map[string]runtime.RuntimeDefaults
	}{
		runtimes: append([]string(nil), runtimes...),
		defaults: defaults,
	})
}

func (emitter *recordingEmitter) EmitDaemonStatus(string, string, string) {}

func (emitter *recordingEmitter) EmitAgentStatusWithInstance(string, string, string, string, *dto.AppAgentInstance) {
}

func (emitter *recordingEmitter) EmitAgentActivityWithInstance(string, string, string, string, string, *dto.AppAgentInstance) {
}

func (emitter *recordingEmitter) EmitAgentSessionWithInstance(_ string, _ string, agentID string, sessionID string, _ *dto.AppAgentInstance) {
	emitter.sessions = append(emitter.sessions, struct{ agentID, sessionID string }{
		agentID:   agentID,
		sessionID: sessionID,
	})
}

func (emitter *recordingEmitter) EmitAgentWakeStatus(_ string, agentID string, agentName string, displayName string, messageID string, wakeToken string, target string, status string, replyMessageID string, detail string, timestamp string) {
	emitter.wakeStatuses = append(emitter.wakeStatuses, struct {
		agentID        string
		agentName      string
		displayName    string
		messageID      string
		wakeToken      string
		target         string
		status         string
		replyMessageID string
		detail         string
		timestamp      string
	}{
		agentID:        agentID,
		agentName:      agentName,
		displayName:    displayName,
		messageID:      messageID,
		wakeToken:      wakeToken,
		target:         target,
		status:         status,
		replyMessageID: replyMessageID,
		detail:         detail,
		timestamp:      timestamp,
	})
}

func (emitter *recordingEmitter) EmitAgentTrajectory(_ string, ownerUserID string, _ string, sessionID string, timestamp int64, entries json.RawMessage) {
	emitter.sessionIDs = append(emitter.sessionIDs, sessionID)
	emitter.trajectoryTimestamps = append(emitter.trajectoryTimestamps, timestamp)
	emitter.trajectoryOwners = append(emitter.trajectoryOwners, ownerUserID)
	emitter.trajectories = append(emitter.trajectories, entries)
}

func TestReadyFiltersUnsupportedRuntimeAliasesAndKeepsKimi(t *testing.T) {
	ctx, pool, dataStore := setupReadyFilterStore(t)
	emitter := &recordingEmitter{}
	hub := &Hub{store: NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart), realtime: emitter}
	client := &clientState{
		machineID:     "mac_ready_filter",
		workspaceID:   "wsp_ready_filter",
		runningAgents: make(map[string]bool),
	}

	payload, err := json.Marshal(map[string]any{
		"type":     "ready",
		"runtimes": []string{"codex", "codex-app-server", "claude", "kimi"},
		"runtimeDefaults": map[string]any{
			"codex-app-server": map[string]any{"model": "gpt-5-app-server", "reasoningEffort": "medium"},
			"codex":            map[string]any{"model": "gpt-5-codex", "reasoningEffort": "high"},
			"claude":           map[string]any{"model": "sonnet", "reasoningEffort": ""},
			"kimi":             map[string]any{"model": "kimi-k2", "reasoningEffort": "high"},
		},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	require.Len(t, emitter.machineCapabilities, 1)
	require.Equal(t, []string{"codex", "claude", "kimi"}, emitter.machineCapabilities[0].runtimes)
	require.NotContains(t, emitter.machineCapabilities[0].defaults, "codex-app-server")
	require.Equal(t, "kimi-k2", emitter.machineCapabilities[0].defaults["kimi"].Model)
	require.Zero(t, loadReadyFilterCapabilityCount(t, ctx, pool, "mac_ready_filter"))
}

func TestReadyDeduplicatesRuntimes(t *testing.T) {
	ctx, pool, dataStore := setupReadyFilterStore(t)
	emitter := &recordingEmitter{}
	hub := &Hub{store: NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart), realtime: emitter}
	client := &clientState{
		machineID:     "mac_ready_filter",
		workspaceID:   "wsp_ready_filter",
		runningAgents: make(map[string]bool),
	}

	payload, err := json.Marshal(map[string]any{
		"type":     "ready",
		"runtimes": []string{"codex", "codex", " codex "},
		"runtimeDefaults": map[string]any{
			"codex": map[string]any{"model": "gpt-5-codex", "reasoningEffort": "high"},
		},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "1.0.0",
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	require.Len(t, emitter.machineCapabilities, 1)
	require.Equal(t, []string{"codex"}, emitter.machineCapabilities[0].runtimes)
	require.Equal(t, "gpt-5-codex", emitter.machineCapabilities[0].defaults["codex"].Model)
	require.Zero(t, loadReadyFilterCapabilityCount(t, ctx, pool, "mac_ready_filter"))
}

func TestReadyPersistsDaemonVersion(t *testing.T) {
	ctx, pool, dataStore := setupReadyFilterStore(t)
	hub := &Hub{store: NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart)}
	client := &clientState{
		machineID:     "mac_ready_filter",
		workspaceID:   "wsp_ready_filter",
		runningAgents: make(map[string]bool),
	}

	payload, err := json.Marshal(map[string]any{
		"type":          "ready",
		"runtimes":      []string{"claude"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "agentteam/v1.0.2-rc.11",
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	var daemonVersion string
	require.NoError(t, pool.QueryRow(ctx, `
SELECT daemon_version
FROM machines
WHERE id = $1
`, "mac_ready_filter").Scan(&daemonVersion))
	require.Equal(t, "agentteam/v1.0.2-rc.11", daemonVersion)
}

func TestReadyTracksReportedRunningAgents(t *testing.T) {
	ctx, pool, dataStore := setupReadyFilterStore(t)
	hub := &Hub{store: NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart)}
	client := &clientState{
		machineID:     "mac_ready_filter",
		workspaceID:   "wsp_ready_filter",
		runningAgents: map[string]bool{"agt_stale": true},
	}

	payload, err := json.Marshal(map[string]any{
		"type":          "ready",
		"capabilities":  []string{"agent:permission-mode"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{"agt_live", " ", "agt_live"},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "agentteam/v1.0.2",
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	require.True(t, client.isAgentRunning("agt_live"))
	require.False(t, client.isAgentRunning("agt_stale"))
}

func TestReadyRunningAgentsLimitAppliesBeforePermissionChecks(t *testing.T) {
	agentIDs := make([]string, 0, daemonReadyRunningAgentsLimit+1)
	for i := 0; i <= daemonReadyRunningAgentsLimit; i++ {
		agentIDs = append(agentIDs, fmt.Sprintf("agt_%d", i))
	}
	client := &clientState{capabilities: map[string]struct{}{}}

	_, _, err := (&Hub{}).filterRunnableAgentIDs(context.Background(), client, agentIDs)

	require.ErrorContains(t, err, "daemon ready runningAgents exceeds limit")
}

func TestStartCommandRequiresMaaSDirectCapability(t *testing.T) {
	client := &clientState{capabilities: map[string]struct{}{}}
	command := agent.StartCommand{
		Config: agent.StartConfig{CredentialKind: agent.DaemonLLMCredentialKindMaaSCapability},
	}

	require.ErrorIs(t, client.requireMaaSDirectCapabilityLocked(command), runtimecontrol.ErrDaemonMaaSDirectUpgradeRequired)

	client.capabilities[agentMaaSDirectCapability] = struct{}{}
	require.NoError(t, client.requireMaaSDirectCapabilityLocked(command))
	require.NoError(t, client.requireMaaSDirectCapabilityLocked(agent.StartCommand{}))
}

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

	runtimes, defaults := filterReadyRuntimePayload(
		[]string{"builtin", "unknown", " legacy ", "opencode"},
		map[string]runtime.RuntimeDefaults{
			"builtin":  {Model: "legacy", ReasoningEffort: "medium"},
			"opencode": {Model: "open", ReasoningEffort: ""},
			"unknown":  {Model: "ignored", ReasoningEffort: ""},
		},
	)

	require.Equal(t, []string{"builtin", "opencode"}, runtimes)
	require.Equal(t, map[string]runtime.RuntimeDefaults{
		"builtin":  {Model: "legacy", ReasoningEffort: "medium"},
		"opencode": {Model: "open", ReasoningEffort: ""},
	}, defaults)
}

func TestWorkspaceTargetPathRequiresDaemonCapability(t *testing.T) {
	hub := &Hub{
		clients: map[string]*clientState{
			"mac_target_path": {
				machineID:    "mac_target_path",
				capabilities: make(map[string]struct{}),
			},
		},
	}

	_, err := hub.ListAgentOutputTargetFiles(
		context.Background(),
		"mac_target_path",
		"agt_target_path",
		"/tmp/review-output",
		"",
	)
	require.ErrorIs(t, err, runtimecontrol.ErrDaemonWorkspaceTargetPathUpgradeRequired)

	_, err = hub.ReadAgentOutputTargetFile(
		context.Background(),
		"mac_target_path",
		"agt_target_path",
		"/tmp/review-output",
		"report.md",
	)
	require.ErrorIs(t, err, runtimecontrol.ErrDaemonWorkspaceTargetPathUpgradeRequired)
}

func TestHandleMessageRejectsNullWorkspaceFileTreeFiles(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	treeCh := make(chan fileTreeResult, 1)
	client := &clientState{pendingTree: treeCh}

	payload, err := json.Marshal(map[string]any{
		"type":    "agent:workspace:file_tree",
		"agentId": "agt_test",
		"files":   nil,
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "daemon protocol violation")

	select {
	case result := <-treeCh:
		require.ErrorContains(t, result.err, "daemon protocol violation")
	default:
		require.FailNow(t, "expected daemon protocol to publish file tree error result")
	}
}

func TestHandleMessagePublishesWorkspaceFileTreeDecodeErrorToWaiter(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	treeCh := make(chan fileTreeResult, 1)
	client := &clientState{pendingTree: treeCh}

	payload, err := json.Marshal(map[string]any{
		"type":    "agent:workspace:file_tree",
		"agentId": "agt_test",
		"files": []map[string]any{
			{
				"name": "missing-is-directory.txt",
				"path": "missing-is-directory.txt",
			},
		},
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "decode workspace file tree")

	select {
	case result := <-treeCh:
		require.ErrorContains(t, result.err, "decode workspace file tree")
	default:
		require.FailNow(t, "expected daemon protocol to publish file tree error result")
	}
}

func TestHandleMessagePublishesWorkspaceFileTreePageMetadata(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	treeCh := make(chan fileTreeResult, 1)
	client := &clientState{pendingTree: treeCh}

	payload, err := json.Marshal(map[string]any{
		"type":       "agent:workspace:file_tree",
		"agentId":    "agt_test",
		"limit":      2,
		"cursor":     "4",
		"hasMore":    true,
		"nextCursor": "6",
		"files": []map[string]any{
			{
				"name":        "notes.md",
				"path":        "notes.md",
				"isDirectory": false,
			},
		},
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	select {
	case result := <-treeCh:
		require.NoError(t, result.err)
		require.Len(t, result.page.Files, 1)
		require.Equal(t, 2, result.page.Limit)
		require.Equal(t, "4", result.page.Cursor)
		require.True(t, result.page.HasMore)
		require.Equal(t, "6", result.page.NextCursor)
	case <-time.After(time.Second):
		require.FailNow(t, "expected daemon protocol to publish file tree result")
	}
}

func TestHandleMessageRejectsNullWorkspaceScanDirectories(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	scanCh := make(chan scanResult, 1)
	client := &clientState{pendingScan: scanCh}

	payload, err := json.Marshal(map[string]any{
		"type":        "machine:workspace:scan_result",
		"directories": nil,
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "daemon protocol violation")

	select {
	case result := <-scanCh:
		require.ErrorContains(t, result.err, "daemon protocol violation")
	default:
		require.FailNow(t, "expected daemon protocol to publish scan error result")
	}
}

func TestHandleMessagePublishesWorkspaceScanPageMetadata(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	scanCh := make(chan scanResult, 1)
	client := &clientState{pendingScan: scanCh}

	payload, err := json.Marshal(map[string]any{
		"type":       "machine:workspace:scan_result",
		"limit":      2,
		"cursor":     "4",
		"hasMore":    true,
		"nextCursor": "6",
		"directories": []map[string]any{
			{
				"directoryName": "agt_test",
				"fileCount":     1,
			},
		},
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	select {
	case result := <-scanCh:
		require.NoError(t, result.err)
		require.Len(t, result.page.Directories, 1)
		require.Equal(t, 2, result.page.Limit)
		require.Equal(t, "4", result.page.Cursor)
		require.True(t, result.page.HasMore)
		require.Equal(t, "6", result.page.NextCursor)
	case <-time.After(time.Second):
		require.FailNow(t, "expected daemon protocol to publish scan result")
	}
}

func TestHandleMessagePublishesWorkspaceFileContentDecodeErrorToWaiter(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	readCh := make(chan fileContentResult, 1)
	client := &clientState{readRequests: map[string]chan fileContentResult{"req_test": readCh}}

	payload, err := json.Marshal(map[string]any{
		"type":      "agent:workspace:file_content",
		"agentId":   "agt_test",
		"requestId": "req_test",
		"content":   map[string]any{"oops": true},
		"binary":    false,
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "decode workspace file content")

	select {
	case result := <-readCh:
		require.ErrorContains(t, result.err, "decode workspace file content")
	default:
		require.FailNow(t, "expected daemon protocol to publish file content error result")
	}
}

func TestHandleMessageRejectsNullAgentTrajectoryEntries(t *testing.T) {
	ctx := context.Background()
	emitter := &recordingEmitter{}
	hub := &Hub{realtime: emitter}
	client := &clientState{workspaceID: "wsp_test"}

	payload, err := json.Marshal(map[string]any{
		"type":    "agent:trajectory",
		"agentId": "agt_test",
		"entries": nil,
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "daemon protocol violation")
	require.Empty(t, emitter.trajectories)
}

func TestHandleMessageEmitsAgentTrajectoryForArrayEntries(t *testing.T) {
	ctx := context.Background()
	emitter := &recordingEmitter{}
	hub := &Hub{store: &recordingStartStore{}, realtime: emitter}
	client := &clientState{
		machineID:       "mac_test",
		workspaceID:     "wsp_test",
		createdByUserID: "usr_owner",
	}

	payload, err := json.Marshal(map[string]any{
		"type":      "agent:trajectory",
		"agentId":   "agt_test",
		"sessionId": "sess-1",
		"timestamp": int64(1780287800000),
		"entries":   []map[string]any{{"kind": "text", "text": "hello"}},
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))
	require.Len(t, emitter.trajectories, 1)
	require.Equal(t, []string{"usr_owner"}, emitter.trajectoryOwners)
	require.Equal(t, []string{"sess-1"}, emitter.sessionIDs)
	require.Equal(t, []int64{1780287800000}, emitter.trajectoryTimestamps)
	require.NotEmpty(t, emitter.trajectories[0])
	require.Equal(t, byte('['), emitter.trajectories[0][0])
}

func TestHandleMessageIgnoresAgentTrajectoryFromUnexpectedMachine(t *testing.T) {
	ctx := context.Background()
	emitter := &recordingEmitter{}
	agentBelongsToMachine := false
	store := &recordingStartStore{agentBelongsToMachine: &agentBelongsToMachine}
	hub := &Hub{store: store, realtime: emitter}
	client := &clientState{
		machineID:   "mac_wrong",
		workspaceID: "wsp_test",
		logger:      slog.New(slog.NewTextHandler(io.Discard, nil)),
	}

	payload, err := json.Marshal(map[string]any{
		"type":      "agent:trajectory",
		"agentId":   " agt_test ",
		"sessionId": "sess-1",
		"timestamp": int64(1780287800000),
		"entries":   []map[string]any{{"kind": "text", "text": "hello"}},
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))
	require.Empty(t, emitter.trajectories)
	require.Empty(t, emitter.sessionIDs)
	require.Empty(t, emitter.trajectoryTimestamps)
}

func TestHandleMessageRejectsIncompleteAgentWakeStatus(t *testing.T) {
	ctx := context.Background()
	emitter := &recordingEmitter{}
	hub := &Hub{realtime: emitter}
	client := &clientState{workspaceID: "wsp_test"}

	payload, err := json.Marshal(map[string]any{
		"type":      "agent:wake_status",
		"agentId":   "agt_test",
		"messageId": "msg-1",
		"status":    "received",
		"timestamp": "2026-03-27T12:00:00Z",
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "daemon protocol violation")
	require.Empty(t, emitter.wakeStatuses)
}

func TestHandleMessageEmitsAgentWakeStatus(t *testing.T) {
	ctx := context.Background()
	emitter := &recordingEmitter{}
	hub := &Hub{realtime: emitter}
	client := &clientState{workspaceID: "wsp_test"}

	payload, err := json.Marshal(map[string]any{
		"type":           "agent:wake_status",
		"agentId":        "agt_test",
		"messageId":      "msg-1",
		"wakeToken":      "wake-1",
		"target":         "#meta-agentteam:00000520",
		"status":         "reply_sent",
		"replyMessageId": "msg-reply",
		"detail":         "confirmed",
		"timestamp":      "2026-03-27T12:00:00Z",
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))
	require.Len(t, emitter.wakeStatuses, 1)
	require.Equal(t, "agt_test", emitter.wakeStatuses[0].agentID)
	require.Equal(t, "msg-1", emitter.wakeStatuses[0].messageID)
	require.Equal(t, "wake-1", emitter.wakeStatuses[0].wakeToken)
	require.Equal(t, "reply_sent", emitter.wakeStatuses[0].status)
	require.Equal(t, "msg-reply", emitter.wakeStatuses[0].replyMessageID)
}

func setupReadyFilterStore(t *testing.T) (context.Context, coredb.DB, *store.Graph) {
	t.Helper()

	ctx := context.Background()
	pool := openDaemonTestTiDB(t)
	archiveStorage, err := archive.NewFilesystem(t.TempDir())
	require.NoError(t, err)
	dataStore := store.New(pool, archiveStorage, nil, 30)

	mustExecProtocolValidation(t, pool, `
INSERT INTO users (id, email, name, display_name, avatar_url, is_admin, password_hash, email_verified_at, created_at, updated_at)
VALUES ('usr_ready_filter', 'ready-filter@example.com', 'ready-filter', 'Ready Filter', '', TRUE, 'hash', NOW(), NOW(), NOW())
`)
	mustExecProtocolValidation(t, pool, `
INSERT INTO organizations (id, slug, name, created_by_user_id, created_at, updated_at)
VALUES ('org_wsp_ready_filter', 'ready-filter', 'ready-filter', 'usr_ready_filter', NOW(), NOW())
`)
	mustExecProtocolValidation(t, pool, `
INSERT INTO workspaces (id, organization_id, slug, name, latest_message_seq, created_at, updated_at)
VALUES ('wsp_ready_filter', 'org_wsp_ready_filter', 'ready-filter', 'ready-filter', 0, NOW(), NOW())
`)
	mustExecProtocolValidation(t, pool, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, created_at, updated_at)
VALUES ('mac_ready_filter', 'wsp_ready_filter', 'usr_ready_filter', 'Ready filter machine', 'machine-key', NOW(), NOW())
`)

	return ctx, pool, dataStore
}

func loadReadyFilterCapabilityCount(t *testing.T, ctx context.Context, pool coredb.Queryer, machineID string) int {
	t.Helper()

	var count int
	err := pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM machine_runtime_capabilities
WHERE machine_id = $1
`, machineID).Scan(&count)
	require.NoError(t, err)
	return count
}

func mustExecProtocolValidation(t *testing.T, pool coredb.QueryExecer, query string, args ...any) {
	t.Helper()

	_, err := pool.Exec(context.Background(), query, args...)
	require.NoError(t, err)
}

func TestHandleMessageRejectsIncompleteAgentRecoveryPointPrepared(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}
	client := &clientState{}

	payload, err := json.Marshal(map[string]any{
		"type":      "agent:recovery_point_prepared",
		"agentId":   "agt_test",
		"messageId": "msg-1",
		"strategy":  "codex_app_server_turns",
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "daemon protocol violation")
}

func TestHandleMessagePublishesRestoreRecoveryPointFailureToWaiter(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	restoreCh := make(chan restoreRecoveryPointResult, 1)
	client := &clientState{restoreRequests: map[string]chan restoreRecoveryPointResult{"restore-1": restoreCh}}

	payload, err := json.Marshal(map[string]any{
		"type":      "agent:restore_recovery_point:result",
		"agentId":   "agt_test",
		"requestId": "restore-1",
		"success":   false,
		"error":     "restore failed",
	})
	require.NoError(t, err)

	require.NoError(t, hub.handleMessage(ctx, client, payload))

	select {
	case result := <-restoreCh:
		require.ErrorContains(t, result.err, "restore failed")
	default:
		require.FailNow(t, "expected daemon protocol to publish restore failure result")
	}
}

func TestHandleMessageFailsAllRestoreWaitersOnMalformedResult(t *testing.T) {
	ctx := context.Background()
	hub := &Hub{}

	restoreCh := make(chan restoreRecoveryPointResult, 1)
	client := &clientState{restoreRequests: map[string]chan restoreRecoveryPointResult{"restore-1": restoreCh}}

	payload, err := json.Marshal(map[string]any{
		"type":    "agent:restore_recovery_point:result",
		"agentId": "agt_test",
		"success": true,
	})
	require.NoError(t, err)

	require.ErrorContains(t, hub.handleMessage(ctx, client, payload), "daemon protocol violation")

	select {
	case result := <-restoreCh:
		require.ErrorContains(t, result.err, "daemon protocol violation")
	default:
		require.FailNow(t, "expected daemon protocol to fail outstanding restore waiters")
	}
}