package agent

import (
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

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

func reclaimedLease(id string, reclaimedAt time.Time) pendingRuntimeInput {
	return pendingRuntimeInput{
		Message:     types.InboxMessage{MessageID: id, TargetKey: "#ops", TargetSeq: 1},
		Input:       drivercontract.RuntimeInput{ID: id},
		reclaimed:   true,
		reclaimedAt: reclaimedAt,
	}
}

// TestSweepReclaimedCorrelationLeasesEvictsExpiredByTTL verifies that a
// reclaimed correlation lease past reclaimedCorrelationLeaseTTL is dropped while
// in-window leases and live (non-reclaimed) entries are retained.
func TestSweepReclaimedCorrelationLeasesEvictsExpiredByTTL(t *testing.T) {
	base := time.Now().UTC()
	ap := &Process{}
	ap.mu.Lock()
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"recent":  reclaimedLease("recent", base.Add(-time.Minute)),
		"expired": reclaimedLease("expired", base.Add(-(reclaimedCorrelationLeaseTTL + time.Second))),
		"live":    {Message: types.InboxMessage{MessageID: "live"}, Input: drivercontract.RuntimeInput{ID: "live"}, reclaimedAt: base.Add(-time.Hour)},
	}
	ap.sweepReclaimedCorrelationLeasesLocked(base)
	require.Contains(t, ap.pendingRuntimeInputs, "recent")
	require.NotContains(t, ap.pendingRuntimeInputs, "expired")
	require.Contains(t, ap.pendingRuntimeInputs, "live")
	ap.mu.Unlock()
}

// TestSweepReclaimedCorrelationLeasesEvictsOldestOverCap verifies the hard
// capacity bound: when reclaimed leases exceed maxReclaimedCorrelationLeases the
// oldest (by reclaimedAt) are evicted first, keeping the newest correlation
// hints most likely to still receive a late acked ACK.
func TestSweepReclaimedCorrelationLeasesEvictsOldestOverCap(t *testing.T) {
	base := time.Now().UTC()
	ap := &Process{}
	ap.mu.Lock()
	total := maxReclaimedCorrelationLeases + 2
	ap.pendingRuntimeInputs = make(map[string]pendingRuntimeInput, total)
	for i := 0; i < total; i++ {
		id := fmt.Sprintf("lease-%03d", i)
		ap.pendingRuntimeInputs[id] = reclaimedLease(id, base.Add(time.Duration(i)*time.Millisecond))
	}
	ap.sweepReclaimedCorrelationLeasesLocked(base)
	require.Len(t, ap.pendingRuntimeInputs, maxReclaimedCorrelationLeases)
	require.NotContains(t, ap.pendingRuntimeInputs, "lease-000")
	require.NotContains(t, ap.pendingRuntimeInputs, "lease-001")
	require.Contains(t, ap.pendingRuntimeInputs, fmt.Sprintf("lease-%03d", total-1))
	ap.mu.Unlock()
}

// TestReclaimedCorrelationLeaseEvictedByProductionDeadlineWithoutSubsequentReclaim
// verifies B1: a reclaimed correlation lease is evicted on its TTL deadline by
// the production sweep timer, with no second applyReclaimedDeliveryOwnership
// call to drive the sweep. Without the timer the lease would outlive its TTL
// until the next reclaim admission (or indefinitely if none arrives).
func TestReclaimedCorrelationLeaseEvictedByProductionDeadlineWithoutSubsequentReclaim(t *testing.T) {
	createdAt := time.Now().UTC()
	old := DeliveryOwnershipRecord{
		Message:        types.InboxMessage{MessageID: "old", Target: "#ops"},
		Input:          drivercontract.RuntimeInput{ID: "old"},
		State:          DeliveryOwnershipInflight,
		InflightReason: DeliveryInflightOutcomeUnknown,
		AcceptedAt:     createdAt,
	}
	ap := newInboxRetryTestProcess(&ackedRuntimeInputStubDriver{supports: true}, time.Hour)
	ap.reclaimedCorrelationLeaseTTL = 100 * time.Millisecond
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"old": {
			Message:        old.Message,
			Input:          old.Input,
			SessionID:      ap.SessionID,
			CreatedAt:      createdAt,
			InflightReason: DeliveryInflightOutcomeUnknown,
		},
	}
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	ap.applyReclaimedDeliveryOwnership([]DeliveryOwnershipRecord{{
		Message:     old.Message,
		Input:       old.Input,
		State:       DeliveryOwnershipTerminal,
		Outcome:     DeliveryOutcomeFailed,
		FailureCode: "ownership_backlog_reclaimed",
		AcceptedAt:  createdAt,
	}})

	ap.mu.Lock()
	_, hasLease := ap.pendingRuntimeInputs["old"]
	ap.mu.Unlock()
	require.True(t, hasLease, "reclaimed correlation lease is retained immediately after projection")

	require.Eventually(t, func() bool {
		ap.mu.Lock()
		defer ap.mu.Unlock()
		_, present := ap.pendingRuntimeInputs["old"]
		return !present
	}, 2*time.Second, 5*time.Millisecond, "reclaimed correlation lease must be evicted by the production deadline sweep without a subsequent reclaim admission")
}