using System.Buffers;
using System.Text;
using Microsoft.Extensions.Logging;

namespace NATS.Server.Protocol;

public enum CommandType
{
    Ping,
    Pong,
    Connect,
    Info,
    Pub,
    HPub,
    Sub,
    Unsub,
    Ok,
    Err,
}

public readonly struct ParsedCommand
{
    public CommandType Type { get; init; }
    public string? Subject { get; init; }
    public string? ReplyTo { get; init; }
    public string? Queue { get; init; }
    public string? Sid { get; init; }
    public int MaxMessages { get; init; }
    public int HeaderSize { get; init; }
    public ReadOnlyMemory<byte> Payload { get; init; }

    public static ParsedCommand Simple(CommandType type) => new() { Type = type, MaxMessages = -1 };
}

public sealed class NatsParser
{
    private static readonly byte[] CrLfBytes = "\r\n"u8.ToArray();
    private readonly int _maxPayload;
    private ILogger? _logger;
    public ILogger? Logger { set => _logger = value; }

    // State for split-packet payload reading
    private bool _awaitingPayload;
    private int _expectedPayloadSize;
    private string? _pendingSubject;
    private string? _pendingReplyTo;
    private int _pendingHeaderSize;
    private CommandType _pendingType;

    public NatsParser(int maxPayload = NatsProtocol.MaxPayloadSize, ILogger? logger = null)
    {
        _maxPayload = maxPayload;
        _logger = logger;
    }

    private void TraceInOp(string op, ReadOnlySpan<byte> arg = default)
    {
        if (_logger == null || !_logger.IsEnabled(LogLevel.Trace))
            return;
        if (arg.IsEmpty)
            _logger.LogTrace("<<- {Op}", op);
        else
            _logger.LogTrace("<<- {Op} {Arg}", op, Encoding.ASCII.GetString(arg));
    }

    public bool TryParse(ref ReadOnlySequence<byte> buffer, out ParsedCommand command)
    {
        command = default;

        if (_awaitingPayload)
            return TryReadPayload(ref buffer, out command);

        // Look for \r\n to find control line
        var reader = new SequenceReader<byte>(buffer);
        if (!reader.TryReadTo(out ReadOnlySequence<byte> line, CrLfBytes.AsSpan()))
            return false;

        // Control line size check
        if (line.Length > NatsProtocol.MaxControlLineSize)
            throw new ProtocolViolationException("Maximum control line exceeded");

        // Get line as contiguous span
        Span<byte> lineSpan = stackalloc byte[(int)line.Length];
        line.CopyTo(lineSpan);

        // Identify command by first bytes
        if (lineSpan.Length < 2)
        {
            if (lineSpan.Length == 1 && lineSpan[0] is (byte)'+')
            {
                // partial -- need more data
                return false;
            }

            throw new ProtocolViolationException("Unknown protocol operation");
        }

        byte b0 = (byte)(lineSpan[0] | 0x20); // lowercase
        byte b1 = (byte)(lineSpan[1] | 0x20);

        switch (b0)
        {
            case (byte)'p':
                if (b1 == (byte)'i') // PING
                {
                    command = ParsedCommand.Simple(CommandType.Ping);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("PING");
                    return true;
                }

                if (b1 == (byte)'o') // PONG
                {
                    command = ParsedCommand.Simple(CommandType.Pong);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("PONG");
                    return true;
                }

                if (b1 == (byte)'u') // PUB
                {
                    return ParsePub(lineSpan, ref buffer, reader.Position, out command);
                }

                break;

            case (byte)'h':
                if (b1 == (byte)'p') // HPUB
                {
                    return ParseHPub(lineSpan, ref buffer, reader.Position, out command);
                }

                break;

            case (byte)'s':
                if (b1 == (byte)'u') // SUB
                {
                    command = ParseSub(lineSpan);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("SUB", lineSpan[4..]);
                    return true;
                }

                break;

            case (byte)'u':
                if (b1 == (byte)'n') // UNSUB
                {
                    command = ParseUnsub(lineSpan);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("UNSUB", lineSpan[6..]);
                    return true;
                }

                break;

            case (byte)'c':
                if (b1 == (byte)'o') // CONNECT
                {
                    command = ParseConnect(lineSpan);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("CONNECT");
                    return true;
                }

                break;

            case (byte)'i':
                if (b1 == (byte)'n') // INFO
                {
                    command = ParseInfo(lineSpan);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("INFO");
                    return true;
                }

                break;

            case (byte)'+': // +OK
                command = ParsedCommand.Simple(CommandType.Ok);
                buffer = buffer.Slice(reader.Position);
                TraceInOp("+OK");
                return true;

            case (byte)'-': // -ERR
                command = ParsedCommand.Simple(CommandType.Err);
                buffer = buffer.Slice(reader.Position);
                TraceInOp("-ERR");
                return true;
        }

        throw new ProtocolViolationException("Unknown protocol operation");
    }

