package agent

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"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"
)

type inboxRetryAckedDriver struct {
	busyDeliveryStubDriver

	mu          sync.Mutex
	failures    map[string]int
	calls       []string
	callTimes   []time.Time
	active      int
	maxActive   int
	block       <-chan struct{}
	started     chan string
	cancelCount int
	failureText string
}

func (d *inboxRetryAckedDriver) SupportsAckedBusyRuntimeInputDelivery() bool { return true }

func (d *inboxRetryAckedDriver) DeliverAckedBusyRuntimeInput(ctx context.Context, input drivercontract.RuntimeInput, _ string) error {
	d.mu.Lock()
	d.calls = append(d.calls, input.ID)
	d.callTimes = append(d.callTimes, time.Now())
	d.active++
	if d.active > d.maxActive {
		d.maxActive = d.active
	}
	d.mu.Unlock()

	defer func() {
		d.mu.Lock()
		d.active--
		d.mu.Unlock()
	}()

	if d.started != nil {
		d.started <- input.ID
	}
	if d.block != nil {
		select {
		case <-d.block:
		case <-ctx.Done():
			d.mu.Lock()
			d.cancelCount++
			d.mu.Unlock()
			return ctx.Err()
		}
	}

	d.mu.Lock()
	remaining := d.failures[input.ID]
	if remaining > 0 {
		d.failures[input.ID] = remaining - 1
	}
	d.mu.Unlock()
	if remaining > 0 {
		if d.failureText != "" {
			return errors.New(d.failureText)
		}
		return errors.New("temporary runtime input rejection")
	}
	return nil
}

func (d *inboxRetryAckedDriver) snapshot() (calls []string, callTimes []time.Time, maxActive int, active int, cancelCount int) {
	d.mu.Lock()
	defer d.mu.Unlock()
	return append([]string(nil), d.calls...), append([]time.Time(nil), d.callTimes...), d.maxActive, d.active, d.cancelCount
}

func newInboxRetryTestProcess(driver drivercontract.Driver, delay time.Duration) *Process {
	return &Process{
		AgentID:   "agent-inbox-retry",
		SessionID: "session-inbox-retry",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
			cancel: func() {},
		},
		processLifecycleState: processLifecycleState{
			state:    processStateRunning,
			isIdle:   false,
			exitCode: -1,
		},
		processDeliveryState: processDeliveryState{
			inboxRetryDelay: func(int) time.Duration { return delay },
		},
		processRecoveryState: processRecoveryState{disableRecovery: true},
	}
}

func stopInboxRetryTestProcess(ap *Process) {
	ap.mu.Lock()
	ap.stopping = true
	ap.state = processStateStopped
	ap.stopInboxRetryLocked()
	ap.cancelReclaimedCorrelationLeaseSweepLocked()
	ap.mu.Unlock()
}

func TestInboxRetryTimerRecoversWithoutRuntimeBoundary(t *testing.T) {
	driver := &inboxRetryAckedDriver{failures: map[string]int{"q1": 2}}
	ap := newInboxRetryTestProcess(driver, 10*time.Millisecond)
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	require.NoError(t, ap.deliverMessage(types.InboxMessage{
		MessageID: "q1",
		TargetKey: "scope",
		TargetSeq: 1,
		Content:   "original q1 body",
	}, false))

	require.Eventually(t, func() bool {
		ap.mu.Lock()
		_, pending := ap.pendingRuntimeInputs["q1"]
		queued := len(ap.inbox)
		ap.mu.Unlock()
		calls, _, _, _, _ := driver.snapshot()
		return pending && queued == 0 && len(calls) == 3
	}, time.Second, 5*time.Millisecond, "timer must retry without SessionInit, TurnEnd, or RuntimeIdle")

	calls, callTimes, maxActive, _, _ := driver.snapshot()
	require.Equal(t, []string{"q1", "q1", "q1"}, calls)
	require.Equal(t, 1, maxActive)
	require.Len(t, callTimes, 3)
	require.GreaterOrEqual(t, callTimes[1].Sub(callTimes[0]), 5*time.Millisecond)
	require.GreaterOrEqual(t, callTimes[2].Sub(callTimes[1]), 5*time.Millisecond)

	// Transport acceptance is not a acked acknowledgement. The correlated
	// owner must remain pending until runtime_input_acked arrives.
	require.Nil(t, ap.ReconcileRuntimeAckedMessages([]string{"q1"}))
}

