// Copyright 2024 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Adapted from server/proto.go in the NATS server Go source.
// Inspired by https://github.com/protocolbuffers/protobuf-go/blob/master/encoding/protowire/wire.go

namespace ZB.MOM.NatsNet.Server.Internal;

/// <summary>
/// Low-level protobuf wire format helpers used internally for JetStream API encoding.
/// Mirrors server/proto.go.
/// </summary>
public static class ProtoWire
{
    private static readonly Exception ErrInsufficient =
        new InvalidOperationException("insufficient data to read a value");

    private static readonly Exception ErrOverflow =
        new InvalidOperationException("too much data for a value");

    private static readonly Exception ErrInvalidFieldNumber =
        new InvalidOperationException("invalid field number");

    // -------------------------------------------------------------------------
    // Field scanning
    // -------------------------------------------------------------------------

    /// <summary>
    /// Reads a complete protobuf field (tag + value) from <paramref name="b"/>
    /// and returns the field number, wire type, and total byte size consumed.
    /// Mirrors <c>protoScanField</c>.
    /// </summary>
    public static (int num, int typ, int size, Exception? err) ScanField(ReadOnlySpan<byte> b)
    {
        var (num, typ, sizeTag, err) = ScanTag(b);
        if (err != null) return (0, 0, 0, err);

        var (sizeValue, err2) = ScanFieldValue(typ, b[sizeTag..]);
        if (err2 != null) return (0, 0, 0, err2);

        return (num, typ, sizeTag + sizeValue, null);
    }

    /// <summary>
    /// Reads a protobuf tag varint and returns field number, wire type, and bytes consumed.
    /// Mirrors <c>protoScanTag</c>.
    /// </summary>
    public static (int num, int typ, int size, Exception? err) ScanTag(ReadOnlySpan<byte> b)
    {
        var (tagint, size, err) = ScanVarint(b);
        if (err != null) return (0, 0, 0, err);

        if ((tagint >> 3) > int.MaxValue)
            return (0, 0, 0, ErrInvalidFieldNumber);

        var num = (int)(tagint >> 3);
        if (num < 1)
            return (0, 0, 0, ErrInvalidFieldNumber);

        var typ = (int)(tagint & 7);
        return (num, typ, size, null);
    }

    /// <summary>
    /// Returns the byte count consumed by a field value with the given wire type.
    /// Mirrors <c>protoScanFieldValue</c>.
    /// </summary>
    public static (int size, Exception? err) ScanFieldValue(int typ, ReadOnlySpan<byte> b)
    {
        switch (typ)
        {
            case 0: // varint
            {
                var (_, size, err) = ScanVarint(b);
                return (size, err);
            }
            case 5: // fixed32
                return (4, null);
            case 1: // fixed64
                return (8, null);
            case 2: // length-delimited
            {
                var (size, err) = ScanBytes(b);
                return (size, err);
            }
            default:
                return (0, new InvalidOperationException($"unsupported type: {typ}"));
        }
    }

    // -------------------------------------------------------------------------
    // Varint decode
    // -------------------------------------------------------------------------

    /// <summary>
    /// Decodes a protobuf varint from <paramref name="b"/>.
    /// Returns (value, bytes_consumed, error).
    /// Mirrors <c>protoScanVarint</c>.
    /// </summary>
    public static (ulong v, int size, Exception? err) ScanVarint(ReadOnlySpan<byte> b)
    {
        if (b.Length < 1) return (0, 0, ErrInsufficient);
        ulong v = b[0];
        if (v < 0x80) return (v, 1, null);
        v -= 0x80;

        if (b.Length < 2) return (0, 0, ErrInsufficient);
        ulong y = b[1];
        v += y << 7;
        if (y < 0x80) return (v, 2, null);
        v -= 0x80UL << 7;

        if (b.Length < 3) return (0, 0, ErrInsufficient);
        y = b[2];
        v += y << 14;
        if (y < 0x80) return (v, 3, null);
        v -= 0x80UL << 14;

        if (b.Length < 4) return (0, 0, ErrInsufficient);
        y = b[3];
        v += y << 21;
        if (y < 0x80) return (v, 4, null);
        v -= 0x80UL << 21;

        if (b.Length < 5) return (0, 0, ErrInsufficient);
        y = b[4];
        v += y << 28;
        if (y < 0x80) return (v, 5, null);
        v -= 0x80UL << 28;

        if (b.Length < 6) return (0, 0, ErrInsufficient);
        y = b[5];
        v += y << 35;
        if (y < 0x80) return (v, 6, null);
        v -= 0x80UL << 35;

        if (b.Length < 7) return (0, 0, ErrInsufficient);
        y = b[6];
        v += y << 42;
        if (y < 0x80) return (v, 7, null);
        v -= 0x80UL << 42;

        if (b.Length < 8) return (0, 0, ErrInsufficient);
        y = b[7];
        v += y << 49;
        if (y < 0x80) return (v, 8, null);
        v -= 0x80UL << 49;

        if (b.Length < 9) return (0, 0, ErrInsufficient);
        y = b[8];
        v += y << 56;
        if (y < 0x80) return (v, 9, null);
        v -= 0x80UL << 56;

        if (b.Length < 10) return (0, 0, ErrInsufficient);
        y = b[9];
        v += y << 63;
        if (y < 2) return (v, 10, null);

        return (0, 0, ErrOverflow);
    }

