package daemon_test

import (
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/runtimecontrol"
	"github.com/tidbcloud/AgentTeam/internal/store/identity"
	"github.com/tidbcloud/AgentTeam/internal/store/message"
)

func TestHubStartWakeMessageDeliversWakePayloadDirectly(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-replay-start")

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	}))

	delivery := sendMentionDeliveryForFixture(t, fixture, "@"+fixture.agentHandle+" replay via start path")
	require.Len(t, delivery.WakeRecipients, 1)

	require.NoError(t, fixture.hub.NotifyAgentMessage(fixture.ctx, delivery.WakeRecipients[0], delivery.Seq, delivery.Message))

	start := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:start", start["type"])
	require.Equal(t, fixture.agentID, start["agentId"])

	wakeMessage, ok := start["wakeMessage"].(map[string]any)
	require.True(t, ok)
	require.Equal(t, delivery.MessageID, wakeMessage["message_id"])
	require.Equal(t, delivery.Message.Content, wakeMessage["content"])
}

func TestHubNotifyAgentMessageBackfillsWakeTargetSeq(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-freshness-seq")

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	}))

	err := fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
		AgentID:         fixture.agentID,
		MachineID:       fixture.machineID,
		AgentMemberID:   "mem_unused",
		AutoWakeEnabled: true,
		DeliveryReason:  "dm_peer",
	}, 42, message.ReceiveMessage{
		MessageID:   "msg_freshness_42",
		TargetKey:   "dm_scope_42",
		ChannelType: "dm",
		ChannelName: fixture.agentHandle + "-owner",
		AuthorProfile: message.AuthorProfile{
			ActorType:   message.AuthorActorTypeHuman,
			TargetID:    "owner",
			DisplayName: "owner",
		},
		Content:   "wake with missing target seq",
		Timestamp: time.Now().UTC(),
	})
	require.NoError(t, err)

	start := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:start", start["type"])
	require.Equal(t, fixture.agentID, start["agentId"])

	wakeMessage, ok := start["wakeMessage"].(map[string]any)
	require.True(t, ok)
	require.Equal(t, "msg_freshness_42", wakeMessage["message_id"])
	require.Equal(t, "dm_scope_42", wakeMessage["target_key"])
	require.Equal(t, float64(42), wakeMessage["target_seq"])
}

func TestHubAcceptedLiveDeliverAckCompletesLiveDelivery(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-replay-deliver")

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	}))
	require.NoError(t, fixture.hub.StartAgent(fixture.ctx, fixture.machineID, fixture.agentID))

	initialStart := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:start", initialStart["type"])
	require.Equal(t, fixture.agentID, initialStart["agentId"])
	require.True(t, fixture.hub.AgentRuntimeRunning(fixture.machineID, fixture.agentID))

	delivery := sendMentionDeliveryForFixture(t, fixture, "@"+fixture.agentHandle+" replay via deliver path")
	require.Len(t, delivery.WakeRecipients, 1)

	sendErrCh := make(chan error, 1)
	go func() {
		sendErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, delivery.WakeRecipients[0], delivery.Seq, delivery.Message)
	}()

	deliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", deliver["type"])
	require.Equal(t, fixture.agentID, deliver["agentId"])

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":      "agent:deliver:ack",
		"agentId":   fixture.agentID,
		"messageId": delivery.MessageID,
		"seq":       int(delivery.Seq),
		"accepted":  true,
	}))

	select {
	case err := <-sendErrCh:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for live deliver notify to finish")
	}
}

func TestHubRunningAgentDeliverIncludesInterruptPolicy(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-interrupt-deliver")

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	}))
	require.NoError(t, fixture.hub.StartAgent(fixture.ctx, fixture.machineID, fixture.agentID))

	initialStart := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:start", initialStart["type"])
	require.Equal(t, fixture.agentID, initialStart["agentId"])

	sendErrCh := make(chan error, 1)
	go func() {
		sendErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_unused",
			AutoWakeEnabled: true,
			DeliveryReason:  "direct_mention",
		}, 77, message.ReceiveMessage{
			MessageID:   "msg_interrupt_77",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:         "interrupt and submit",
			Timestamp:       time.Now().UTC(),
			InterruptPolicy: message.AgentInterruptPolicyCurrentRun,
		})
	}()

	deliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", deliver["type"])
	require.Equal(t, fixture.agentID, deliver["agentId"])
	require.Equal(t, message.AgentInterruptPolicyCurrentRun, deliver["interruptPolicy"])

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":            "agent:deliver:ack",
		"agentId":         fixture.agentID,
		"messageId":       "msg_interrupt_77",
		"seq":             int(deliver["seq"].(float64)),
		"interruptPolicy": message.AgentInterruptPolicyCurrentRun,
		"accepted":        true,
	}))

	select {
	case err := <-sendErrCh:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for interrupt routed message send to finish")
	}
}

