namespace NATS.Server.Protocol;

public static class ProtoWire
{
    public const string ErrProtoInsufficient = "insufficient data to read a value";
    public const string ErrProtoOverflow = "too much data for a value";
    public const string ErrProtoInvalidFieldNumber = "invalid field number";

    public static (int Number, int WireType, int Size) ScanField(ReadOnlySpan<byte> buffer)
    {
        var (number, wireType, tagSize) = ScanTag(buffer);
        var valueSize = ScanFieldValue(wireType, buffer[tagSize..]);
        return (number, wireType, tagSize + valueSize);
    }

    public static (int Number, int WireType, int Size) ScanTag(ReadOnlySpan<byte> buffer)
    {
        var (tag, size) = ScanVarint(buffer);
        var fieldNumber = tag >> 3;
        if (fieldNumber > int.MaxValue || fieldNumber < 1)
            throw new ProtoWireException(ErrProtoInvalidFieldNumber);

        return ((int)fieldNumber, (int)(tag & 0x7), size);
    }

    public static int ScanFieldValue(int wireType, ReadOnlySpan<byte> buffer)
    {
        return wireType switch
        {
            0 => ScanVarint(buffer).Size,
            5 => 4,
            1 => 8,
            2 => ScanBytes(buffer),
            _ => throw new ProtoWireException($"unsupported type: {wireType}"),
        };
    }

    public static (ulong Value, int Size) ScanVarint(ReadOnlySpan<byte> buffer)
    {
        ulong value = 0;

        for (var i = 0; i < 10; i++)
        {
            if (i >= buffer.Length)
                throw new ProtoWireException(ErrProtoInsufficient);

            var b = buffer[i];
            if (i == 9)
            {
                if (b > 1)
                    throw new ProtoWireException(ErrProtoOverflow);

                value |= (ulong)b << 63;
                return (value, 10);
            }

            value |= (ulong)(b & 0x7F) << (i * 7);
            if ((b & 0x80) == 0)
                return (value, i + 1);
        }

        throw new ProtoWireException(ErrProtoOverflow);
    }

    public static int ScanBytes(ReadOnlySpan<byte> buffer)
    {
        var (length, lenSize) = ScanVarint(buffer);
        if (length > (ulong)buffer[lenSize..].Length)
            throw new ProtoWireException(ErrProtoInsufficient);

        return lenSize + (int)length;
    }

    public static byte[] EncodeVarint(ulong value)
    {
        Span<byte> scratch = stackalloc byte[10];
        var i = 0;
        while (value >= 0x80)
        {
            scratch[i++] = (byte)((value & 0x7F) | 0x80);
            value >>= 7;
        }

        scratch[i++] = (byte)value;
        return scratch[..i].ToArray();
    }
}

public sealed class ProtoWireException(string message) : Exception(message);
