package daemon

import (
	"context"
	"io"
	"log/slog"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/archive"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
)

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

	ctx := context.Background()
	pool := openDaemonTestTiDB(t)

	archiveStorage, err := archive.NewFilesystem(t.TempDir())
	require.NoError(t, err)

	dataStore := store.New(pool, archiveStorage, nil, 30)
	emitter := &recordingEmitter{}
	hub := &Hub{
		logger:        slog.New(slog.NewTextHandler(io.Discard, nil)),
		store:         NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart),
		runtimeEvents: dataStore.AgentOps,
		realtime:      emitter,
	}

	now := time.Now().UTC()
	workspaceID := "wsp_test"
	userID := "usr_test"
	memberID := "mem_agent"
	expectedMachineID := "mac_expected"
	otherMachineID := "mac_other"
	agentID := "agt_expected"

	mustExecDaemonTest(t, pool, `
INSERT INTO users (id, email, name, display_name, avatar_url, is_admin, password_hash, email_verified_at, created_at, updated_at)
VALUES ($1, 'test@example.com', 'test', 'Test', '', TRUE, 'hash', $2, $2, $2)
`, userID, now)
	organizationID := seedDaemonTestWorkspaceOrganization(t, pool, workspaceID, "test", userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO workspaces (id, organization_id, slug, name, latest_message_seq, created_at, updated_at)
VALUES ($1, $2, 'test', 'Test', 0, $3, $3)
`, workspaceID, organizationID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO members (id, workspace_id, type, display_name, created_at, updated_at)
VALUES ($1, $2, 'agent', 'Agent', $3, $3)
`, memberID, workspaceID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, created_at, updated_at)
VALUES ($1, $2, $3, 'Expected machine', 'key-1', $4, $4)
`, expectedMachineID, workspaceID, userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, created_at, updated_at)
VALUES ($1, $2, $3, 'Other machine', 'key-2', $4, $4)
`, otherMachineID, workspaceID, userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO agents (id, workspace_id, machine_id, member_id, name, runtime, model, description, auth_token_hash, auto_wake_enabled, status, activity, activity_detail, created_at, updated_at)
VALUES ($1, $2, $3, $4, 'agent-1', 'codex', 'gpt-5-codex', '', 'hash', TRUE, 'active', 'online', '', $5, $5)
`, agentID, workspaceID, expectedMachineID, memberID, now)

	payload := []byte(`{"type":"agent:wake_status","agentId":"agt_expected","messageId":"msg-1","wakeToken":"wake-1","status":"reply_sent","timestamp":"2026-03-27T12:00:00Z"}`)
	client := &clientState{
		workspaceID: workspaceID,
		machineID:   otherMachineID,
		logger:      slog.New(slog.NewTextHandler(io.Discard, nil)),
	}

	require.NoError(t, hub.handleMessage(ctx, client, payload))
	require.Empty(t, emitter.wakeStatuses)
}

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

	ctx := context.Background()
	pool := openDaemonTestTiDB(t)

	archiveStorage, err := archive.NewFilesystem(t.TempDir())
	require.NoError(t, err)

	dataStore := store.New(pool, archiveStorage, nil, 30)
	emitter := &recordingEmitter{}
	hub := &Hub{
		logger:        slog.New(slog.NewTextHandler(io.Discard, nil)),
		store:         NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart),
		runtimeEvents: dataStore.AgentOps,
		realtime:      emitter,
	}

	now := time.Now().UTC()
	workspaceID := "wsp_test"
	userID := "usr_test"
	memberID := "mem_agent"
	machineID := "mac_expected"
	agentID := "agt_expected"

	mustExecDaemonTest(t, pool, `
INSERT INTO users (id, email, name, display_name, avatar_url, is_admin, password_hash, email_verified_at, created_at, updated_at)
VALUES ($1, 'test@example.com', 'test', 'Test', '', TRUE, 'hash', $2, $2, $2)
`, userID, now)
	organizationID := seedDaemonTestWorkspaceOrganization(t, pool, workspaceID, "test", userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO workspaces (id, organization_id, slug, name, latest_message_seq, created_at, updated_at)
VALUES ($1, $2, 'test', 'Test', 0, $3, $3)
`, workspaceID, organizationID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO members (id, workspace_id, type, display_name, created_at, updated_at)
VALUES ($1, $2, 'agent', 'Agent', $3, $3)
`, memberID, workspaceID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, created_at, updated_at)
VALUES ($1, $2, $3, 'Expected machine', 'key-1', $4, $4)
`, machineID, workspaceID, userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO agents (id, workspace_id, machine_id, member_id, name, runtime, model, description, auth_token_hash, auto_wake_enabled, status, activity, activity_detail, created_at, updated_at)
VALUES ($1, $2, $3, $4, 'agent-1', 'codex', 'gpt-5-codex', '', 'hash', TRUE, 'active', 'online', '', $5, $5)
`, agentID, workspaceID, machineID, memberID, now)

	payload := []byte(`{"type":"agent:wake_status","agentId":"agt_expected","messageId":"msg-1","wakeToken":"wake-1","target":"#all","status":"reply_sent","replyMessageId":"msg-reply","timestamp":"2026-04-01T12:00:00Z"}`)
	client := &clientState{
		workspaceID: workspaceID,
		machineID:   machineID,
		logger:      slog.New(slog.NewTextHandler(io.Discard, nil)),
	}

	require.NoError(t, hub.handleMessage(ctx, client, payload))
	require.Len(t, emitter.wakeStatuses, 1)
	require.Equal(t, "agent-1", emitter.wakeStatuses[0].agentName)
	require.Equal(t, "Agent", emitter.wakeStatuses[0].displayName)

	var eventClass string
	var storedMessageID string
	var storedWakeToken string
	var storedReplyMessageID string
	require.NoError(t, pool.QueryRow(ctx, `
