using System.Buffers.Binary;
using System.Text;

namespace NATS.Server.Mqtt;

public static class MqttPacketWriter
{
    public static byte[] WriteString(string value)
        => WriteBytes(Encoding.UTF8.GetBytes(value));

    public static byte[] WriteBytes(ReadOnlySpan<byte> bytes)
    {
        if (bytes.Length > ushort.MaxValue)
            throw new ArgumentOutOfRangeException(nameof(bytes), "MQTT length-prefixed field cannot exceed 65535 bytes.");

        var buffer = new byte[2 + bytes.Length];
        BinaryPrimitives.WriteUInt16BigEndian(buffer.AsSpan(0, 2), (ushort)bytes.Length);
        bytes.CopyTo(buffer.AsSpan(2));
        return buffer;
    }

    public static byte[] Write(MqttControlPacketType type, ReadOnlySpan<byte> payload, byte flags = 0)
    {
        if (type == MqttControlPacketType.Reserved)
            throw new ArgumentOutOfRangeException(nameof(type), "MQTT control packet type must be non-zero.");

        var remainingLength = payload.Length;
        var encodedRemainingLength = EncodeRemainingLength(remainingLength);
        var buffer = new byte[1 + encodedRemainingLength.Length + remainingLength];
        buffer[0] = (byte)(((byte)type << 4) | (flags & 0x0F));
        encodedRemainingLength.CopyTo(buffer.AsSpan(1));
        payload.CopyTo(buffer.AsSpan(1 + encodedRemainingLength.Length));
        return buffer;
    }

    internal static byte[] EncodeRemainingLength(int value)
    {
        if (value < 0 || value > MqttProtocolConstants.MaxPayloadSize)
            throw new ArgumentOutOfRangeException(
                nameof(value),
                $"MQTT remaining length must be between 0 and {MqttProtocolConstants.MaxPayloadSize}.");

        Span<byte> scratch = stackalloc byte[4];
        var index = 0;

        do
        {
            var digit = (byte)(value % 128);
            value /= 128;
            if (value > 0)
                digit |= 0x80;
            scratch[index++] = digit;
        } while (value > 0);

        return scratch[..index].ToArray();
    }
}
