using System.Buffers.Binary;

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

/// <summary>
/// Memory and encoding-optimized set for storing unsigned 64-bit integers.
/// Implemented as an AVL tree where each node holds bitmasks for set membership.
/// Approximately 80-100x more memory-efficient than a <see cref="HashSet{T}"/>.
/// Not thread-safe.
/// </summary>
public sealed class SequenceSet
{
    private const int BitsPerBucket = 64;
    private const int NumBuckets = 32;
    internal const int NumEntries = NumBuckets * BitsPerBucket; // 2048

    private const byte MagicByte = 22;
    private const byte CurrentVersion = 2;
    private const int HdrLen = 2;
    private const int MinLen = 2 + 8; // magic + version + num_nodes(4) + num_entries(4)

    private Node? _root;
    private int _size;
    private int _nodes;
    private bool _changed;

    // --- Errors ---

    public static readonly Exception ErrBadEncoding = new InvalidDataException("ss: bad encoding");
    public static readonly Exception ErrBadVersion = new InvalidDataException("ss: bad version");
    public static readonly Exception ErrSetNotEmpty = new InvalidOperationException("ss: set not empty");

    // --- Internal access for testing ---

    internal Node? Root => _root;

    // --- Public API ---

    /// <summary>Inserts a sequence number into the set. Tree is balanced inline.</summary>
    public void Insert(ulong seq)
    {
        _root = Node.Insert(_root, seq, ref _changed, ref _nodes);
        if (_changed)
        {
            _changed = false;
            _size++;
        }
    }

    /// <summary>Returns true if the sequence is a member of the set.</summary>
    public bool Exists(ulong seq)
    {
        for (var n = _root; n != null;)
        {
            if (seq < n.Base)
            {
                n = n.Left;
            }
            else if (seq >= n.Base + NumEntries)
            {
                n = n.Right;
            }
            else
            {
                return n.ExistsBit(seq);
            }
        }
        return false;
    }

    /// <summary>
    /// Sets the initial minimum sequence when known. More effectively utilizes space.
    /// The set must be empty.
    /// </summary>
    public void SetInitialMin(ulong min)
    {
        if (!IsEmpty)
            throw (InvalidOperationException)ErrSetNotEmpty;

        _root = new Node(min);
        _nodes = 1;
    }

    /// <summary>
    /// Removes the sequence from the set. Returns true if the sequence was present.
    /// </summary>
    public bool Delete(ulong seq)
    {
        if (_root == null) return false;

        _root = Node.Delete(_root, seq, ref _changed, ref _nodes);
        if (_changed)
        {
            _changed = false;
            _size--;
            if (_size == 0)
                Empty();
            return true;
        }
        return false;
    }

    /// <summary>Returns the number of items in the set.</summary>
    public int Size => _size;

    /// <summary>Returns the number of nodes in the AVL tree.</summary>
    public int Nodes => _nodes;

    /// <summary>Clears all items from the set.</summary>
    public void Empty()
    {
        _root = null;
        _size = 0;
        _nodes = 0;
    }

    /// <summary>Returns true if the set contains no items.</summary>
    public bool IsEmpty => _root == null;

    /// <summary>
    /// Invokes the callback for each item in ascending order.
    /// Stops early if the callback returns false.
    /// </summary>
    public void Range(Func<ulong, bool> f) => Node.Iter(_root, f);

    /// <summary>Returns the heights of the left and right subtrees of the root.</summary>
    public (int Left, int Right) Heights()
    {
        if (_root == null) return (0, 0);
        return (_root.Left?.Height ?? 0, _root.Right?.Height ?? 0);
    }

    /// <summary>Returns min, max, and count of set items.</summary>
    public (ulong Min, ulong Max, ulong Count) State()
    {
        if (_root == null) return (0, 0, 0);
        var (min, max) = MinMax();
        return (min, max, (ulong)_size);
    }

    /// <summary>Returns the minimum and maximum values in the set.</summary>
    public (ulong Min, ulong Max) MinMax()
    {
        if (_root == null) return (0, 0);

        ulong min = 0;
        for (var l = _root; l != null; l = l.Left)
            if (l.Left == null) min = l.Min();

        ulong max = 0;
        for (var r = _root; r != null; r = r.Right)
            if (r.Right == null) max = r.Max();

        return (min, max);
    }

