// Go: consumer.go (processAckMsg, processNak, processTerm, processAckProgress)
namespace NATS.Server.JetStream.Consumers;

// Go: consumer.go:2550 — ack type prefix constants (+ACK, -NAK, +TERM, +WPI)
public enum AckType { Ack, Nak, Term, Progress, Unknown }

/// <summary>
/// Policy applied when a sequence exceeds its max delivery count.
/// Go reference: consumer.go maxDeliver config.
/// </summary>
public enum DeliveryExceededPolicy
{
    /// <summary>Drop the message — no further action.</summary>
    Drop,
    /// <summary>Move to a dead-letter queue subject.</summary>
    DeadLetter,
}

public sealed class AckProcessor
{
    // Go: consumer.go — ackTerminatedFlag marks sequences that must not be redelivered
    private readonly HashSet<ulong> _terminated = new();
    private readonly Dictionary<ulong, PendingState> _pending = new();
    private readonly int[]? _backoffMs;
    private int _ackWaitMs;

    // Fields for the RedeliveryTracker constructor overload
    private readonly RedeliveryTracker? _tracker;
    private readonly int _maxAckPending;

    // Stores deliver subjects keyed by sequence for tracker-based registrations
    private readonly Dictionary<ulong, string>? _deliverSubjects;

    // Go: consumer.go — maxDeliver config; sequences that exceeded the limit
    private int _maxDeliver;
    private readonly List<ulong> _exceededSequences = new();

    public ulong AckFloor { get; private set; }
    public int TerminatedCount { get; private set; }

    /// <summary>
    /// Sets the max delivery limit. 0 or negative means unlimited.
    /// Go reference: consumer.go maxDeliver config.
    /// </summary>
    public int MaxDeliver
    {
        get => _maxDeliver;
        set => _maxDeliver = Math.Max(value, 0);
    }

    /// <summary>Number of sequences that exceeded max deliveries.</summary>
    public int ExceededCount => _exceededSequences.Count;

    /// <summary>Returns sequences that exceeded max deliveries since last drain.</summary>
    public IReadOnlyList<ulong> GetExceededSequences() => _exceededSequences.AsReadOnly();

    /// <summary>Clears the exceeded sequences list (after advisories are sent).</summary>
    public void DrainExceeded() => _exceededSequences.Clear();

    /// <summary>Policy applied when a sequence exceeds its max delivery count.</summary>
    public DeliveryExceededPolicy ExceededPolicy { get; set; } = DeliveryExceededPolicy.Drop;

    public AckProcessor(int[]? backoffMs = null)
    {
        _backoffMs = backoffMs;
    }

    // Go: consumer.go — ConsumerConfig maxAckPending + RedeliveryTracker integration
    public AckProcessor(RedeliveryTracker tracker, int maxAckPending = 0)
    {
        _tracker = tracker;
        _maxAckPending = maxAckPending;
        _deliverSubjects = new Dictionary<ulong, string>();
        _backoffMs = null;
    }

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

        if (_pending.ContainsKey(sequence))
            return;

        _ackWaitMs = ackWaitMs;

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

    // Go: consumer.go — register with deliver subject; ackWait comes from the tracker
    public void Register(ulong sequence, string deliverSubject)
    {
        if (_tracker is null)
            throw new InvalidOperationException("Register(ulong, string) requires a RedeliveryTracker constructor.");

        var ackWaitMs = (int)Math.Max(_tracker.GetBackoffDelay(1), 1);
        Register(sequence, ackWaitMs);

        if (_deliverSubjects is not null)
            _deliverSubjects[sequence] = deliverSubject;
    }

    // Go: consumer.go — processAck without payload: plain +ACK, also notifies tracker
    public void ProcessAck(ulong seq)
    {
        AckSequence(seq);
        _tracker?.Acknowledge(seq);
    }

    // Go: consumer.go — returns ack deadline for a pending sequence; MinValue if not tracked
    public DateTimeOffset GetDeadline(ulong seq)
    {
        if (_pending.TryGetValue(seq, out var state))
            return new DateTimeOffset(state.DeadlineUtc, TimeSpan.Zero);

        return DateTimeOffset.MinValue;
    }

    // Go: consumer.go — maxAckPending=0 means unlimited; otherwise cap pending registrations
    public bool CanRegister() => _maxAckPending <= 0 || _pending.Count < _maxAckPending;

