package agent

import (
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

func newWakeQueueTestProcess(store DeliveryOwnershipStore) *Process {
	return &Process{
		AgentID:   "agent-wake-queue",
		SessionID: "session-wake-queue",
		processRuntimeBinding: processRuntimeBinding{
			driver: &busyDeliveryStubDriver{},
			cancel: func() {},
		},
		processLifecycleState: processLifecycleState{
			state:    processStateRunning,
			exitCode: -1,
			exited:   make(chan struct{}),
		},
		processDeliveryState: processDeliveryState{deliveryOwnershipStore: store},
		processRecoveryState: processRecoveryState{disableRecovery: true},
	}
}

func acknowledgeCorrelatedDelivery(t *testing.T, ap *Process, store DeliveryOwnershipStore, msg types.InboxMessage, token string) {
	t.Helper()
	pending := pendingRuntimeInput{
		Message:        msg,
		Input:          testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: msg.MessageID}, msg.Content),
		SessionID:      ap.SessionID,
		WakeToken:      token,
		CreatedAt:      time.Now().UTC(),
		InflightReason: DeliveryInflightDispatching,
	}
	_, err := store.PutDeliveryOwnership(ap.ownershipRecord(pending, deliveryOwnershipDispatching, 0, nil))
	require.NoError(t, err)
	ap.rememberPendingRuntimeInput(pending)
	ap.handleParsedEvent(drivercontract.ParsedEvent{
		Kind:           consts.ParsedEventKindRuntimeInputAcked,
		RuntimeInputID: msg.MessageID,
		SessionID:      ap.SessionID,
	})
}

func TestPendingDeliveriesDeduplicateAndBatchSettleResolvedTarget(t *testing.T) {
	store := &memoryDeliveryOwnershipStore{}
	ap := newWakeQueueTestProcess(store)
	var events []WakeStatusEvent
	ap.onWakeStatus = func(event WakeStatusEvent) { events = append(events, event) }

	for _, input := range []struct {
		id  string
		seq int64
	}{
		{"q1", 1}, {"q2", 2}, {"q3", 3},
	} {
		acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
			MessageID: input.id, Target: "#ops:thread", TargetKey: "target-ops", TargetSeq: input.seq,
			Content: "inbound",
		}, "token-"+input.id)
	}
	// A duplicate runtime acknowledgement is not a second obligation.
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID: "q2", Target: "#ops:thread", TargetKey: "target-ops", TargetSeq: 2,
	}, "token-q2-duplicate")

	require.True(t, ap.ConfirmVisibleReplyForTarget("target-ops", 3, "reply-1"))
	ap.mu.Lock()
	require.Nil(t, ap.pendingWake)
	require.Empty(t, ap.pendingDeliveryQueue)
	ap.mu.Unlock()

	var sent []WakeStatusEvent
	for _, event := range events {
		if event.Status == consts.WakeStatusReplySent {
			sent = append(sent, event)
		}
	}
	require.Equal(t, []string{"q1", "q2", "q3"}, []string{sent[0].MessageID, sent[1].MessageID, sent[2].MessageID})
	require.Len(t, sent, 3)
	require.Equal(t, "reply-1", sent[0].ReplyMessageID)
}

func TestPendingDeliveriesKeepTargetsIndependent(t *testing.T) {
	store := &memoryDeliveryOwnershipStore{}
	ap := newWakeQueueTestProcess(store)
	var events []WakeStatusEvent
	ap.onWakeStatus = func(event WakeStatusEvent) { events = append(events, event) }
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID: "ops-1", Target: "#ops", TargetKey: "ops", TargetSeq: 1,
	}, "token-ops")
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID: "dm-1", Target: "dm:@owner", TargetKey: "dm", TargetSeq: 1,
	}, "token-dm")

	require.True(t, ap.ConfirmVisibleReplyForTarget("ops", 1, "reply-ops"))
	ap.mu.Lock()
	require.NotNil(t, ap.pendingWake)
	require.Equal(t, "dm-1", ap.pendingWake.Wake.MessageID)
	ap.mu.Unlock()
	require.False(t, ap.ConfirmVisibleReplyForTarget("unknown", 10, "reply-unknown"))
	require.True(t, ap.ConfirmVisibleReplyForTarget("dm", 1, "reply-dm"))

	var sent []WakeStatusEvent
	for _, event := range events {
		if event.Status == consts.WakeStatusReplySent {
			sent = append(sent, event)
		}
	}
	require.Equal(t, []string{"ops-1", "dm-1"}, []string{sent[0].MessageID, sent[1].MessageID})
}

