// Go: consumer.go (trackPending, processNak, rdc map, addToRedeliverQueue ~line 5540)
// RedeliveryTracker manages sequences waiting for redelivery after a NAK or ack-wait
// expiry.  It mirrors the Go consumer's rdc (redelivery count) map combined with the
// rdq (redelivery queue) priority ordering.
namespace NATS.Server.JetStream.Consumers;

public sealed class RedeliveryTracker
{
    private readonly int[] _backoffMs;
    private readonly long[]? _backoffMsLong;

    // Go: consumer.go — pending maps sseq → (deadline, deliveries)
    private readonly Dictionary<ulong, RedeliveryEntry> _entries = new();

    // Go: consumer.go — rdc map tracks per-sequence delivery counts
    private readonly Dictionary<ulong, int> _deliveryCounts = new();

    // Go: consumer.go — rdq priority queue ordered by deadline for efficient dispatch
    private readonly PriorityQueue<ulong, DateTimeOffset> _priorityQueue = new();

    // Stored config for the new constructor overload
    private readonly int _maxDeliveries;
    private readonly long _ackWaitMs;

    // Go: consumer.go:100 — BackOff []time.Duration in ConsumerConfig; empty falls back to ackWait
    public RedeliveryTracker(int[] backoffMs)
    {
        _backoffMs = backoffMs;
        _backoffMsLong = null;
        _maxDeliveries = 0;
        _ackWaitMs = 0;
    }

    // Go: consumer.go — ConsumerConfig maxDeliver + ackWait + backoff, new overload storing config fields
    public RedeliveryTracker(int maxDeliveries, long ackWaitMs, long[]? backoffMs = null)
    {
        _backoffMs = [];
        _backoffMsLong = backoffMs;
        _maxDeliveries = maxDeliveries;
        _ackWaitMs = ackWaitMs;
    }

    // Go: consumer.go:5540 — trackPending records delivery count and schedules deadline
    // using the backoff array indexed by (deliveryCount-1), clamped at last entry.
    // Returns the UTC time at which the sequence next becomes eligible for redelivery.
    public DateTime Schedule(ulong seq, int deliveryCount, int ackWaitMs = 0)
    {
        var delayMs = ResolveDelay(deliveryCount, ackWaitMs);
        var deadline = DateTime.UtcNow.AddMilliseconds(Math.Max(delayMs, 1));

        _entries[seq] = new RedeliveryEntry
        {
            DeadlineUtc = deadline,
            DeliveryCount = deliveryCount,
        };

        return deadline;
    }

    // Go: consumer.go — schedule with an explicit deadline into the priority queue
    public void Schedule(ulong seq, DateTimeOffset deadline)
    {
        _deliveryCounts.TryAdd(seq, 0);
        _priorityQueue.Enqueue(seq, deadline);
    }

    // Go: consumer.go — rdq entries are dispatched once their deadline has passed
    public IReadOnlyList<ulong> GetDue()
    {
        var now = DateTime.UtcNow;
        List<ulong>? due = null;

        foreach (var (seq, entry) in _entries)
        {
            if (entry.DeadlineUtc <= now)
            {
                due ??= [];
                due.Add(seq);
            }
        }

        return due ?? (IReadOnlyList<ulong>)[];
    }

    // Go: consumer.go — drain the rdq priority queue of all entries whose deadline <= now,
    // returning them in deadline order (earliest first).
    public IEnumerable<ulong> GetDue(DateTimeOffset now)
    {
        List<(ulong seq, DateTimeOffset deadline)>? dequeued = null;
        List<(ulong seq, DateTimeOffset deadline)>? future = null;

        // Drain the entire queue, separating due from future
        while (_priorityQueue.TryDequeue(out var seq, out var deadline))
        {
            // Skip sequences that were acknowledged
            if (!_deliveryCounts.ContainsKey(seq))
                continue;

            if (deadline <= now)
            {
                dequeued ??= [];
                dequeued.Add((seq, deadline));
            }
            else
            {
                future ??= [];
                future.Add((seq, deadline));
            }
        }

        // Re-enqueue future items
        if (future is not null)
        {
            foreach (var (seq, deadline) in future)
                _priorityQueue.Enqueue(seq, deadline);
        }

        if (dequeued is null)
            return [];

        // Already extracted in priority order since PriorityQueue dequeues min first
        return dequeued.Select(x => x.seq);
    }

    // Go: consumer.go — acking a sequence removes it from the pending redelivery set
    public void Acknowledge(ulong seq)
    {
        _entries.Remove(seq);
        _deliveryCounts.Remove(seq);
        // Priority queue entries are lazily skipped in GetDue when seq not in _deliveryCounts
    }

    // Go: consumer.go — maxdeliver check: drop sequence once delivery count exceeds max
    public bool IsMaxDeliveries(ulong seq, int maxDeliver)
    {
        if (maxDeliver <= 0)
            return false;

        if (!_entries.TryGetValue(seq, out var entry))
            return false;

        return entry.DeliveryCount >= maxDeliver;
    }

    // Go: consumer.go — maxdeliver check using the stored _maxDeliveries from new constructor
    public bool IsMaxDeliveries(ulong seq)
    {
        if (_maxDeliveries <= 0)
            return false;

        _deliveryCounts.TryGetValue(seq, out var count);
        return count >= _maxDeliveries;
    }

    // Go: consumer.go — rdc map increment: track how many times a sequence has been delivered
    public void IncrementDeliveryCount(ulong seq)
    {
        _deliveryCounts[seq] = _deliveryCounts.TryGetValue(seq, out var count) ? count + 1 : 1;
    }

    // Go: consumer.go — backoff delay lookup: index by deliveryCount, clamp to last entry,
    // fall back to ackWait when no backoff array is configured.
    public long GetBackoffDelay(int deliveryCount)
    {
        if (_backoffMsLong is { Length: > 0 })
        {
            var idx = Math.Min(deliveryCount - 1, _backoffMsLong.Length - 1);
            if (idx < 0) idx = 0;
            return _backoffMsLong[idx];
        }

        return _ackWaitMs;
    }

    public bool IsTracking(ulong seq) => _entries.ContainsKey(seq);

    public int TrackedCount => _entries.Count;

    // Go: consumer.go — backoff index = min(deliveries-1, len(backoff)-1);
    // falls back to ackWaitMs when the backoff array is empty.
    private int ResolveDelay(int deliveryCount, int ackWaitMs)
    {
        if (_backoffMs.Length == 0)
            return Math.Max(ackWaitMs, 1);

        var idx = Math.Min(deliveryCount - 1, _backoffMs.Length - 1);
        if (idx < 0)
            idx = 0;
        return _backoffMs[idx];
    }

    private sealed class RedeliveryEntry
    {
        public DateTime DeadlineUtc { get; set; }
        public int DeliveryCount { get; set; }
    }
}