    // -------------------------------------------------------------------------
    // Length-delimited decode
    // -------------------------------------------------------------------------

    /// <summary>
    /// Returns the total byte count consumed by a length-delimited field
    /// (length varint + content).
    /// Mirrors <c>protoScanBytes</c>.
    /// </summary>
    public static (int size, Exception? err) ScanBytes(ReadOnlySpan<byte> b)
    {
        var (l, lenSize, err) = ScanVarint(b);
        if (err != null) return (0, err);

        if (l > (ulong)(b.Length - lenSize))
            return (0, ErrInsufficient);

        return (lenSize + (int)l, null);
    }

    // -------------------------------------------------------------------------
    // Varint encode
    // -------------------------------------------------------------------------

    /// <summary>
    /// Encodes a <see cref="ulong"/> as a protobuf varint.
    /// Mirrors <c>protoEncodeVarint</c>.
    /// </summary>
    public static byte[] EncodeVarint(ulong v)
    {
        if (v < 1UL << 7)
            return [(byte)v];

        if (v < 1UL << 14)
            return [(byte)((v >> 0) & 0x7F | 0x80), (byte)(v >> 7)];

        if (v < 1UL << 21)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)(v >> 14)];

        if (v < 1UL << 28)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)((v >> 14) & 0x7F | 0x80),
                (byte)(v >> 21)];

        if (v < 1UL << 35)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)((v >> 14) & 0x7F | 0x80),
                (byte)((v >> 21) & 0x7F | 0x80),
                (byte)(v >> 28)];

        if (v < 1UL << 42)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)((v >> 14) & 0x7F | 0x80),
                (byte)((v >> 21) & 0x7F | 0x80),
                (byte)((v >> 28) & 0x7F | 0x80),
                (byte)(v >> 35)];

        if (v < 1UL << 49)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)((v >> 14) & 0x7F | 0x80),
                (byte)((v >> 21) & 0x7F | 0x80),
                (byte)((v >> 28) & 0x7F | 0x80),
                (byte)((v >> 35) & 0x7F | 0x80),
                (byte)(v >> 42)];

        if (v < 1UL << 56)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)((v >> 14) & 0x7F | 0x80),
                (byte)((v >> 21) & 0x7F | 0x80),
                (byte)((v >> 28) & 0x7F | 0x80),
                (byte)((v >> 35) & 0x7F | 0x80),
                (byte)((v >> 42) & 0x7F | 0x80),
                (byte)(v >> 49)];

        if (v < 1UL << 63)
            return [
                (byte)((v >> 0) & 0x7F | 0x80),
                (byte)((v >> 7) & 0x7F | 0x80),
                (byte)((v >> 14) & 0x7F | 0x80),
                (byte)((v >> 21) & 0x7F | 0x80),
                (byte)((v >> 28) & 0x7F | 0x80),
                (byte)((v >> 35) & 0x7F | 0x80),
                (byte)((v >> 42) & 0x7F | 0x80),
                (byte)((v >> 49) & 0x7F | 0x80),
                (byte)(v >> 56)];

        return [
            (byte)((v >> 0) & 0x7F | 0x80),
            (byte)((v >> 7) & 0x7F | 0x80),
            (byte)((v >> 14) & 0x7F | 0x80),
            (byte)((v >> 21) & 0x7F | 0x80),
            (byte)((v >> 28) & 0x7F | 0x80),
            (byte)((v >> 35) & 0x7F | 0x80),
            (byte)((v >> 42) & 0x7F | 0x80),
            (byte)((v >> 49) & 0x7F | 0x80),
            (byte)((v >> 56) & 0x7F | 0x80),
            1];
    }
}