func TestServerFrontierLeavesHigherSequenceForFreshCheck(t *testing.T) {
	store := &memoryDeliveryOwnershipStore{}
	ap := newWakeQueueTestProcess(store)
	var events []WakeStatusEvent
	ap.onWakeStatus = func(event WakeStatusEvent) { events = append(events, event) }
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID: "old", Target: "#ops:thread", TargetKey: "ops", TargetSeq: 1,
	}, "token-old")
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID: "new", Target: "#ops:thread", TargetKey: "ops", TargetSeq: 3,
	}, "token-new")

	// A server frontier that covers only the older message cannot settle the
	// newer obligation and must make it eligible for a subsequent check.
	require.True(t, ap.ConfirmVisibleReplyForTarget("ops", 1, "reply-old"))
	ap.mu.Lock()
	require.NotNil(t, ap.pendingWake)
	require.Equal(t, "new", ap.pendingWake.Wake.MessageID)
	require.Equal(t, pendingDeliveryAwaitingCheck, ap.pendingWake.Status)
	ap.mu.Unlock()

	stdin := &recordingWriteCloser{}
	ap.stdinPipe = stdin
	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	require.Contains(t, stdin.String(), `exact target "#ops:thread"`)
	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	ap.mu.Lock()
	require.Nil(t, ap.pendingWake)
	ap.mu.Unlock()
	var completed []WakeStatusEvent
	for _, event := range events {
		if event.Status == consts.WakeStatusCompletedWithoutVisibleReply {
			completed = append(completed, event)
		}
	}
	require.Len(t, completed, 1)
	require.Equal(t, "new", completed[0].MessageID)
}

func TestContextRecoveryVisibleReplySettlementKeepsTargetFrontiersIndependent(t *testing.T) {
	store := &memoryDeliveryOwnershipStore{}
	ap := newWakeQueueTestProcess(store)
	var events []WakeStatusEvent
	ap.onWakeStatus = func(event WakeStatusEvent) { events = append(events, event) }

	for _, msg := range []types.InboxMessage{
		{MessageID: "ops-1", Target: "#ops", TargetKey: "ops", TargetSeq: 1, MessageType: "context_recovery"},
		{MessageID: "ops-3", Target: "#ops", TargetKey: "ops", TargetSeq: 3, MessageType: "context_recovery"},
		{MessageID: "dm-2", Target: "dm:@alice", TargetKey: "dm", TargetSeq: 2, MessageType: "context_recovery"},
	} {
		acknowledgeCorrelatedDelivery(t, ap, store, msg, "token-"+msg.MessageID)
	}

	require.True(t, ap.ConfirmVisibleReplyForTarget("ops", 1, "reply-ops-old"))
	ap.RememberFreshnessHeld("ops", 1, 3)
	require.False(t, ap.ConfirmVisibleReplyForTarget("ops", 2, "reply-ops-held"), "held/older frontier must not settle the higher recovery obligation")
	require.True(t, ap.ConfirmVisibleReplyForTarget("dm", 2, "reply-dm"))

	ap.mu.Lock()
	remaining := ap.pendingDeliveriesLocked()
	ap.mu.Unlock()
	require.Len(t, remaining, 1)
	require.Equal(t, "ops-3", remaining[0].Wake.MessageID)
	require.Equal(t, int64(3), remaining[0].Wake.TargetSeq)

	require.True(t, ap.ConfirmVisibleReplyForTarget("ops", 3, "reply-ops-new"))
	var sent []WakeStatusEvent
	for _, event := range events {
		if event.Status == consts.WakeStatusReplySent {
			sent = append(sent, event)
		}
	}
	require.Equal(t, []string{"ops-1", "dm-2", "ops-3"}, []string{sent[0].MessageID, sent[1].MessageID, sent[2].MessageID})
}