func TestHubRunningAgentAcksDoNotCrossNormalAndUrgentSameMessage(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-interrupt-ack-key")

	fixture.writeReady(t, map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{fixture.agentID},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	})

	normalErrCh := make(chan error, 1)
	go func() {
		normalErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_normal",
			AutoWakeEnabled: true,
			DeliveryReason:  message.AgentDeliveryReasonDirectMention,
		}, 88, message.ReceiveMessage{
			MessageID:   "msg_ack_key_88",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:   "normal delivery",
			Timestamp: time.Now().UTC(),
		})
	}()

	normalDeliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", normalDeliver["type"])
	require.Equal(t, fixture.agentID, normalDeliver["agentId"])
	require.Empty(t, normalDeliver["interruptPolicy"])

	urgentErrCh := make(chan error, 1)
	go func() {
		urgentErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_urgent",
			AutoWakeEnabled: false,
			DeliveryReason:  message.AgentDeliveryReasonSpaceDefault,
		}, 88, message.ReceiveMessage{
			MessageID:   "msg_ack_key_88",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:         "urgent delivery",
			Timestamp:       time.Now().UTC(),
			InterruptPolicy: message.AgentInterruptPolicyCurrentRun,
		})
	}()

	urgentDeliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", urgentDeliver["type"])
	require.Equal(t, fixture.agentID, urgentDeliver["agentId"])
	require.Equal(t, message.AgentInterruptPolicyCurrentRun, urgentDeliver["interruptPolicy"])

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":      "agent:deliver:ack",
		"agentId":   fixture.agentID,
		"messageId": "msg_ack_key_88",
		"seq":       int(normalDeliver["seq"].(float64)),
		"accepted":  true,
	}))

	select {
	case err := <-normalErrCh:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for normal delivery ack")
	}
	select {
	case err := <-urgentErrCh:
		t.Fatalf("urgent delivery completed from normal ack: %v", err)
	case <-time.After(150 * time.Millisecond):
	}

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":            "agent:deliver:ack",
		"agentId":         fixture.agentID,
		"messageId":       "msg_ack_key_88",
		"seq":             int(urgentDeliver["seq"].(float64)),
		"interruptPolicy": message.AgentInterruptPolicyCurrentRun,
		"accepted":        true,
	}))

	select {
	case err := <-urgentErrCh:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for urgent delivery ack")
	}
}

func TestHubUrgentReplayRejectedAckReturnsError(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-interrupt-rejected")

	fixture.writeReady(t, map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{fixture.agentID},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	})

	sendErrCh := make(chan error, 1)
	go func() {
		sendErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_no_wake_default",
			AutoWakeEnabled: false,
			DeliveryReason:  message.AgentDeliveryReasonSpaceDefault,
		}, 89, message.ReceiveMessage{
			MessageID:   "msg_interrupt_rejected_89",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:         "urgent replay rejected by desktop",
			Timestamp:       time.Now().UTC(),
			InterruptPolicy: message.AgentInterruptPolicyCurrentRun,
		})
	}()

	deliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", deliver["type"])
	require.Equal(t, fixture.agentID, deliver["agentId"])
	require.Equal(t, message.AgentInterruptPolicyCurrentRun, deliver["interruptPolicy"])

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":            "agent:deliver:ack",
		"agentId":         fixture.agentID,
		"messageId":       "msg_interrupt_rejected_89",
		"seq":             int(deliver["seq"].(float64)),
		"interruptPolicy": message.AgentInterruptPolicyCurrentRun,
		"accepted":        false,
		"detail":          "interrupt delivery failed",
	}))

	select {
	case err := <-sendErrCh:
		require.Error(t, err)
		require.ErrorContains(t, err, "urgent delivery rejected")
		require.ErrorContains(t, err, "interrupt delivery failed")
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for urgent delivery rejection")
	}
}

func TestHubUrgentReplayBypassesAutoWakePreference(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-interrupt-no-wake")

	fixture.writeReady(t, map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	})
	require.NoError(t, fixture.hub.StartAgent(fixture.ctx, fixture.machineID, fixture.agentID))

	initialStart := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:start", initialStart["type"])
	require.Equal(t, fixture.agentID, initialStart["agentId"])

	sendErrCh := make(chan error, 1)
	go func() {
		sendErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_no_wake_default",
			AutoWakeEnabled: false,
			DeliveryReason:  message.AgentDeliveryReasonSpaceDefault,
		}, 78, message.ReceiveMessage{
			MessageID:   "msg_interrupt_no_wake_78",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:         "urgent replay of a no-wake visible message",
			Timestamp:       time.Now().UTC(),
			InterruptPolicy: message.AgentInterruptPolicyCurrentRun,
		})
	}()

	deliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", deliver["type"])
	require.Equal(t, fixture.agentID, deliver["agentId"])
	require.Equal(t, message.AgentInterruptPolicyCurrentRun, deliver["interruptPolicy"])

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":            "agent:deliver:ack",
		"agentId":         fixture.agentID,
		"messageId":       "msg_interrupt_no_wake_78",
		"seq":             int(deliver["seq"].(float64)),
		"interruptPolicy": message.AgentInterruptPolicyCurrentRun,
		"accepted":        true,
	}))

	select {
	case err := <-sendErrCh:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for no-wake interrupt replay to finish")
	}
}