func TestInboxRetryWaitsForRuntimeBoundaryWhenBusyTransportCannotAccept(t *testing.T) {
	driver := &protocolDeliveryStubDriver{runtimeName: "codex", supportsBusy: false}
	ap := newInboxRetryTestProcess(driver, time.Millisecond)
	ap.inbox = []types.InboxMessage{{MessageID: "q2", Content: "wait behind active q1"}}
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"q1": {
			Message:   types.InboxMessage{MessageID: "q1", Content: "active"},
			Input:     testRuntimeInputWithConfig(drivercontract.RuntimeInput{ID: "q1"}, "active"),
			CreatedAt: time.Now(),
		},
	}
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	ap.scheduleInboxRetry(true)
	time.Sleep(20 * time.Millisecond)
	ap.mu.Lock()
	require.Nil(t, ap.inboxRetryTimer)
	require.False(t, ap.inboxRetryRunning)
	require.Len(t, ap.inbox, 1)
	require.Equal(t, "q2", ap.inbox[0].MessageID)
	ap.mu.Unlock()

	// A later authoritative idle boundary makes the same durable head
	// deliverable and explicitly kicks the worker.
	ap.mu.Lock()
	ap.isIdle = true
	ap.mu.Unlock()
	ap.scheduleInboxRetry(true)
	require.Eventually(t, func() bool {
		ap.mu.Lock()
		defer ap.mu.Unlock()
		_, acked := ap.ackedSet["q2"]
		return acked && len(ap.inbox) == 0
	}, time.Second, 5*time.Millisecond)
}

func TestInboxRetryTerminalFailurePromotesSafetyOnlyAfterOwnerIsQuiescent(t *testing.T) {
	driver := &protocolRecoveryStubDriver{deliverErr: errors.New("permanent direct failure")}
	ap := &Process{
		AgentID:   "agent-terminal-safety",
		SessionID: "thread-1",
		processRuntimeBinding: processRuntimeBinding{
			driver: driver,
			cancel: func() {},
		},
		processLifecycleState: processLifecycleState{
			state:    processStateRunning,
			isIdle:   true,
			exitCode: -1,
		},
		processDeliveryState: processDeliveryState{
			inboxRetryMaxAttempts:  1,
			inboxRetryDelay:        func(int) time.Duration { return time.Millisecond },
			deliveryOwnershipStore: &memoryDeliveryOwnershipStore{},
		},
		processRecoveryState: processRecoveryState{},
	}
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	var safetyMu sync.Mutex
	var safety []bool
	require.NoError(t, ap.SetOnResumeSafety(func(_ uint64, safe bool, _ string) error {
		safetyMu.Lock()
		safety = append(safety, safe)
		safetyMu.Unlock()
		return nil
	}))
	require.NoError(t, ap.markResumeSafety(true))

	require.NoError(t, ap.deliverMessage(types.InboxMessage{
		MessageID: "q1",
		Target:    "#ops",
		Content:   "poison",
	}, false))
	require.Eventually(t, func() bool {
		ap.mu.Lock()
		_, failed := ap.failedSet["q1"]
		inboxEmpty := len(ap.inbox) == 0
		resumeSafe := ap.resumeSafety
		ap.mu.Unlock()
		return failed && inboxEmpty && resumeSafe
	}, time.Second, 5*time.Millisecond)

	safetyMu.Lock()
	got := append([]bool(nil), safety...)
	safetyMu.Unlock()
	require.GreaterOrEqual(t, len(got), 3)
	require.True(t, got[0])
	require.True(t, got[len(got)-1],
		"safety may return only after the failed tombstone is durable and the local owner is removed")
	for _, safe := range got[1 : len(got)-1] {
		require.False(t, safe, "no retry may publish a transient safe boundary while the owner is live")
	}
}

