namespace ZB.MOM.NatsNet.Server;

internal sealed partial class NatsConsumer
{
    private static readonly byte[] AckAck = "+ACK"u8.ToArray();
    private static readonly byte[] AckOk = "+OK"u8.ToArray();
    private static readonly byte[] AckNak = "-NAK"u8.ToArray();
    private static readonly byte[] AckProgress = "+WPI"u8.ToArray();
    private static readonly byte[] AckNext = "+NXT"u8.ToArray();

    private readonly Queue<JsAckMsg> _ackQueue = new();
    private string? _lastAckReplySubject;

    internal void SendAckReply(string subject)
    {
        if (string.IsNullOrWhiteSpace(subject))
            return;

        _mu.EnterWriteLock();
        try
        {
            _lastAckReplySubject = subject;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal static JsAckMsg NewJSAckMsg(string subject, string reply, int headerBytes, byte[] msg)
        => new()
        {
            Subject = subject,
            Reply = reply,
            HeaderBytes = headerBytes,
            Msg = msg,
        };

    internal void PushAck(string subject, string reply, int headerBytes, byte[] rawMessage)
    {
        ArgumentNullException.ThrowIfNull(rawMessage);
        _mu.EnterWriteLock();
        try
        {
            _ackQueue.Enqueue(NewJSAckMsg(subject, reply, headerBytes, (byte[])rawMessage.Clone()));
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal void ProcessAck(string subject, string reply, int headerBytes, byte[] rawMessage)
    {
        ArgumentNullException.ThrowIfNull(rawMessage);

        var msg = headerBytes > 0 && headerBytes <= rawMessage.Length
            ? rawMessage[headerBytes..]
            : rawMessage;

        var (streamSeq, deliverySeq, deliveryCount) = AckReplyInfo(subject);
        var skipAckReply = streamSeq == 0;

        if (msg.Length == 0 || msg.SequenceEqual(AckAck) || msg.SequenceEqual(AckOk))
        {
            ProcessAckMessage(streamSeq, deliverySeq, deliveryCount, reply);
        }
        else if (StartsWith(msg, AckNext))
        {
            ProcessAckMessage(streamSeq, deliverySeq, deliveryCount, string.Empty);
            skipAckReply = true;
        }
        else if (StartsWith(msg, AckNak))
        {
            _state.Redelivered ??= new Dictionary<ulong, ulong>();
            _state.Redelivered[streamSeq] = _state.Redelivered.TryGetValue(streamSeq, out var redeliveries)
                ? redeliveries + 1
                : 1UL;
            skipAckReply = true;
        }
        else if (msg.SequenceEqual(AckProgress))
        {
            ProgressUpdate(streamSeq);
        }

        if (!string.IsNullOrWhiteSpace(reply) && !skipAckReply)
            SendAckReply(reply);
    }

    internal bool ProcessAckMsg(ulong streamSequence, ulong deliverySequence, ulong deliveryCount, string reply, bool doSample)
    {
        ProcessAckMessage(streamSequence, deliverySequence, deliveryCount, reply);
        if (doSample && NeedAck())
            SampleAck(reply);
        return true;
    }

    internal bool ProcessNak(ulong streamSequence, ulong deliverySequence, ulong deliveryCount, byte[] message)
    {
        ArgumentNullException.ThrowIfNull(message);
        _mu.EnterWriteLock();
        try
        {
            _state.Redelivered ??= [];
            _state.Redelivered[streamSequence] = Math.Max(deliveryCount + 1, _state.Redelivered.GetValueOrDefault(streamSequence));
            _state.Pending ??= [];
            _state.Pending[streamSequence] = new Pending
            {
                Sequence = deliverySequence,
                Timestamp = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(),
            };
            return true;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal bool ProcessTerm(ulong streamSequence, ulong deliverySequence, ulong deliveryCount, string reason, string reply)
    {
        _mu.EnterWriteLock();
        try
        {
            _state.Pending?.Remove(streamSequence);
            _state.Redelivered ??= [];
            _state.Redelivered[streamSequence] = Math.Max(deliveryCount, _state.Redelivered.GetValueOrDefault(streamSequence));
            if (!string.IsNullOrWhiteSpace(reply))
                _lastAckReplySubject = reply;
            _ = reason;
            _ = deliverySequence;
            return true;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal TimeSpan AckWait(TimeSpan backOff)
    {
        if (backOff > TimeSpan.Zero)
            return backOff;

        return Config.AckWait > TimeSpan.Zero ? Config.AckWait : DefaultAckWait;
    }

    internal bool CheckRedelivered(ulong streamSequence)
    {
        _mu.EnterReadLock();
        try
        {
            return _state.Redelivered?.TryGetValue(streamSequence, out var count) == true && count > 1;
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }

    internal void ProgressUpdate(ulong sequence)
    {
        _mu.EnterWriteLock();
        try
        {
            _state.Pending ??= new Dictionary<ulong, Pending>();
            if (_state.Pending.TryGetValue(sequence, out var pending))
            {
                pending.Timestamp = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
                _state.Pending[sequence] = pending;
            }
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal void UpdateSkipped(ulong sequence)
    {
        _mu.EnterWriteLock();
        try
        {
            _state.AckFloor.Stream = Math.Max(_state.AckFloor.Stream, sequence > 0 ? sequence - 1 : 0);
            _state.AckFloor.Consumer = Math.Max(_state.AckFloor.Consumer, sequence > 0 ? sequence - 1 : 0);
            _updateChannel.Writer.TryWrite(true);
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    private void ProcessAckMessage(ulong streamSeq, ulong deliverySeq, ulong deliveryCount, string reply)
    {
        _mu.EnterWriteLock();
        try
        {
            _state.Pending ??= new Dictionary<ulong, Pending>();
            _state.Pending.Remove(streamSeq);
            _state.AckFloor.Stream = Math.Max(_state.AckFloor.Stream, streamSeq);
            _state.AckFloor.Consumer = Math.Max(_state.AckFloor.Consumer, deliverySeq);
            _state.Delivered.Stream = Math.Max(_state.Delivered.Stream, streamSeq);
            _state.Delivered.Consumer = Math.Max(_state.Delivered.Consumer, deliverySeq);

            _state.Redelivered ??= new Dictionary<ulong, ulong>();
            if (deliveryCount > 1)
                _state.Redelivered[streamSeq] = deliveryCount;

            if (!string.IsNullOrWhiteSpace(reply))
                _lastAckReplySubject = reply;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    private static (ulong StreamSequence, ulong DeliverySequence, ulong DeliveryCount) AckReplyInfo(string subject)
    {
        if (string.IsNullOrWhiteSpace(subject))
            return (0, 0, 0);

        var numbers = subject
            .Split('.', StringSplitOptions.RemoveEmptyEntries)
            .Select(static token => ulong.TryParse(token, out var value) ? (ulong?)value : null)
            .Where(static value => value.HasValue)
            .Select(static value => value!.Value)
            .ToArray();

        if (numbers.Length >= 3)
            return (numbers[^2], numbers[^3], numbers[^1]);
        if (numbers.Length == 2)
            return (numbers[1], numbers[0], 1);
        if (numbers.Length == 1)
            return (numbers[0], numbers[0], 1);

        return (0, 0, 0);
    }

    private static bool StartsWith(byte[] message, byte[] prefix)
    {
        if (message.Length < prefix.Length)
            return false;

        for (var i = 0; i < prefix.Length; i++)
        {
            if (message[i] != prefix[i])
                return false;
        }

        return true;
    }
}

internal sealed class JsAckMsg
{
    internal string Subject { get; set; } = string.Empty;
    internal string Reply { get; set; } = string.Empty;
    internal int HeaderBytes { get; set; }
    internal byte[] Msg { get; set; } = [];
}