    /// <summary>Returns a deep clone of this set.</summary>
    public SequenceSet Clone()
    {
        var css = new SequenceSet { _nodes = _nodes, _size = _size };
        css._root = Node.Clone(_root);
        return css;
    }

    /// <summary>Unions one or more sequence sets into this set.</summary>
    public void Union(params SequenceSet[] sets)
    {
        foreach (var sa in sets)
        {
            Node.NodeIter(sa._root, n =>
            {
                for (var nb = 0; nb < NumBuckets; nb++)
                {
                    var b = n.Bits[nb];
                    for (var pos = 0UL; b != 0; pos++)
                    {
                        if ((b & 1) == 1)
                        {
                            var seq = n.Base + ((ulong)nb * BitsPerBucket) + pos;
                            Insert(seq);
                        }
                        b >>= 1;
                    }
                }
            });
        }
    }

    /// <summary>Returns the union of all given sets.</summary>
    public static SequenceSet? UnionSets(params SequenceSet[] sets)
    {
        if (sets.Length == 0) return null;

        // Clone the largest set first for efficiency.
        Array.Sort(sets, (a, b) => b.Size.CompareTo(a.Size));
        var ss = sets[0].Clone();
        for (var i = 1; i < sets.Length; i++)
        {
            sets[i].Range(n =>
            {
                ss.Insert(n);
                return true;
            });
        }
        return ss;
    }

    /// <summary>Returns the number of bytes needed to encode this set.</summary>
    public int EncodeLen() => MinLen + (Nodes * ((NumBuckets + 1) * 8 + 2));

    /// <summary>
    /// Encodes this set into a compact binary representation.
    /// Reuses the provided buffer if it is large enough.
    /// </summary>
    public byte[] Encode(byte[]? buf)
    {
        var nn = Nodes;
        var encLen = EncodeLen();

        if (buf == null || buf.Length < encLen)
            buf = new byte[encLen];

        buf[0] = MagicByte;
        buf[1] = CurrentVersion;

        var i = HdrLen;
        BinaryPrimitives.WriteUInt32LittleEndian(buf.AsSpan(i), (uint)nn);
        BinaryPrimitives.WriteUInt32LittleEndian(buf.AsSpan(i + 4), (uint)_size);
        i += 8;

        Node.NodeIter(_root, n =>
        {
            BinaryPrimitives.WriteUInt64LittleEndian(buf.AsSpan(i), n.Base);
            i += 8;
            foreach (var b in n.Bits)
            {
                BinaryPrimitives.WriteUInt64LittleEndian(buf.AsSpan(i), b);
                i += 8;
            }
            BinaryPrimitives.WriteUInt16LittleEndian(buf.AsSpan(i), (ushort)n.Height);
            i += 2;
        });

        return buf[..i];
    }

    /// <summary>
    /// Decodes a sequence set from the binary representation.
    /// Returns the set and the number of bytes consumed.
    /// Throws <see cref="InvalidDataException"/> on malformed input.
    /// </summary>
    public static (SequenceSet Set, int BytesRead) Decode(ReadOnlySpan<byte> buf)
    {
        if (buf.Length < MinLen || buf[0] != MagicByte)
            throw (InvalidDataException)ErrBadEncoding;

        return buf[1] switch
        {
            1 => Decodev1(buf),
            2 => Decodev2(buf),
            _ => throw (InvalidDataException)ErrBadVersion
        };
    }

    // --- Internal tree helpers ---

    /// <summary>Inserts a pre-built node directly into the tree (used during Decode).</summary>
    internal void InsertNode(Node n)
    {
        _nodes++;
        if (_root == null)
        {
            _root = n;
            return;
        }
        for (var p = _root; p != null;)
        {
            if (n.Base < p.Base)
            {
                if (p.Left == null) { p.Left = n; return; }
                p = p.Left;
            }
            else
            {
                if (p.Right == null) { p.Right = n; return; }
                p = p.Right;
            }
        }
    }

