using System.Buffers.Binary;

namespace NATS.Server.JetStream.Storage;

// Go: server/store.go:397 encodeConsumerState / server/filestore.go:12216 decodeConsumerState
/// <summary>
/// Binary encode/decode of ConsumerState matching the Go wire format.
///
/// Wire layout (version 2):
///   [0]  magic   = 22 (0x16)
///   [1]  version = 2
///   varuint: AckFloor.Consumer
///   varuint: AckFloor.Stream
///   varuint: Delivered.Consumer
///   varuint: Delivered.Stream
///   varuint: len(Pending)
///   if len(Pending) > 0:
///       varint:  mints (base Unix seconds, now rounded to second)
///       for each pending entry:
///           varuint: sseq - AckFloor.Stream
///           varuint: dseq - AckFloor.Consumer
///           varint:  mints - (ts / 1e9)   (delta seconds)
///   varuint: len(Redelivered)
///   for each redelivered entry:
///       varuint: sseq - AckFloor.Stream
///       varuint: count
///
/// Reference: golang/nats-server/server/store.go:397  (encodeConsumerState)
///            golang/nats-server/server/filestore.go:12216 (decodeConsumerState)
/// </summary>
public static class ConsumerStateCodec
{
    // Go: filestore.go:285
    private const byte Magic = 22;
    // Go: filestore.go:291  hdrLen = 2
    private const int HdrLen = 2;
    // Go: filestore.go:11852  seqsHdrSize = 6*binary.MaxVarintLen64 + hdrLen
    // binary.MaxVarintLen64 = 10 bytes
    private const int SeqsHdrSize = 6 * 10 + HdrLen;

    private const ulong HighBit = 1UL << 63;

    /// <summary>
    /// Encodes consumer state into the Go-compatible binary format.
    /// Reference: golang/nats-server/server/store.go:397
    /// </summary>
    public static byte[] Encode(ConsumerState state)
    {
        // Upper-bound the buffer size.
        var maxSize = SeqsHdrSize;
        if (state.Pending?.Count > 0)
            maxSize += state.Pending.Count * (3 * 10) + 10;
        if (state.Redelivered?.Count > 0)
            maxSize += state.Redelivered.Count * (2 * 10) + 10;

        var buf = new byte[maxSize];
        buf[0] = Magic;
        buf[1] = 2; // version 2

        var n = HdrLen;
        n += PutUvarint(buf, n, state.AckFloor.Consumer);
        n += PutUvarint(buf, n, state.AckFloor.Stream);
        n += PutUvarint(buf, n, state.Delivered.Consumer);
        n += PutUvarint(buf, n, state.Delivered.Stream);
        n += PutUvarint(buf, n, (ulong)(state.Pending?.Count ?? 0));

        var asflr = state.AckFloor.Stream;
        var adflr = state.AckFloor.Consumer;

        if (state.Pending?.Count > 0)
        {
            // Base timestamp: now rounded to the nearest second (Unix seconds).
            var mints = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
            n += PutVarint(buf, n, mints);

            foreach (var kv in state.Pending)
            {
                var sseq = kv.Key;
                var p = kv.Value;
                n += PutUvarint(buf, n, sseq - asflr);
                n += PutUvarint(buf, n, p.Sequence - adflr);
                // Downsample timestamp to seconds, store delta from mints.
                var ts = p.Timestamp / (long)TimeSpan.TicksPerSecond; // ns -> seconds? No — ns / 1e9
                // p.Timestamp is Unix nanoseconds; convert to seconds.
                ts = p.Timestamp / 1_000_000_000L;
                n += PutVarint(buf, n, mints - ts);
            }
        }

        n += PutUvarint(buf, n, (ulong)(state.Redelivered?.Count ?? 0));

        if (state.Redelivered?.Count > 0)
        {
            foreach (var kv in state.Redelivered)
            {
                n += PutUvarint(buf, n, kv.Key - asflr);
                n += PutUvarint(buf, n, kv.Value);
            }
        }

        return buf[..n];
    }

    /// <summary>
    /// Decodes consumer state from the Go-compatible binary format.
    /// Reference: golang/nats-server/server/filestore.go:12216
    /// </summary>
    public static ConsumerState Decode(ReadOnlySpan<byte> buf)
    {
        // Copy to array first so lambdas can capture without ref-type restrictions.
        return DecodeFromArray(buf.ToArray());
    }