func TestInboxRetryTerminationPolicyDistinguishesRuntimeBoundary(t *testing.T) {
	ap := newInboxRetryTestProcess(&inboxRetryAckedDriver{}, time.Millisecond)
	ap.inboxRetryMessageID = "q1"
	ap.inboxRetryLastError = "blocked"
	ap.inboxRetryAttempt = 100
	ap.inboxRetryMaxAttempts = 1
	ap.inboxRetryStartedAt = time.Now()

	ap.inboxRetryLastClass = deliveryRetryPermanent
	require.True(t, ap.shouldTerminateInboxRetryHeadLocked("q1"), "permanent failures terminate immediately")

	ap.inboxRetryLastClass = deliveryRetryTransient
	require.True(t, ap.shouldTerminateInboxRetryHeadLocked("q1"), "transient failures use the short attempt budget")

	ap.inboxRetryLastClass = deliveryRetryBoundary
	ap.inboxRetryBoundaryMaxElapsed = time.Hour
	require.False(t, ap.shouldTerminateInboxRetryHeadLocked("q1"), "review/approval/compaction boundaries ignore the short attempt budget")
	ap.inboxRetryBoundaryMaxElapsed = time.Millisecond
	ap.inboxRetryStartedAt = time.Now().Add(-time.Second)
	require.True(t, ap.shouldTerminateInboxRetryHeadLocked("q1"), "a boundary still has a long elapsed terminal bound")
}

func TestClassifyDeliveryRetryErrorRecognizesRuntimeBoundaries(t *testing.T) {
	require.Equal(t, deliveryRetryBoundary, classifyDeliveryRetryError(errors.New("cannot accept input during review turn")))
	require.Equal(t, deliveryRetryBoundary, classifyDeliveryRetryError(errors.New("cannot accept input during compaction")))
	require.Equal(t, deliveryRetryBoundary, classifyDeliveryRetryError(errors.New("runtime input mutation was not attempted: resolve codex delivery target: Codex turn dispatch crossed a lifecycle boundary: codex thread thread-1 is active but its active turn id is unavailable after authoritative read")))
	require.Equal(t, deliveryRetryPermanent, classifyDeliveryRetryError(errors.New("delivery produced empty input")))
	require.Equal(t, deliveryRetryTransient, classifyDeliveryRetryError(errors.New("temporary runtime input rejection")))
	require.Equal(t, deliveryRetryBoundary, classifyDeliveryRetryError(&appServerRequestError{
		Method:  "turn/steer",
		Message: "request rejected",
		Data: json.RawMessage(`{
			"codexErrorInfo":{"activeTurnNotSteerable":{"turnKind":"review"}}
		}`),
	}), "structured app-server data must win over its generic message")
}

func TestInboxRetryUsesSingleWorkerWhenKickedConcurrently(t *testing.T) {
	release := make(chan struct{})
	driver := &inboxRetryAckedDriver{
		failures: map[string]int{},
		block:    release,
		started:  make(chan string, 4),
	}
	ap := newInboxRetryTestProcess(driver, time.Millisecond)
	ap.inbox = []types.InboxMessage{{MessageID: "q1", Content: "first"}}
	t.Cleanup(func() {
		select {
		case <-release:
		default:
			close(release)
		}
		stopInboxRetryTestProcess(ap)
	})

	ap.scheduleInboxRetry(true)
	require.Equal(t, "q1", <-driver.started)

	var kicks sync.WaitGroup
	for i := 0; i < 20; i++ {
		kicks.Add(1)
		go func() {
			defer kicks.Done()
			ap.scheduleInboxRetry(true)
		}()
	}
	kicks.Wait()
	time.Sleep(10 * time.Millisecond)
	calls, _, maxActive, active, _ := driver.snapshot()
	require.Equal(t, []string{"q1"}, calls)
	require.Equal(t, 1, maxActive)
	require.Equal(t, 1, active)

	close(release)
	require.Eventually(t, func() bool {
		_, _, _, active, _ := driver.snapshot()
		ap.mu.Lock()
		_, pending := ap.pendingRuntimeInputs["q1"]
		ap.mu.Unlock()
		return active == 0 && pending
	}, time.Second, 5*time.Millisecond)
	calls, _, maxActive, _, _ = driver.snapshot()
	require.Equal(t, []string{"q1"}, calls)
	require.Equal(t, 1, maxActive)
}

