using System.Buffers;
using System.Text;
using System.Text.Json;
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? Operation { 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, string operation) =>
        new() { Type = type, Operation = operation, MaxMessages = -1 };
}

public sealed class NatsParser
{
    private static ReadOnlySpan<byte> CrLfBytes => "\r\n"u8;
    private ILogger? _logger;
    public ILogger? Logger { set => _logger = value; }

    // State for split-packet payload reading
    private bool _awaitingPayload;
    private int _expectedPayloadSize;
    private byte[]? _pendingSubject;
    private byte[]? _pendingReplyTo;
    private int _pendingHeaderSize;
    private CommandType _pendingType;
    private string _pendingOperation = string.Empty;

    public NatsParser(int maxPayload = NatsProtocol.MaxPayloadSize, ILogger? logger = null)
    {
        _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 (!TryParseView(ref buffer, out var view))
            return false;

        command = view.Materialize();
        return true;
    }

    internal bool TryParseView(ref ReadOnlySequence<byte> buffer, out ParsedCommandView 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))
            return false;

        // Control line size check
        if (line.Length > NatsProtocol.MaxControlLineSize)
        {
            var snippetLength = (int)Math.Min(line.Length, NatsProtocol.MaxControlLineSnippetSize);
            var snippetBytes = new byte[snippetLength];
            line.Slice(0, snippetLength).CopyTo(snippetBytes);
            var snippet = ProtoSnippet(0, NatsProtocol.MaxControlLineSnippetSize, snippetBytes);
            throw new ProtocolViolationException(
                $"Maximum control line exceeded (max={NatsProtocol.MaxControlLineSize}, len={line.Length}, snip={snippet}...)");
        }

        // 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: {ProtoSnippet(lineSpan)}");
        }

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

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

                if (b1 == (byte)'o') // PONG
                {
                    command = ParsedCommandView.Simple(CommandType.Pong, "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(line);
                    buffer = buffer.Slice(reader.Position);
                    TraceInOp("CONNECT");
                    return true;
                }

                break;

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

                break;

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

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

        throw new ProtocolViolationException($"Unknown protocol operation: {ProtoSnippet(lineSpan)}");
    }

    // Go reference: parser.go protoSnippet(start, max, buf).
    internal static string ProtoSnippet(int start, int max, ReadOnlySpan<byte> buffer)
    {
        if (start >= buffer.Length)
            return "\"\"";

        var stop = start + max;
        if (stop > buffer.Length)
            stop = buffer.Length - 1;

        if (stop <= start)
            return "\"\"";

        var slice = buffer[start..stop];
        return JsonSerializer.Serialize(Encoding.ASCII.GetString(slice));
    }

    internal static string ProtoSnippet(ReadOnlySpan<byte> buffer) =>
        ProtoSnippet(0, NatsProtocol.ProtoSnippetSize, buffer);

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

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

        byte[] subject;
        byte[]? reply = null;
        int size;

        if (argCount == 2)
        {
            subject = argsSpan[ranges[0]].ToArray();
            size = ParseSize(argsSpan[ranges[1]]);
        }
        else if (argCount == 3)
        {
            subject = argsSpan[ranges[0]].ToArray();
            reply = argsSpan[ranges[1]].ToArray();
            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;
        _pendingOperation = "PUB";

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

    private bool ParseHPub(
        Span<byte> line,
        ref ReadOnlySequence<byte> buffer,
        SequencePosition afterLine,
        out ParsedCommandView 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);

        byte[] subject;
        byte[]? reply = null;
        int hdrSize, totalSize;

        if (argCount == 3)
        {
            subject = argsSpan[ranges[0]].ToArray();
            hdrSize = ParseSize(argsSpan[ranges[1]]);
            totalSize = ParseSize(argsSpan[ranges[2]]);
        }
        else if (argCount == 4)
        {
            subject = argsSpan[ranges[0]].ToArray();
            reply = argsSpan[ranges[1]].ToArray();
            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;
        _pendingOperation = "HPUB";

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

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

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

        var payloadSlice = buffer.Slice(0, _expectedPayloadSize);

        // 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 ParsedCommandView
        {
            Type = _pendingType,
            Operation = _pendingOperation,
            Subject = _pendingSubject,
            ReplyTo = _pendingReplyTo,
            Payload = payloadSlice,
            HeaderSize = _pendingHeaderSize,
            MaxMessages = -1,
        };

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

    private static ParsedCommandView ParseSub(Span<byte> line)
    {
        // SUB subject [queue] sid -- skip "SUB "
        if (line.Length < 5)
            throw new ProtocolViolationException("Invalid SUB arguments");
        Span<Range> ranges = stackalloc Range[4];
        var argsSpan = line[4..];
        int argCount = SplitArgs(argsSpan, ranges);

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

    private static ParsedCommandView ParseUnsub(Span<byte> line)
    {
        // UNSUB sid [max_msgs] -- skip "UNSUB "
        if (line.Length < 7)
            throw new ProtocolViolationException("Invalid UNSUB arguments");
        Span<Range> ranges = stackalloc Range[3];
        var argsSpan = line[6..];
        int argCount = SplitArgs(argsSpan, ranges);

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

    private static ParsedCommandView ParseConnect(ReadOnlySequence<byte> line)
    {
        var reader = new SequenceReader<byte>(line);
        if (!reader.TryAdvanceTo((byte)' ', advancePastDelimiter: true))
            throw new ProtocolViolationException("Invalid CONNECT");

        var json = line.Slice(reader.Position);
        return new ParsedCommandView
        {
            Type = CommandType.Connect,
            Operation = "CONNECT",
            Payload = json,
            MaxMessages = -1,
        };
    }

    private static ParsedCommandView ParseInfo(ReadOnlySequence<byte> line)
    {
        var reader = new SequenceReader<byte>(line);
        if (!reader.TryAdvanceTo((byte)' ', advancePastDelimiter: true))
            throw new ProtocolViolationException("Invalid INFO");

        var json = line.Slice(reader.Position);
        return new ParsedCommandView
        {
            Type = CommandType.Info,
            Operation = "INFO",
            Payload = json,
            MaxMessages = -1,
        };
    }

    private static ReadOnlyMemory<byte> CopyBytes(ReadOnlySpan<byte> value) =>
        value.IsEmpty ? ReadOnlyMemory<byte>.Empty : value.ToArray();

    /// <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)
    {
    }
}