    private bool ParsePub(
        Span<byte> line,
        ref ReadOnlySequence<byte> buffer,
        SequencePosition afterLine,
        out ParsedCommand command)
    {
        command = default;

        // PUB subject [reply] size -- skip "PUB "
        Span<Range> ranges = stackalloc Range[4];
        var argsSpan = line[4..];
        int argCount = SplitArgs(argsSpan, ranges);

        string subject;
        string? reply = null;
        int size;

        if (argCount == 2)
        {
            subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
            size = ParseSize(argsSpan[ranges[1]]);
        }
        else if (argCount == 3)
        {
            subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
            reply = Encoding.ASCII.GetString(argsSpan[ranges[1]]);
            size = ParseSize(argsSpan[ranges[2]]);
        }
        else
        {
            throw new ProtocolViolationException("Invalid PUB arguments");
        }

        if (size < 0)
            throw new ProtocolViolationException("Invalid payload size");

        // Now read payload + \r\n (max payload enforcement is done at the client level)
        buffer = buffer.Slice(afterLine);
        _awaitingPayload = true;
        _expectedPayloadSize = size;
        _pendingSubject = subject;
        _pendingReplyTo = reply;
        _pendingHeaderSize = -1;
        _pendingType = CommandType.Pub;

        TraceInOp("PUB", argsSpan);
        return TryReadPayload(ref buffer, out command);
    }