func TestInboxRetryStopCancelsTimerAndInflightWorker(t *testing.T) {
	t.Run("timer", func(t *testing.T) {
		driver := &inboxRetryAckedDriver{failures: map[string]int{}}
		ap := newInboxRetryTestProcess(driver, time.Hour)
		ap.inbox = []types.InboxMessage{{MessageID: "q1", Content: "first"}}
		ap.scheduleInboxRetry(false)

		ap.mu.Lock()
		generation := ap.inboxRetryGeneration
		require.NotNil(t, ap.inboxRetryTimer)
		ap.mu.Unlock()
		ap.Stop()

		time.Sleep(20 * time.Millisecond)
		calls, _, _, _, _ := driver.snapshot()
		require.Empty(t, calls)
		ap.mu.Lock()
		require.Nil(t, ap.inboxRetryTimer)
		require.Greater(t, ap.inboxRetryGeneration, generation)
		require.Len(t, ap.inbox, 1)
		ap.mu.Unlock()
	})

	t.Run("inflight", func(t *testing.T) {
		release := make(chan struct{})
		driver := &inboxRetryAckedDriver{
			failures: map[string]int{},
			block:    release,
			started:  make(chan string, 1),
		}
		ap := newInboxRetryTestProcess(driver, time.Millisecond)
		ap.inbox = []types.InboxMessage{{MessageID: "q1", Content: "first"}}
		ap.scheduleInboxRetry(true)
		require.Equal(t, "q1", <-driver.started)

		ap.Stop()
		require.Eventually(t, func() bool {
			_, _, _, active, cancelCount := driver.snapshot()
			return active == 0 && cancelCount == 1
		}, time.Second, 5*time.Millisecond)
		time.Sleep(10 * time.Millisecond)

		calls, _, maxActive, _, _ := driver.snapshot()
		require.Equal(t, []string{"q1"}, calls)
		require.Equal(t, 1, maxActive)
		ap.mu.Lock()
		require.Nil(t, ap.inboxRetryTimer)
		require.False(t, ap.inboxRetryRunning)
		require.Len(t, ap.inbox, 1, "canceled in-flight delivery must return to daemon ownership")
		require.Equal(t, "q1", ap.inbox[0].MessageID)
		ap.mu.Unlock()
	})
}

func TestInboxRetryRejectsStaleTimerGeneration(t *testing.T) {
	driver := &inboxRetryAckedDriver{failures: map[string]int{}}
	ap := newInboxRetryTestProcess(driver, time.Hour)
	ap.inbox = []types.InboxMessage{{MessageID: "q1", Content: "first"}}
	ap.inboxRetryGeneration = 2

	ap.runInboxRetry(1)

	calls, _, _, _, _ := driver.snapshot()
	require.Empty(t, calls)
	ap.mu.Lock()
	require.Len(t, ap.inbox, 1)
	require.Equal(t, uint64(2), ap.inboxRetryGeneration)
	ap.mu.Unlock()
}

func TestInboxRetryDoesNotReclaimAcceptedPendingInputAtRuntimeBoundary(t *testing.T) {
	driver := &inboxRetryAckedDriver{failures: map[string]int{}}
	ap := newInboxRetryTestProcess(driver, time.Millisecond)
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	require.NoError(t, ap.deliverMessage(types.InboxMessage{
		MessageID: "q1",
		Content:   "accepted but not acknowledged",
	}, false))

	for _, kind := range []string{
		consts.ParsedEventKindTurnEnd,
		consts.ParsedEventKindRuntimeIdle,
	} {
		ap.handleParsedEvent(drivercontract.ParsedEvent{Kind: kind, SessionID: "session-inbox-retry"})
	}
	time.Sleep(20 * time.Millisecond)

	calls, _, _, _, _ := driver.snapshot()
	require.Equal(t, []string{"q1"}, calls, "normal boundaries must not resubmit runtime-owned pending input")
	ap.mu.Lock()
	require.Empty(t, ap.inbox)
	require.Contains(t, ap.pendingRuntimeInputs, "q1")
	ap.mu.Unlock()
	require.Nil(t, ap.ReconcileRuntimeAckedMessages([]string{"q1"}))
}