    private static (SequenceSet Set, int BytesRead) Decodev2(ReadOnlySpan<byte> buf)
    {
        var index = 2;
        var nn = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[index..]);
        var sz = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[(index + 4)..]);
        index += 8;

        var expectedLen = MinLen + (nn * ((NumBuckets + 1) * 8 + 2));
        if (buf.Length < expectedLen)
            throw (InvalidDataException)ErrBadEncoding;

        var ss = new SequenceSet { _size = sz };
        var nodes = new Node[nn];

        for (var i = 0; i < nn; i++)
        {
            var n = new Node(BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]));
            index += 8;
            for (var bi = 0; bi < NumBuckets; bi++)
            {
                n.Bits[bi] = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
                index += 8;
            }
            n.Height = (int)BinaryPrimitives.ReadUInt16LittleEndian(buf[index..]);
            index += 2;
            nodes[i] = n;
            ss.InsertNode(n);
        }

        return (ss, index);
    }

    private static (SequenceSet Set, int BytesRead) Decodev1(ReadOnlySpan<byte> buf)
    {
        const int v1NumBuckets = 64;

        var index = 2;
        var nn = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[index..]);
        var sz = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[(index + 4)..]);
        index += 8;

        var expectedLen = MinLen + (nn * ((v1NumBuckets + 1) * 8 + 2));
        if (buf.Length < expectedLen)
            throw (InvalidDataException)ErrBadEncoding;

        var ss = new SequenceSet();
        for (var i = 0; i < nn; i++)
        {
            var baseVal = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
            index += 8;
            for (var nb = 0UL; nb < v1NumBuckets; nb++)
            {
                var n = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
                for (var pos = 0UL; n != 0; pos++)
                {
                    if ((n & 1) == 1)
                    {
                        var seq = baseVal + (nb * BitsPerBucket) + pos;
                        ss.Insert(seq);
                    }
                    n >>= 1;
                }
                index += 8;
            }
            // Skip encoded height.
            index += 2;
        }

        if (ss.Size != sz)
            throw (InvalidDataException)ErrBadEncoding;

        return (ss, index);
    }

    // -------------------------------------------------------------------------
    // Internal Node class
    // -------------------------------------------------------------------------

    internal sealed class Node
    {
        public ulong Base;
        public readonly ulong[] Bits = new ulong[NumBuckets];
        public Node? Left;
        public Node? Right;
        public int Height;

        public Node(ulong baseVal)
        {
            Base = baseVal;
            Height = 1;
        }

        // Sets the bit for seq. seq must be within [Base, Base+NumEntries).
        public void SetBit(ulong seq, ref bool inserted)
        {
            var offset = seq - Base;
            var i = (int)(offset / BitsPerBucket);
            var mask = 1UL << (int)(offset % BitsPerBucket);
            if ((Bits[i] & mask) == 0)
            {
                Bits[i] |= mask;
                inserted = true;
            }
        }

        public bool ExistsBit(ulong seq)
        {
            var offset = seq - Base;
            var i = (int)(offset / BitsPerBucket);
            var mask = 1UL << (int)(offset % BitsPerBucket);
            return (Bits[i] & mask) != 0;
        }

        // Clears the bit for seq. Returns true if the node is now empty.
        public bool ClearBit(ulong seq, ref bool deleted)
        {
            var offset = seq - Base;
            var i = (int)(offset / BitsPerBucket);
            var mask = 1UL << (int)(offset % BitsPerBucket);
            if ((Bits[i] & mask) != 0)
            {
                Bits[i] &= ~mask;
                deleted = true;
            }
            foreach (var b in Bits)
                if (b != 0) return false;
            return true;
        }

        public ulong Min()
        {
            for (var i = 0; i < NumBuckets; i++)
            {
                if (Bits[i] != 0)
                    return Base + (ulong)(i * BitsPerBucket) + (ulong)System.Numerics.BitOperations.TrailingZeroCount(Bits[i]);
            }
            return 0;
        }

        public ulong Max()
        {
            for (var i = NumBuckets - 1; i >= 0; i--)
            {
                if (Bits[i] != 0)
                    return Base + (ulong)(i * BitsPerBucket) +
                           (ulong)(BitsPerBucket - System.Numerics.BitOperations.LeadingZeroCount(Bits[i] >> 1));
            }
            return 0;
        }

        // Static AVL helpers

        public static int BalanceFactor(Node? n)
        {
            if (n == null) return 0;
            return (n.Left?.Height ?? 0) - (n.Right?.Height ?? 0);
        }

        private static int MaxH(Node? n)
        {
            if (n == null) return 0;
            return Math.Max(n.Left?.Height ?? 0, n.Right?.Height ?? 0);
        }

        public static Node Insert(Node? n, ulong seq, ref bool inserted, ref int nodes)
        {
            if (n == null)
            {
                var baseVal = (seq / NumEntries) * NumEntries;
                var newNode = new Node(baseVal);
                newNode.SetBit(seq, ref inserted);
                nodes++;
                return newNode;
            }

            if (seq < n.Base)
                n.Left = Insert(n.Left, seq, ref inserted, ref nodes);
            else if (seq >= n.Base + NumEntries)
                n.Right = Insert(n.Right, seq, ref inserted, ref nodes);
            else
                n.SetBit(seq, ref inserted);

            n.Height = MaxH(n) + 1;

            var bf = BalanceFactor(n);
            if (bf > 1)
            {
                if (BalanceFactor(n.Left) < 0)
                    n.Left = n.Left!.RotateLeft();
                return n.RotateRight();
            }
            if (bf < -1)
            {
                if (BalanceFactor(n.Right) > 0)
                    n.Right = n.Right!.RotateRight();
                return n.RotateLeft();
            }
            return n;
        }

        public static Node? Delete(Node? n, ulong seq, ref bool deleted, ref int nodes)
        {
            if (n == null) return null;

            if (seq < n.Base)
                n.Left = Delete(n.Left, seq, ref deleted, ref nodes);
            else if (seq >= n.Base + NumEntries)
                n.Right = Delete(n.Right, seq, ref deleted, ref nodes);
            else if (n.ClearBit(seq, ref deleted))
            {
                nodes--;
                if (n.Left == null)
                    n = n.Right;
                else if (n.Right == null)
                    n = n.Left;
                else
                {
                    n.Right = n.Right.InsertNodePrev(n.Left);
                    n = n.Right;
                }
            }

            if (n == null) return null;

            n.Height = MaxH(n) + 1;

            var bf = BalanceFactor(n);
            if (bf > 1)
            {
                if (BalanceFactor(n.Left) < 0)
                    n.Left = n.Left!.RotateLeft();
                return n.RotateRight();
            }
            if (bf < -1)
            {
                if (BalanceFactor(n.Right) > 0)
                    n.Right = n.Right!.RotateRight();
                return n.RotateLeft();
            }
            return n;
        }

        private Node RotateLeft()
        {
            var r = Right;
            if (r != null)
            {
                Right = r.Left;
                r.Left = this;
                Height = MaxH(this) + 1;
                r.Height = MaxH(r) + 1;
            }
            else
            {
                Right = null;
                Height = MaxH(this) + 1;
            }
            return r!;
        }

        private Node RotateRight()
        {
            var l = Left;
            if (l != null)
            {
                Left = l.Right;
                l.Right = this;
                Height = MaxH(this) + 1;
                l.Height = MaxH(l) + 1;
            }
            else
            {
                Left = null;
                Height = MaxH(this) + 1;
            }
            return l!;
        }

        // Inserts nn into this subtree assuming nn.Base < all nodes in this subtree.
        public Node InsertNodePrev(Node nn)
        {
            if (Left == null)
                Left = nn;
            else
                Left = Left.InsertNodePrev(nn);

            Height = MaxH(this) + 1;

            var bf = BalanceFactor(this);
            if (bf > 1)
            {
                if (BalanceFactor(Left) < 0)
                    Left = Left!.RotateLeft();
                return RotateRight();
            }
            if (bf < -1)
            {
                if (BalanceFactor(Right) > 0)
                    Right = Right!.RotateRight();
                return RotateLeft();
            }
            return this;
        }

        // Iterates nodes in tree order (pre-order: root → left → right).
        public static void NodeIter(Node? n, Action<Node> f)
        {
            if (n == null) return;
            f(n);
            NodeIter(n.Left, f);
            NodeIter(n.Right, f);
        }

        // Iterates items in ascending order (in-order traversal).
        // Returns false if the callback returns false.
        public static bool Iter(Node? n, Func<ulong, bool> f)
        {
            if (n == null) return true;

            if (!Iter(n.Left, f)) return false;

            for (var num = n.Base; num < n.Base + NumEntries; num++)
            {
                if (n.ExistsBit(num))
                    if (!f(num)) return false;
            }

            return Iter(n.Right, f);
        }

        public static Node? Clone(Node? src)
        {
            if (src == null) return null;
            var n = new Node(src.Base) { Height = src.Height };
            src.Bits.CopyTo(n.Bits, 0);
            n.Left = Clone(src.Left);
            n.Right = Clone(src.Right);
            return n;
        }
    }
}