func TestSignalOnlyDeliveryDoesNotTriggerOrdinaryFinalCheck(t *testing.T) {
	for _, test := range []struct {
		name string
		msg  types.InboxMessage
	}{
		{name: "reaction", msg: types.InboxMessage{MessageType: "reaction"}},
		{name: "system", msg: types.InboxMessage{AuthorProfile: types.AuthorProfile{ActorType: "system"}}},
		{name: "workflow callback", msg: types.InboxMessage{WorkflowContext: types.WorkflowContext{
			RunID: "run", EffectID: "effect", NodeID: "agent", CallbackExpected: true,
		}}},
	} {
		t.Run(test.name, func(t *testing.T) {
			store := &memoryDeliveryOwnershipStore{}
			ap := newWakeQueueTestProcess(store)
			stdin := &recordingWriteCloser{}
			ap.stdinPipe = stdin
			msg := test.msg
			msg.MessageID = "signal-1"
			msg.Target = "#ops:signal"
			msg.TargetKey = "ops"
			msg.TargetSeq = 1
			acknowledgeCorrelatedDelivery(t, ap, store, msg, "token-signal")

			var events []WakeStatusEvent
			ap.onWakeStatus = func(event WakeStatusEvent) { events = append(events, event) }
			ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
			require.NotContains(t, stdin.String(), "Final delivery check")
			ap.mu.Lock()
			require.Nil(t, ap.pendingWake)
			ap.mu.Unlock()
			var completed bool
			for _, event := range events {
				completed = completed || event.Status == consts.WakeStatusCompletedWithoutVisibleReply
			}
			require.True(t, completed)
		})
	}
}

func TestContextRecoveryDeliveryUsesExactTargetFinalCheck(t *testing.T) {
	store := &memoryDeliveryOwnershipStore{}
	ap := newWakeQueueTestProcess(store)
	stdin := &recordingWriteCloser{}
	ap.stdinPipe = stdin
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID:   "ctx-recovery-1",
		Target:      "#ops:deadbeef",
		TargetKey:   "ops",
		TargetSeq:   42,
		MessageType: "context_recovery",
	}, "token-context-recovery")

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	require.Contains(t, stdin.String(), `exact target "#ops:deadbeef"`)
	require.Contains(t, stdin.String(), "Final delivery check")
}

func TestFinalCheckThenNoReplyTerminatesWithoutLoop(t *testing.T) {
	store := &memoryDeliveryOwnershipStore{}
	ap := newWakeQueueTestProcess(store)
	stdin := &recordingWriteCloser{}
	ap.stdinPipe = stdin
	var events []WakeStatusEvent
	ap.onWakeStatus = func(event WakeStatusEvent) { events = append(events, event) }
	acknowledgeCorrelatedDelivery(t, ap, store, types.InboxMessage{
		MessageID: "q1", Target: "#ops:q1", TargetKey: "ops", TargetSeq: 1,
	}, "token-q1")

	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	require.Equal(t, 1, strings.Count(stdin.String(), "Final delivery check"))
	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: consts.ParsedEventKindTurnEnd})
	require.Equal(t, 1, strings.Count(stdin.String(), "Final delivery check"))
	ap.mu.Lock()
	require.Nil(t, ap.pendingWake)
	ap.mu.Unlock()
	var completed int
	for _, event := range events {
		if event.Status == consts.WakeStatusCompletedWithoutVisibleReply {
			completed++
			require.Equal(t, "reason=no_reply", event.Detail)
		}
	}
	require.Equal(t, 1, completed)
}

func TestFinalCheckAggregatesResolvedTargets(t *testing.T) {
	ap := newWakeQueueTestProcess(&memoryDeliveryOwnershipStore{})
	text := ap.finalVisibleReplyCheckPromptForDeliveries([]pendingDelivery{
		{Wake: WakeInfo{MessageID: "a1", Target: "#ops:thread", TargetKey: "thread", TargetSeq: 1, DeliveryClass: DeliveryClassVisibleReply}},
		{Wake: WakeInfo{MessageID: "a2", Target: "#ops:thread", TargetKey: "thread", TargetSeq: 2, DeliveryClass: DeliveryClassVisibleReply}},
		{Wake: WakeInfo{MessageID: "b1", Target: "dm:@owner", TargetKey: "dm", TargetSeq: 1, DeliveryClass: DeliveryClassVisibleReply}},
		{Wake: WakeInfo{MessageID: "signal", Target: "#signal", DeliveryClass: DeliveryClassSignalOnly}},
	})

	require.Equal(t, 1, strings.Count(text, `- exact target "#ops:thread"`))
	require.Equal(t, 1, strings.Count(text, `- exact target "dm:@owner"`))
	require.NotContains(t, text, `exact target "#signal"`)
	require.NotContains(t, text, "wake message")
}