    // Go: consumer.go:2550 — parse ack type prefix from raw payload bytes
    public static AckType ParseAckType(ReadOnlySpan<byte> data)
    {
        if (data.StartsWith("+ACK"u8))
            return AckType.Ack;
        if (data.StartsWith("-NAK"u8))
            return AckType.Nak;
        if (data.StartsWith("+TERM"u8))
            return AckType.Term;
        if (data.StartsWith("+WPI"u8))
            return AckType.Progress;
        return AckType.Unknown;
    }

    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;
    }

    // Go: consumer.go:2550 (processAck)
    // Dispatches to the appropriate ack handler based on ack type prefix.
    // Empty or "+ACK" → ack single; "-NAK" → schedule redelivery; "+TERM" → terminate; "+WPI" → progress reset.
    public void ProcessAck(ulong seq, ReadOnlySpan<byte> payload)
    {
        if (payload.IsEmpty || payload.SequenceEqual("+ACK"u8))
        {
            AckSequence(seq);
            return;
        }

        if (payload.StartsWith("-NAK"u8))
        {
            // Go: consumer.go — parseNak extracts optional delay from "-NAK {delay}"
            var delayMs = 0;
            var rest = payload["-NAK"u8.Length..];
            if (!rest.IsEmpty && rest[0] == (byte)' ')
            {
                var delaySpan = rest[1..];
                if (TryParseInt(delaySpan, out var parsed))
                    delayMs = parsed;
            }
            ProcessNak(seq, delayMs);
            return;
        }

        if (payload.StartsWith("+TERM"u8))
        {
            ProcessTerm(seq);
            return;
        }

        if (payload.StartsWith("+WPI"u8))
        {
            ProcessProgress(seq);
            return;
        }

        // Unknown ack type — treat as plain ack per Go behavior
        AckSequence(seq);
    }

    // Go: consumer.go — processAck for "+ACK": removes from pending and advances AckFloor when contiguous
    public void AckSequence(ulong seq)
    {
        _pending.Remove(seq);
        _terminated.Remove(seq);

        // Advance floor while the next-in-order sequences are no longer pending
        if (seq == AckFloor + 1)
        {
            AckFloor = seq;
            while (_pending.Count > 0)
            {
                var next = AckFloor + 1;
                if (_pending.ContainsKey(next))
                    break;
                // Only advance if next is definitely below any pending sequence
                // Stop when we hit a gap or run out of sequences to check
                if (!HasSequenceBelow(next))
                    break;
                AckFloor = next;
            }
        }
    }

    // Go: consumer.go — processNak: schedules redelivery with optional explicit delay or backoff array
    public void ProcessNak(ulong seq, int delayMs = 0)
    {
        if (_terminated.Contains(seq))
            return;

        if (!_pending.TryGetValue(seq, out var state))
            return;

        int effectiveDelay;
        if (delayMs > 0)
        {
            effectiveDelay = delayMs;
        }
        else if (_backoffMs is { Length: > 0 })
        {
            // Go: consumer.go — backoff array clamps at last entry for high delivery counts
            var idx = Math.Min(state.Deliveries - 1, _backoffMs.Length - 1);
            effectiveDelay = _backoffMs[idx];
        }
        else
        {
            effectiveDelay = Math.Max(_ackWaitMs, 1);
        }

        ScheduleRedelivery(seq, effectiveDelay);
    }

    // Go: consumer.go — processTerm: removes from pending permanently; sequence is never redelivered
    public void ProcessTerm(ulong seq)
    {
        if (_pending.Remove(seq))
        {
            _terminated.Add(seq);
            TerminatedCount++;
        }
    }

    // Go: consumer.go — processAckProgress (+WPI): resets ack deadline to original ackWait without bumping delivery count
    public void ProcessProgress(ulong seq)
    {
        if (!_pending.TryGetValue(seq, out var state))
            return;

        state.DeadlineUtc = DateTime.UtcNow.AddMilliseconds(Math.Max(_ackWaitMs, 1));
        _pending[seq] = state;
    }

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

        state.Deliveries++;

        // Go: consumer.go — check maxDeliver; Deliveries > maxDeliver means the limit is exceeded
        if (_maxDeliver > 0 && state.Deliveries > _maxDeliver)
        {
            _exceededSequences.Add(sequence);
            _pending.Remove(sequence);
            _terminated.Add(sequence);
            TerminatedCount++;
            return;
        }

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

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

    /// <summary>
    /// Clears all pending acks, terminated set, and resets the ack floor.
    /// Used during consumer reset to specific sequence.
    /// Go reference: consumer.go processResetReq — clear all tracking state.
    /// </summary>
    public void ClearAll()
    {
        _pending.Clear();
        _terminated.Clear();
        AckFloor = 0;
        TerminatedCount = 0;
        _exceededSequences.Clear();
    }

    /// <summary>
    /// Resets the ack floor to the specified value.
    /// Used during consumer reset.
    /// </summary>
    public void SetAckFloor(ulong floor)
    {
        AckFloor = floor;
        // Remove any pending entries below or at the new floor
        foreach (var key in _pending.Keys.Where(k => k <= floor).ToArray())
            _pending.Remove(key);
    }

    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 bool HasSequenceBelow(ulong upTo)
    {
        foreach (var key in _pending.Keys)
        {
            if (key < upTo)
                return true;
        }
        return false;
    }

    private static bool TryParseInt(ReadOnlySpan<byte> span, out int value)
    {
        value = 0;
        if (span.IsEmpty)
            return false;

        foreach (var b in span)
        {
            if (b < (byte)'0' || b > (byte)'9')
                return false;
            value = value * 10 + (b - '0');
        }
        return true;
    }

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