func TestHubUrgentReplayReresolvesClientAfterAckWaitReconnect(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-interrupt-reconnect")

	fixture.writeReady(t, map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{fixture.agentID},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	})

	sendErrCh := make(chan error, 1)
	go func() {
		sendErrCh <- fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_no_wake_default",
			AutoWakeEnabled: false,
			DeliveryReason:  message.AgentDeliveryReasonSpaceDefault,
		}, 79, message.ReceiveMessage{
			MessageID:   "msg_interrupt_reconnect_79",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:         "urgent replay while reconnecting",
			Timestamp:       time.Now().UTC(),
			InterruptPolicy: message.AgentInterruptPolicyCurrentRun,
		})
	}()

	initialDeliver := readWSJSONMessage(t, fixture.conn, 2*time.Second)
	require.Equal(t, "agent:deliver", initialDeliver["type"])
	require.Equal(t, fixture.agentID, initialDeliver["agentId"])
	require.Equal(t, message.AgentInterruptPolicyCurrentRun, initialDeliver["interruptPolicy"])

	reconnected := fixture.connect(t)
	fixture.writeReadyOn(t, reconnected, map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{fixture.agentID},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.1",
	})

	replayedDeliver := readWSJSONMessage(t, reconnected, 4*time.Second)
	require.Equal(t, "agent:deliver", replayedDeliver["type"])
	require.Equal(t, fixture.agentID, replayedDeliver["agentId"])
	require.Equal(t, message.AgentInterruptPolicyCurrentRun, replayedDeliver["interruptPolicy"])

	require.NoError(t, reconnected.WriteJSON(map[string]any{
		"type":            "agent:deliver:ack",
		"agentId":         fixture.agentID,
		"messageId":       "msg_interrupt_reconnect_79",
		"seq":             int(replayedDeliver["seq"].(float64)),
		"interruptPolicy": message.AgentInterruptPolicyCurrentRun,
		"accepted":        true,
	}))

	select {
	case err := <-sendErrCh:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for interrupt replay after reconnect")
	}
}

func TestHubUrgentReplayRequiresConnectedDaemon(t *testing.T) {
	t.Parallel()
	fixture := setupAutoWakeHubFixture(t, "agent-interrupt-offline")

	require.NoError(t, fixture.conn.WriteJSON(map[string]any{
		"type":          "ready",
		"capabilities":  []string{"workspace"},
		"runtimes":      []string{"codex"},
		"runningAgents": []string{},
		"hostname":      "test-host",
		"os":            "linux",
		"daemonVersion": "0.9.0",
	}))
	require.NoError(t, fixture.conn.Close())

	var err error
	require.Eventually(t, func() bool {
		err = fixture.hub.NotifyAgentMessage(fixture.ctx, message.AgentRecipient{
			AgentID:         fixture.agentID,
			MachineID:       fixture.machineID,
			AgentMemberID:   "mem_offline_default",
			AutoWakeEnabled: false,
			DeliveryReason:  message.AgentDeliveryReasonSpaceDefault,
		}, 79, message.ReceiveMessage{
			MessageID:   "msg_interrupt_offline_79",
			ChannelType: "channel",
			ChannelName: "ops",
			AuthorProfile: message.AuthorProfile{
				ActorType:   message.AuthorActorTypeHuman,
				TargetID:    "owner",
				DisplayName: "owner",
			},
			Content:         "urgent replay while daemon websocket is offline",
			Timestamp:       time.Now().UTC(),
			InterruptPolicy: message.AgentInterruptPolicyCurrentRun,
		})
		return errors.Is(err, runtimecontrol.ErrMachineOffline)
	}, 2*time.Second, 10*time.Millisecond)
	require.ErrorIs(t, err, runtimecontrol.ErrMachineOffline)
}

func sendMentionDeliveryForFixture(t *testing.T, fixture autoWakeHubFixture, content string) message.MessageDelivery {
	t.Helper()

	_, delivery, err := fixture.dataStore.App().SendMessageForUser(
		fixture.ctx,
		identity.User{ID: fixture.userID},
		fixture.workspaceID,
		fixture.channelID,
		content,
		nil,
	)
	require.NoError(t, err)
	return delivery
}