using System.Buffers;

namespace NATS.Server.Mqtt;

public enum MqttControlPacketType : byte
{
    Reserved = 0,
    Connect = 1,
    ConnAck = 2,
    Publish = 3,
    PubAck = 4,
    PubRec = 5,
    PubRel = 6,
    PubComp = 7,
    Subscribe = 8,
    SubAck = 9,
    Unsubscribe = 10,
    UnsubAck = 11,
    PingReq = 12,
    PingResp = 13,
    Disconnect = 14,
}

public sealed record MqttControlPacket(
    MqttControlPacketType Type,
    byte Flags,
    int RemainingLength,
    ReadOnlyMemory<byte> Payload);

public static class MqttPacketReader
{
    /// <summary>
    /// Parses a complete MQTT control packet from a contiguous span.
    /// </summary>
    public static MqttControlPacket Read(ReadOnlySpan<byte> buffer)
    {
        if (buffer.Length < 2)
            throw new FormatException("MQTT packet is shorter than fixed header.");

        var first = buffer[0];
        var type = (MqttControlPacketType)(first >> 4);
        var flags = (byte)(first & 0x0F);
        var remainingLength = DecodeRemainingLength(buffer[1..], out var consumed);
        if (remainingLength > MqttProtocolConstants.MaxPayloadSize)
            throw new FormatException("MQTT packet remaining length exceeds protocol maximum.");
        var payloadStart = 1 + consumed;
        var totalLength = payloadStart + remainingLength;
        if (remainingLength < 0 || totalLength > buffer.Length)
            throw new FormatException("MQTT packet remaining length exceeds available bytes.");

        var payload = buffer[payloadStart..totalLength].ToArray();
        return new MqttControlPacket(type, flags, remainingLength, payload);
    }

    /// <summary>
    /// Attempts to read a complete MQTT control packet from a <see cref="ReadOnlySequence{T}"/>.
    /// Returns false if more data is needed (partial read). Advances <paramref name="consumed"/>
    /// past the packet bytes on success.
    /// Used with <see cref="System.IO.Pipelines.PipeReader"/> for incremental parsing.
    /// </summary>
    public static bool TryRead(ReadOnlySequence<byte> buffer, out MqttControlPacket? packet, out SequencePosition consumed)
    {
        packet = null;
        consumed = buffer.Start;

        if (buffer.Length < 2)
            return false;

        // Read the fixed header byte
        var reader = new SequenceReader<byte>(buffer);
        reader.TryRead(out var firstByte);
        var type = (MqttControlPacketType)(firstByte >> 4);
        var flags = (byte)(firstByte & 0x0F);

        // Decode remaining length (variable 1-4 bytes)
        var multiplier = 1;
        var remainingLength = 0;
        var lengthBytesConsumed = 0;

        while (lengthBytesConsumed < 4)
        {
            if (!reader.TryRead(out var digit))
                return false; // need more data

            lengthBytesConsumed++;
            remainingLength += (digit & 0x7F) * multiplier;

            if ((digit & 0x80) == 0)
                break;

            multiplier *= 128;

            if (lengthBytesConsumed == 4)
                throw new FormatException("Invalid MQTT remaining length encoding.");
        }

        if (remainingLength > MqttProtocolConstants.MaxPayloadSize)
            throw new FormatException("MQTT packet remaining length exceeds protocol maximum.");

        // Check if we have the full payload
        var headerSize = 1 + lengthBytesConsumed;
        var totalPacketSize = headerSize + remainingLength;
        if (buffer.Length < totalPacketSize)
            return false; // need more data

        // Extract payload
        byte[] payload;
        if (remainingLength == 0)
        {
            payload = [];
        }
        else
        {
            payload = new byte[remainingLength];
            buffer.Slice(headerSize, remainingLength).CopyTo(payload);
        }

        packet = new MqttControlPacket(type, flags, remainingLength, payload);
        consumed = buffer.GetPosition(totalPacketSize);
        return true;
    }

    internal static int DecodeRemainingLength(ReadOnlySpan<byte> encoded, out int consumed)
    {
        var multiplier = 1;
        var value = 0;
        consumed = 0;

        for (var i = 0; i < encoded.Length && i < 4; i++)
        {
            var digit = encoded[i];
            consumed++;
            value += (digit & 0x7F) * multiplier;

            if ((digit & 0x80) == 0)
                return value;

            multiplier *= 128;
        }

        throw new FormatException("Invalid MQTT remaining length encoding.");
    }
}