    private bool ParseHPub(
        Span<byte> line,
        ref ReadOnlySequence<byte> buffer,
        SequencePosition afterLine,
        out ParsedCommand command)
    {
        command = default;

        // HPUB subject [reply] hdr_size total_size -- skip "HPUB "
        Span<Range> ranges = stackalloc Range[5];
        var argsSpan = line[5..];
        int argCount = SplitArgs(argsSpan, ranges);

        string subject;
        string? reply = null;
        int hdrSize, totalSize;

        if (argCount == 3)
        {
            subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
            hdrSize = ParseSize(argsSpan[ranges[1]]);
            totalSize = ParseSize(argsSpan[ranges[2]]);
        }
        else if (argCount == 4)
        {
            subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]);
            reply = Encoding.ASCII.GetString(argsSpan[ranges[1]]);
            hdrSize = ParseSize(argsSpan[ranges[2]]);
            totalSize = ParseSize(argsSpan[ranges[3]]);
        }
        else
        {
            throw new ProtocolViolationException("Invalid HPUB arguments");
        }

        if (hdrSize < 0 || totalSize < 0 || hdrSize > totalSize)
            throw new ProtocolViolationException("Invalid HPUB sizes");

        buffer = buffer.Slice(afterLine);
        _awaitingPayload = true;
        _expectedPayloadSize = totalSize;
        _pendingSubject = subject;
        _pendingReplyTo = reply;
        _pendingHeaderSize = hdrSize;
        _pendingType = CommandType.HPub;

        TraceInOp("HPUB", argsSpan);
        return TryReadPayload(ref buffer, out command);
    }

    private bool TryReadPayload(ref ReadOnlySequence<byte> buffer, out ParsedCommand command)
    {
        command = default;

        // Need: _expectedPayloadSize bytes + \r\n
        long needed = _expectedPayloadSize + 2; // payload + \r\n
        if (buffer.Length < needed)
            return false;

        // Extract payload
        var payloadSlice = buffer.Slice(0, _expectedPayloadSize);
        var payload = new byte[_expectedPayloadSize];
        payloadSlice.CopyTo(payload);

        // Verify \r\n after payload
        var trailer = buffer.Slice(_expectedPayloadSize, 2);
        Span<byte> trailerBytes = stackalloc byte[2];
        trailer.CopyTo(trailerBytes);
        if (trailerBytes[0] != (byte)'\r' || trailerBytes[1] != (byte)'\n')
            throw new ProtocolViolationException("Expected \\r\\n after payload");

        command = new ParsedCommand
        {
            Type = _pendingType,
            Subject = _pendingSubject,
            ReplyTo = _pendingReplyTo,
            Payload = payload,
            HeaderSize = _pendingHeaderSize,
            MaxMessages = -1,
        };

        buffer = buffer.Slice(buffer.GetPosition(needed));
        _awaitingPayload = false;
        return true;
    }

    private static ParsedCommand ParseSub(Span<byte> line)
    {
        // SUB subject [queue] sid -- skip "SUB "
        Span<Range> ranges = stackalloc Range[4];
        var argsSpan = line[4..];
        int argCount = SplitArgs(argsSpan, ranges);

        return argCount switch
        {
            2 => new ParsedCommand
            {
                Type = CommandType.Sub,
                Subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
                Sid = Encoding.ASCII.GetString(argsSpan[ranges[1]]),
                MaxMessages = -1,
            },
            3 => new ParsedCommand
            {
                Type = CommandType.Sub,
                Subject = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
                Queue = Encoding.ASCII.GetString(argsSpan[ranges[1]]),
                Sid = Encoding.ASCII.GetString(argsSpan[ranges[2]]),
                MaxMessages = -1,
            },
            _ => throw new ProtocolViolationException("Invalid SUB arguments"),
        };
    }

    private static ParsedCommand ParseUnsub(Span<byte> line)
    {
        // UNSUB sid [max_msgs] -- skip "UNSUB "
        Span<Range> ranges = stackalloc Range[3];
        var argsSpan = line[6..];
        int argCount = SplitArgs(argsSpan, ranges);

        return argCount switch
        {
            1 => new ParsedCommand
            {
                Type = CommandType.Unsub,
                Sid = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
                MaxMessages = -1,
            },
            2 => new ParsedCommand
            {
                Type = CommandType.Unsub,
                Sid = Encoding.ASCII.GetString(argsSpan[ranges[0]]),
                MaxMessages = ParseSize(argsSpan[ranges[1]]),
            },
            _ => throw new ProtocolViolationException("Invalid UNSUB arguments"),
        };
    }

    private static ParsedCommand ParseConnect(Span<byte> line)
    {
        // CONNECT {json} -- find first space after command
        int spaceIdx = line.IndexOf((byte)' ');
        if (spaceIdx < 0)
            throw new ProtocolViolationException("Invalid CONNECT");

        var json = line[(spaceIdx + 1)..];
        return new ParsedCommand
        {
            Type = CommandType.Connect,
            Payload = json.ToArray(),
            MaxMessages = -1,
        };
    }

    private static ParsedCommand ParseInfo(Span<byte> line)
    {
        // INFO {json} -- find first space after command
        int spaceIdx = line.IndexOf((byte)' ');
        if (spaceIdx < 0)
            throw new ProtocolViolationException("Invalid INFO");

        var json = line[(spaceIdx + 1)..];
        return new ParsedCommand
        {
            Type = CommandType.Info,
            Payload = json.ToArray(),
            MaxMessages = -1,
        };
    }

    /// <summary>
    /// Parse a decimal integer from ASCII bytes. Returns -1 on error.
    /// </summary>
    internal static int ParseSize(Span<byte> data)
    {
        if (data.Length == 0 || data.Length > 9)
            return -1;
        int n = 0;
        foreach (byte b in data)
        {
            if (b < (byte)'0' || b > (byte)'9')
                return -1;
            n = n * 10 + (b - '0');
        }

        return n;
    }

    /// <summary>
    /// Split by spaces/tabs into argument ranges. Returns the number of arguments found.
    /// Uses Span&lt;Range&gt; for zero-allocation argument splitting.
    /// </summary>
    internal static int SplitArgs(Span<byte> data, Span<Range> ranges)
    {
        int count = 0;
        int start = -1;

        for (int i = 0; i < data.Length; i++)
        {
            byte b = data[i];
            if (b is (byte)' ' or (byte)'\t')
            {
                if (start >= 0)
                {
                    if (count >= ranges.Length)
                        throw new ProtocolViolationException("Too many arguments");
                    ranges[count++] = start..i;
                    start = -1;
                }
            }
            else
            {
                if (start < 0)
                    start = i;
            }
        }

        if (start >= 0)
        {
            if (count >= ranges.Length)
                throw new ProtocolViolationException("Too many arguments");
            ranges[count++] = start..data.Length;
        }

        return count;
    }
}

public class ProtocolViolationException : Exception
{
    public ProtocolViolationException(string message)
        : base(message)
    {
    }
}