SELECT event_class, COALESCE(message_id, ''), COALESCE(wake_token, ''), COALESCE(reply_message_id, '')
FROM agent_runtime_events
WHERE workspace_id = $1 AND agent_id = $2
ORDER BY occurred_at DESC, created_at DESC
LIMIT 1
`, workspaceID, agentID).Scan(&eventClass, &storedMessageID, &storedWakeToken, &storedReplyMessageID))
	require.Equal(t, agent.AgentRuntimeEventClassVisibleReplySent, eventClass)
	require.Equal(t, "msg-1", storedMessageID)
	require.Equal(t, "wake-1", storedWakeToken)
	require.Equal(t, "msg-reply", storedReplyMessageID)
}

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

	ctx := context.Background()
	pool := openDaemonTestTiDB(t)

	archiveStorage, err := archive.NewFilesystem(t.TempDir())
	require.NoError(t, err)

	dataStore := store.New(pool, archiveStorage, nil, 30)
	emitter := &recordingEmitter{}
	hub := &Hub{
		logger:        slog.New(slog.NewTextHandler(io.Discard, nil)),
		store:         NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart),
		runtimeEvents: dataStore.AgentOps,
		realtime:      emitter,
	}

	now := time.Now().UTC()
	workspaceID := "wsp_transient"
	userID := "usr_transient"
	memberID := "mem_transient"
	machineID := "mac_transient"
	agentID := "agt_transient"

	mustExecDaemonTest(t, pool, `
INSERT INTO users (id, email, name, display_name, avatar_url, is_admin, password_hash, email_verified_at, created_at, updated_at)
VALUES ($1, 'transient@example.com', 'transient', 'Transient', '', TRUE, 'hash', $2, $2, $2)
`, userID, now)
	organizationID := seedDaemonTestWorkspaceOrganization(t, pool, workspaceID, "transient", userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO workspaces (id, organization_id, slug, name, latest_message_seq, created_at, updated_at)
VALUES ($1, $2, 'transient', 'Transient', 0, $3, $3)
`, workspaceID, organizationID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO members (id, workspace_id, type, display_name, created_at, updated_at)
VALUES ($1, $2, 'agent', 'Transient Agent', $3, $3)
`, memberID, workspaceID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, created_at, updated_at)
VALUES ($1, $2, $3, 'Transient machine', 'key-transient', $4, $4)
`, machineID, workspaceID, userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO agents (id, workspace_id, machine_id, member_id, name, runtime, model, description, auth_token_hash, auto_wake_enabled, status, activity, activity_detail, created_at, updated_at)
VALUES ($1, $2, $3, $4, 'transient-agent', 'codex', 'gpt-5-codex', '', 'hash', TRUE, 'active', 'online', '', $5, $5)
`, agentID, workspaceID, machineID, memberID, now)

	client := &clientState{
		workspaceID: workspaceID,
		machineID:   machineID,
		logger:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		capabilities: map[string]struct{}{
			transientWakeStatusDesktopCapability: {},
		},
	}
	require.NoError(t, hub.handleMessage(ctx, client, []byte(`{
		"type":"agent:wake_status",
		"agentId":"agt_transient",
		"messageId":"msg-transient-negotiated",
		"wakeToken":"wake-transient-negotiated",
		"target":"#all",
		"status":"reply_sent",
		"replyMessageId":"msg-reply-negotiated",
		"timestamp":"2026-08-19T06:00:00Z",
		"transient":true
	}`)))

	require.Len(t, emitter.wakeStatuses, 1)
	require.Equal(t, "msg-transient-negotiated", emitter.wakeStatuses[0].messageID)
	require.Equal(t, int64(0), countWakeStatusRuntimeEvents(
		t,
		ctx,
		pool,
		workspaceID,
		agentID,
		"msg-transient-negotiated",
		"wake-transient-negotiated",
	))

	err = hub.handleMessage(ctx, client, []byte(`{
		"type":"agent:wake_status",
		"agentId":"agt_transient",
		"messageId":"msg-transient-invalid-time",
		"wakeToken":"wake-transient-invalid-time",
		"target":"#all",
		"status":"reply_sent",
		"timestamp":"not-a-timestamp",
		"transient":true
	}`))
	require.ErrorContains(t, err, "parse agent wake status timestamp")
	require.Len(t, emitter.wakeStatuses, 1)

	client.setCapabilities(nil)
	require.NoError(t, hub.handleMessage(ctx, client, []byte(`{
		"type":"agent:wake_status",
		"agentId":"agt_transient",
		"messageId":"msg-transient-forged",
		"wakeToken":"wake-transient-forged",
		"target":"#all",
		"status":"reply_sent",
		"replyMessageId":"msg-reply-forged",
		"timestamp":"2026-08-19T06:01:00Z",
		"transient":true
	}`)))

	require.Len(t, emitter.wakeStatuses, 2)
	require.Equal(t, "msg-transient-forged", emitter.wakeStatuses[1].messageID)
	require.Equal(t, int64(1), countWakeStatusRuntimeEvents(
		t,
		ctx,
		pool,
		workspaceID,
		agentID,
		"msg-transient-forged",
		"wake-transient-forged",
	))
}

func countWakeStatusRuntimeEvents(
	t *testing.T,
	ctx context.Context,
	pool coredb.QueryExecer,
	workspaceID string,
	agentID string,
	messageID string,
	wakeToken string,
) int64 {
	t.Helper()

	var count int64
	require.NoError(t, pool.QueryRow(ctx, `