    private static ConsumerState DecodeFromArray(byte[] buf)
    {
        if (buf.Length < 2 || buf[0] != Magic)
            throw new InvalidDataException("Corrupt consumer state: bad magic or too short.");

        var version = buf[1];
        if (version != 1 && version != 2)
            throw new InvalidDataException($"Unsupported consumer state version: {version}.");

        var bi = HdrLen;

        ulong ReadSeq()
        {
            if (bi < 0) return 0;
            var (v, n) = GetUvarint(buf.AsSpan(bi));
            if (n <= 0) { bi = -1; return 0; }
            bi += n;
            return v;
        }

        long ReadTimestamp()
        {
            if (bi < 0) return 0;
            var (v, n) = GetVarint(buf.AsSpan(bi));
            if (n <= 0) { bi = -1; return -1; }
            bi += n;
            return v;
        }

        var state = new ConsumerState();
        state.AckFloor = state.AckFloor with { Consumer = ReadSeq() };
        state.AckFloor = state.AckFloor with { Stream = ReadSeq() };
        state.Delivered = state.Delivered with { Consumer = ReadSeq() };
        state.Delivered = state.Delivered with { Stream = ReadSeq() };

        if (bi == -1)
            throw new InvalidDataException("Corrupt consumer state: truncated header.");

        // Version 1: delivered was stored as "next to be delivered", adjust back.
        if (version == 1)
        {
            if (state.AckFloor.Consumer > 1)
                state.Delivered = state.Delivered with { Consumer = state.Delivered.Consumer + state.AckFloor.Consumer - 1 };
            if (state.AckFloor.Stream > 1)
                state.Delivered = state.Delivered with { Stream = state.Delivered.Stream + state.AckFloor.Stream - 1 };
        }

        // Sanity check high-bit guard.
        if ((state.AckFloor.Stream & HighBit) != 0 || (state.Delivered.Stream & HighBit) != 0)
            throw new InvalidDataException("Corrupt consumer state: sequence high-bit set.");

        var numPending = ReadSeq();
        if (numPending > 0)
        {
            var mints = ReadTimestamp();
            state.Pending = new Dictionary<ulong, Pending>((int)numPending);

            for (var i = 0UL; i < numPending; i++)
            {
                var sseq = ReadSeq();
                ulong dseq = 0;
                if (version == 2)
                    dseq = ReadSeq();
                var ts = ReadTimestamp();

                if (bi == -1)
                    throw new InvalidDataException("Corrupt consumer state: truncated pending entry.");

                sseq += state.AckFloor.Stream;
                if (sseq == 0)
                    throw new InvalidDataException("Corrupt consumer state: zero sseq in pending.");

                if (version == 2)
                    dseq += state.AckFloor.Consumer;

                // Reconstruct timestamp (nanoseconds).
                long tsNs;
                if (version == 1)
                    tsNs = (ts + mints) * 1_000_000_000L;
                else
                    tsNs = (mints - ts) * 1_000_000_000L;

                state.Pending[sseq] = new Pending(dseq, tsNs);
            }
        }

        var numRedelivered = ReadSeq();
        if (numRedelivered > 0)
        {
            state.Redelivered = new Dictionary<ulong, ulong>((int)numRedelivered);
            for (var i = 0UL; i < numRedelivered; i++)
            {
                var seq = ReadSeq();
                var count = ReadSeq();
                if (seq > 0 && count > 0)
                    state.Redelivered[seq + state.AckFloor.Stream] = count;
            }
        }

        return state;
    }

    // -------------------------------------------------------------------------
    // Varint helpers (Go's encoding/binary LEB128 format)
    // -------------------------------------------------------------------------

    private static int PutUvarint(byte[] buf, int offset, ulong v)
    {
        var n = 0;
        while (v >= 0x80)
        {
            buf[offset + n] = (byte)(v | 0x80);
            v >>= 7;
            n++;
        }
        buf[offset + n] = (byte)v;
        return n + 1;
    }

    private static int PutVarint(byte[] buf, int offset, long v)
    {
        var uv = v < 0 ? ~((ulong)v << 1) : (ulong)v << 1;
        return PutUvarint(buf, offset, uv);
    }

    private static (ulong Value, int BytesRead) GetUvarint(ReadOnlySpan<byte> buf)
    {
        ulong x = 0;
        var s = 0;
        for (var i = 0; i < buf.Length; i++)
        {
            var b = buf[i];
            if (b < 0x80)
            {
                if (i == 9 && b > 1)
                    return (0, -(i + 1)); // overflow
                return (x | (ulong)b << s, i + 1);
            }
            x |= (ulong)(b & 0x7f) << s;
            s += 7;
        }
        return (0, 0);
    }

    private static (long Value, int BytesRead) GetVarint(ReadOnlySpan<byte> buf)
    {
        var (uv, n) = GetUvarint(buf);
        long x = (long)(uv >> 1);
        if ((uv & 1) != 0)
            x = ~x;
        return (x, n);
    }
}
