namespace NATS.Server.JetStream.Consumers;

public sealed class AckProcessor
{
    private readonly Dictionary<ulong, PendingState> _pending = new();
    public ulong AckFloor { get; private set; }

    public void Register(ulong sequence, int ackWaitMs)
    {
        if (sequence <= AckFloor)
            return;

        if (_pending.ContainsKey(sequence))
            return;

        _pending[sequence] = new PendingState
        {
            DeadlineUtc = DateTime.UtcNow.AddMilliseconds(Math.Max(ackWaitMs, 1)),
            Deliveries = 1,
        };
    }

    public bool TryGetExpired(out ulong sequence, out int deliveries)
    {
        foreach (var (seq, state) in _pending)
        {
            if (DateTime.UtcNow >= state.DeadlineUtc)
            {
                sequence = seq;
                deliveries = state.Deliveries;
                return true;
            }
        }

        sequence = 0;
        deliveries = 0;
        return false;
    }

    public void ScheduleRedelivery(ulong sequence, int delayMs)
    {
        if (!_pending.TryGetValue(sequence, out var state))
            return;

        state.Deliveries++;
        state.DeadlineUtc = DateTime.UtcNow.AddMilliseconds(Math.Max(delayMs, 1));
        _pending[sequence] = state;
    }

    public void Drop(ulong sequence)
    {
        _pending.Remove(sequence);
    }

    public bool HasPending => _pending.Count > 0;
    public int PendingCount => _pending.Count;

    public void AckAll(ulong sequence)
    {
        foreach (var key in _pending.Keys.Where(k => k <= sequence).ToArray())
            _pending.Remove(key);

        if (sequence > AckFloor)
            AckFloor = sequence;
    }

    private sealed class PendingState
    {
        public DateTime DeadlineUtc { get; set; }
        public int Deliveries { get; set; }
    }
}