SELECT COUNT(*)
FROM agent_runtime_events
WHERE workspace_id = $1
  AND agent_id = $2
  AND message_id = $3
  AND wake_token = $4
`, workspaceID, agentID, messageID, wakeToken).Scan(&count))
	return count
}

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

	ctx := context.Background()
	pool := openDaemonTestTiDB(t)

	archiveStorage, err := archive.NewFilesystem(t.TempDir())
	require.NoError(t, err)

	dataStore := store.New(pool, archiveStorage, nil, 30)
	hub := &Hub{
		logger:        slog.New(slog.NewTextHandler(io.Discard, nil)),
		store:         NewStore(newTestAgentRuntime(pool, dataStore).Agents(), newTestAgentRuntime(pool, dataStore).Recovery(), newTestAgentRuntime(pool, dataStore).PrepareAgentStart),
		runtimeEvents: dataStore.AgentOps,
	}

	now := time.Now().UTC()
	workspaceID := "wsp_test"
	userID := "usr_test"
	memberID := "mem_agent"
	machineID := "mac_expected"
	agentID := "agt_expected"

	mustExecDaemonTest(t, pool, `
INSERT INTO users (id, email, name, display_name, avatar_url, is_admin, password_hash, email_verified_at, created_at, updated_at)
VALUES ($1, 'test@example.com', 'test', 'Test', '', TRUE, 'hash', $2, $2, $2)
`, userID, now)
	organizationID := seedDaemonTestWorkspaceOrganization(t, pool, workspaceID, "test", userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO workspaces (id, organization_id, slug, name, latest_message_seq, created_at, updated_at)
VALUES ($1, $2, 'test', 'Test', 0, $3, $3)
`, workspaceID, organizationID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO members (id, workspace_id, type, display_name, created_at, updated_at)
VALUES ($1, $2, 'agent', 'Agent', $3, $3)
`, memberID, workspaceID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO machines (id, workspace_id, created_by_user_id, name, api_key, created_at, updated_at)
VALUES ($1, $2, $3, 'Expected machine', 'key-1', $4, $4)
`, machineID, workspaceID, userID, now)
	mustExecDaemonTest(t, pool, `
INSERT INTO agents (id, workspace_id, machine_id, member_id, name, runtime, model, description, auth_token_hash, auto_wake_enabled, status, activity, activity_detail, created_at, updated_at)
VALUES ($1, $2, $3, $4, 'agent-1', 'codex', 'gpt-5-codex', '', 'hash', TRUE, 'active', 'online', '', $5, $5)
`, agentID, workspaceID, machineID, memberID, now)

	payload := []byte(`{"type":"agent:runtime_event","agentId":"agt_expected","messageId":"msg-1","eventClass":"runtime_retrying","target":"#all","timestamp":"2026-04-01T12:00:00Z"}`)
	client := &clientState{
		workspaceID: workspaceID,
		machineID:   machineID,
		logger:      slog.New(slog.NewTextHandler(io.Discard, nil)),
	}

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

	var eventClass string
	var storedTarget string
	require.NoError(t, pool.QueryRow(ctx, `
SELECT event_class, target
FROM agent_runtime_events
WHERE workspace_id = $1 AND agent_id = $2
ORDER BY occurred_at DESC, created_at DESC
LIMIT 1
`, workspaceID, agentID).Scan(&eventClass, &storedTarget))
	require.Equal(t, agent.AgentRuntimeEventClassRuntimeRetrying, eventClass)
	require.Equal(t, "#all", storedTarget)
}

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

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

func seedDaemonTestWorkspaceOrganization(t *testing.T, pool coredb.QueryExecer, workspaceID string, slug string, createdByUserID string, now time.Time) string {
	t.Helper()

	organizationID := "org_" + workspaceID
	mustExecDaemonTest(t, pool, `
INSERT INTO organizations (id, slug, name, created_by_user_id, created_at, updated_at)
VALUES ($1, $2, $2, $3, $4, $4)
`, organizationID, slug, createdByUserID, now)
	return organizationID
}