using System.Buffers.Binary;

namespace ZB.MOM.NatsNet.Server.Internal.DataStructures;

/// <summary>
/// A time-based hash wheel for efficiently scheduling and expiring timer tasks keyed by sequence number.
/// Each slot covers a 1-second window; the wheel has 4096 slots (covering ~68 minutes before wrapping).
/// Not thread-safe.
/// </summary>
/// <remarks>
/// Mirrors the Go <c>thw.HashWheel</c> type. Timestamps are Unix nanoseconds (<see cref="long"/>).
/// </remarks>
public sealed class HashWheel
{
    /// <summary>Slot width in nanoseconds (1 second).</summary>
    private const long TickDuration = 1_000_000_000L;
    private const int WheelBits = 12;
    private const int WheelSize = 1 << WheelBits; // 4096
    private const int WheelMask = WheelSize - 1;
    private const int HeaderLen = 17; // 1 magic + 8 count + 8 highSeq

    public static readonly Exception ErrTaskNotFound = new InvalidOperationException("thw: task not found");
    public static readonly Exception ErrInvalidVersion = new InvalidDataException("thw: encoded version not known");

    private readonly Slot?[] _wheel = new Slot?[WheelSize];
    private long _lowest = long.MaxValue;
    private ulong _count;

    // --- Slot ---

    private sealed class Slot
    {
        public readonly Dictionary<ulong, long> Entries = new();
        public long Lowest = long.MaxValue;
    }

    /// <summary>Creates a new empty <see cref="HashWheel"/>.</summary>
    public static HashWheel NewHashWheel() => new();

    private static Slot NewSlot() => new();

    private long GetPosition(long expires) => (expires / TickDuration) & WheelMask;

    // --- Public API ---

    /// <summary>Returns the number of tasks currently scheduled.</summary>
    public ulong Count => _count;

    /// <summary>Schedules a new timer task.</summary>
    public void Add(ulong seq, long expires)
    {
        var pos = (int)GetPosition(expires);
        _wheel[pos] ??= NewSlot();
        var slot = _wheel[pos]!;

        if (!slot.Entries.ContainsKey(seq))
            _count++;
        slot.Entries[seq] = expires;

        if (expires < slot.Lowest)
        {
            slot.Lowest = expires;
            if (expires < _lowest)
                _lowest = expires;
        }
    }

    /// <summary>Removes a timer task.</summary>
    /// <exception cref="InvalidOperationException">Thrown (as <see cref="ErrTaskNotFound"/>) when not found.</exception>
    public void Remove(ulong seq, long expires)
    {
        var pos = (int)GetPosition(expires);
        var slot = _wheel[pos];
        if (slot is null || !slot.Entries.ContainsKey(seq))
            throw ErrTaskNotFound;

        slot.Entries.Remove(seq);
        _count--;
        if (slot.Entries.Count == 0)
            _wheel[pos] = null;
    }

    /// <summary>Updates the expiration time of an existing timer task.</summary>
    public void Update(ulong seq, long oldExpires, long newExpires)
    {
        Remove(seq, oldExpires);
        Add(seq, newExpires);
    }

    /// <summary>
    /// Expires all tasks whose timestamp is &lt;= now. The callback receives each task;
    /// if it returns <see langword="true"/> the task is removed, otherwise it is kept.
    /// </summary>
    public void ExpireTasks(Func<ulong, long, bool> callback)
    {
        var now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000_000L;
        ExpireTasksInternal(now, callback);
    }

    internal void ExpireTasksInternal(long ts, Func<ulong, long, bool> callback)
    {
        if (_lowest > ts)
            return;

        var globalLowest = long.MaxValue;
        for (var pos = 0; pos < WheelSize; pos++)
        {
            var slot = _wheel[pos];
            if (slot is null || slot.Lowest > ts)
            {
                if (slot is not null && slot.Lowest < globalLowest)
                    globalLowest = slot.Lowest;
                continue;
            }

            var slotLowest = long.MaxValue;
            // Snapshot keys to allow removal during iteration.
            var keys = slot.Entries.Keys.ToArray();
            foreach (var seq in keys)
            {
                var exp = slot.Entries[seq];
                if (exp <= ts && callback(seq, exp))
                {
                    slot.Entries.Remove(seq);
                    _count--;
                    continue;
                }
                if (exp < slotLowest)
                    slotLowest = exp;
            }

            if (slot.Entries.Count == 0)
            {
                _wheel[pos] = null;
            }
            else
            {
                slot.Lowest = slotLowest;
                if (slotLowest < globalLowest)
                    globalLowest = slotLowest;
            }
        }
        _lowest = globalLowest;
    }

    /// <summary>
    /// Returns the earliest expiration timestamp before <paramref name="before"/>,
    /// or <see cref="long.MaxValue"/> if none.
    /// </summary>
    public long GetNextExpiration(long before) =>
        _lowest < before ? _lowest : long.MaxValue;

    // --- Encode / Decode ---

    /// <summary>
    /// Serializes the wheel to a byte array. <paramref name="highSeq"/> is stored
    /// in the header and returned by <see cref="Decode"/>.
    /// </summary>
    public byte[] Encode(ulong highSeq)
    {
        // Preallocate conservatively: header + up to 2 varints per entry.
        var buf = new List<byte>(HeaderLen + (int)(_count * 16));
        buf.Add(1); // magic version
        AppendUint64LE(buf, _count);
        AppendUint64LE(buf, highSeq);

        foreach (var slot in _wheel)
        {
            if (slot is null)
                continue;
            foreach (var (seq, ts) in slot.Entries)
            {
                AppendVarint(buf, ts);
                AppendUvarint(buf, seq);
            }
        }
        return buf.ToArray();
    }

    /// <summary>
    /// Replaces this wheel's contents with those from a binary snapshot.
    /// Returns the <c>highSeq</c> stored in the header.
    /// </summary>
    public ulong Decode(ReadOnlySpan<byte> b)
    {
        if (b.Length < HeaderLen)
            throw (InvalidDataException)ErrInvalidVersion;
        if (b[0] != 1)
            throw (InvalidDataException)ErrInvalidVersion;

        // Reset wheel.
        Array.Clear(_wheel);
        _lowest = long.MaxValue;
        _count = 0;

        var count = BinaryPrimitives.ReadUInt64LittleEndian(b[1..]);
        var stamp = BinaryPrimitives.ReadUInt64LittleEndian(b[9..]);
        var pos = HeaderLen;

        for (ulong i = 0; i < count; i++)
        {
            var ts = ReadVarint(b, ref pos);
            var seq = ReadUvarint(b, ref pos);
            Add(seq, ts);
        }
        return stamp;
    }

    // --- Encoding helpers ---

    private static void AppendUint64LE(List<byte> buf, ulong v)
    {
        buf.Add((byte)v);
        buf.Add((byte)(v >> 8));
        buf.Add((byte)(v >> 16));
        buf.Add((byte)(v >> 24));
        buf.Add((byte)(v >> 32));
        buf.Add((byte)(v >> 40));
        buf.Add((byte)(v >> 48));
        buf.Add((byte)(v >> 56));
    }

    private static void AppendVarint(List<byte> buf, long v)
    {
        // ZigZag encode like Go's binary.AppendVarint.
        var uv = (ulong)((v << 1) ^ (v >> 63));
        AppendUvarint(buf, uv);
    }

    private static void AppendUvarint(List<byte> buf, ulong v)
    {
        while (v >= 0x80)
        {
            buf.Add((byte)(v | 0x80));
            v >>= 7;
        }
        buf.Add((byte)v);
    }

    private static long ReadVarint(ReadOnlySpan<byte> b, ref int pos)
    {
        var uv = ReadUvarint(b, ref pos);
        var v = (long)(uv >> 1);
        if ((uv & 1) != 0)
            v = ~v;
        return v;
    }

    private static ulong ReadUvarint(ReadOnlySpan<byte> b, ref int pos)
    {
        ulong x = 0;
        int s = 0;
        while (pos < b.Length)
        {
            var by = b[pos++];
            x |= (ulong)(by & 0x7F) << s;
            if ((by & 0x80) == 0)
                return x;
            s += 7;
        }
        throw new InvalidDataException("thw: unexpected EOF in varint");
    }
}
