// 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.

// Go reference: server/avl/seqset.go

using System.Buffers.Binary;
using System.Numerics;

namespace NATS.Server.Internal.Avl;

/// <summary>
/// SequenceSet is a memory and encoding optimized set for storing unsigned ints.
/// Uses an AVL tree with nodes that hold bitmasks for set membership.
/// Not thread safe.
/// </summary>
public class SequenceSet
{
    internal const int BitsPerBucket = 64;
    internal const int NumBuckets = 32;
    internal const int NumEntries = NumBuckets * BitsPerBucket; // 2048

    private const byte Magic = 22;
    private const byte Version = 2;
    private const int HdrLen = 2;
    private const int MinLen = HdrLen + 8; // magic + version + num nodes + num entries

    internal Node? Root;
    private int _size;
    private int _nodes;
    private bool _changed;

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

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

    /// <summary>Fast check of the set being empty.</summary>
    public bool IsEmpty => Root == null;

    /// <summary>Insert will insert the sequence into the set. The tree will be 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 this set.</summary>
    public bool Exists(ulong seq)
    {
        var n = Root;
        while (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 new InvalidOperationException("Set not empty");
        }

        Root = new Node { Base = min, Height = 1 };
        _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>Clears all items from the set.</summary>
    public void Empty()
    {
        Root = null;
        _size = 0;
        _nodes = 0;
    }

    /// <summary>
    /// Invokes the callback for each item in ascending order.
    /// If the callback returns false, iteration terminates.
    /// </summary>
    public void Range(Func<ulong, bool> callback) => Node.Iter(Root, callback);

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

        var l = Root.Left?.Height ?? 0;
        var r = Root.Right?.Height ?? 0;
        return (l, r);
    }

    /// <summary>Returns min, max, and number of set items.</summary>
    public (ulong Min, ulong Max, ulong Num) 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 SequenceSet.</summary>
    public SequenceSet Clone()
    {
        var css = new SequenceSet { _nodes = _nodes, _size = _size };
        css.Root = CloneNode(Root);
        return css;
    }

    /// <summary>Unions this set with one or more other sets by inserting all their elements.</summary>
    public void Union(params SequenceSet[] others)
    {
        foreach (var other in others)
        {
            Node.NodeIter(other.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 a union of all provided sets.</summary>
    public static SequenceSet CreateUnion(params SequenceSet[] sets)
    {
        if (sets.Length == 0)
        {
            return new SequenceSet();
        }

        // Sort descending by size so we clone the largest.
        var sorted = sets.OrderByDescending(s => s.Size).ToArray();
        var ss = sorted[0].Clone();

        for (var i = 1; i < sorted.Length; i++)
        {
            sorted[i].Range(n =>
            {
                ss.Insert(n);
                return true;
            });
        }

        return ss;
    }

    /// <summary>Returns the bytes needed for encoding.</summary>
    public int EncodeLength() => MinLen + (_nodes * ((NumBuckets + 1) * 8 + 2));

    /// <summary>Encodes the set to a compact binary format.</summary>
    public byte[] Encode()
    {
        var buf = new byte[EncodeLength()];
        var written = Encode(buf);
        return buf.AsSpan(0, written).ToArray();
    }

    /// <summary>
    /// Encodes the set into a caller-provided buffer.
    /// Returns the number of bytes written.
    /// </summary>
    public int Encode(byte[] destination)
    {
        var encLen = EncodeLength();
        if (destination.Length < encLen)
            throw new ArgumentException("Destination buffer too small for encoded SequenceSet.", nameof(destination));

        destination[0] = Magic;
        destination[1] = Version;
        var i = HdrLen;
        BinaryPrimitives.WriteUInt32LittleEndian(destination.AsSpan(i), (uint)_nodes);
        BinaryPrimitives.WriteUInt32LittleEndian(destination.AsSpan(i + 4), (uint)_size);
        i += 8;

        Node.NodeIter(Root, n =>
        {
            BinaryPrimitives.WriteUInt64LittleEndian(destination.AsSpan(i), n.Base);
            i += 8;
            for (var bi = 0; bi < NumBuckets; bi++)
            {
                BinaryPrimitives.WriteUInt64LittleEndian(destination.AsSpan(i), n.Bits[bi]);
                i += 8;
            }

            BinaryPrimitives.WriteUInt16LittleEndian(destination.AsSpan(i), (ushort)n.Height);
            i += 2;
        });

        return i;
    }

    /// <summary>Decodes a SequenceSet from a binary buffer. Returns the set and number of bytes read.</summary>
    public static (SequenceSet Set, int BytesRead) Decode(ReadOnlySpan<byte> buf)
    {
        if (buf.Length < MinLen || buf[0] != Magic)
        {
            throw new InvalidOperationException("Bad encoding");
        }

        return buf[1] switch
        {
            1 => DecodeV1(buf),
            2 => DecodeV2(buf),
            _ => throw new InvalidOperationException("Bad version"),
        };
    }

    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 new InvalidOperationException("Bad encoding");
        }

        var ss = new SequenceSet { _size = sz };

        for (var i = 0; i < nn; i++)
        {
            var n = new Node
            {
                Base = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]),
            };
            index += 8;

            for (var bi = 0; bi < NumBuckets; bi++)
            {
                n.Bits[bi] = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
                index += 8;
            }

            n.Height = BinaryPrimitives.ReadUInt16LittleEndian(buf[index..]);
            index += 2;

            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 new InvalidOperationException("Bad encoding");
        }

        var ss = new SequenceSet();

        for (var i = 0; i < nn; i++)
        {
            var nodeBase = 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 = nodeBase + (nb * BitsPerBucket) + pos;
                        ss.Insert(seq);
                    }

                    n >>= 1;
                }

                index += 8;
            }

            // Skip encoded height.
            index += 2;
        }

        if (ss.Size != sz)
        {
            throw new InvalidOperationException("Bad encoding");
        }

        return (ss, index);
    }

    /// <summary>Inserts a decoded node directly into the tree (no rebalancing needed for ordered inserts).</summary>
    private void InsertNode(Node n)
    {
        _nodes++;

        if (Root == null)
        {
            Root = n;
            return;
        }

        for (var p = Root; ;)
        {
            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 Node? CloneNode(Node? src)
    {
        if (src == null)
        {
            return null;
        }

        var n = new Node { Base = src.Base, Height = src.Height };
        Array.Copy(src.Bits, n.Bits, NumBuckets);
        n.Left = CloneNode(src.Left);
        n.Right = CloneNode(src.Right);
        return n;
    }

    /// <summary>AVL tree node that stores a bitmask covering NumEntries (2048) consecutive sequences.</summary>
    internal sealed class Node
    {
        public ulong Base;
        public readonly ulong[] Bits = new ulong[NumBuckets];
        public Node? Left;
        public Node? Right;
        public int Height;

        /// <summary>Sets the bit for the given sequence. Reports whether it was newly inserted.</summary>
        public void SetBit(ulong seq, ref bool inserted)
        {
            seq -= Base;
            var i = seq / BitsPerBucket;
            var mask = 1UL << (int)(seq % BitsPerBucket);
            if ((Bits[i] & mask) == 0)
            {
                Bits[i] |= mask;
                inserted = true;
            }
        }

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

            for (var b = 0; b < NumBuckets; b++)
            {
                if (Bits[b] != 0)
                {
                    return false;
                }
            }

            return true;
        }

        /// <summary>Checks if the bit for the given sequence is set.</summary>
        public bool ExistsBit(ulong seq)
        {
            seq -= Base;
            var i = seq / BitsPerBucket;
            var mask = 1UL << (int)(seq % BitsPerBucket);
            return (Bits[i] & mask) != 0;
        }

        /// <summary>Returns the minimum sequence in this node (node must not be empty).</summary>
        public ulong Min()
        {
            for (var i = 0; i < NumBuckets; i++)
            {
                if (Bits[i] != 0)
                {
                    return Base + (ulong)(i * BitsPerBucket) + (ulong)BitOperations.TrailingZeroCount(Bits[i]);
                }
            }

            return 0;
        }

        /// <summary>Returns the maximum sequence in this node (node must not be empty).</summary>
        public ulong Max()
        {
            for (var i = NumBuckets - 1; i >= 0; i--)
            {
                if (Bits[i] != 0)
                {
                    return Base + (ulong)(i * BitsPerBucket) + (ulong)(BitsPerBucket - BitOperations.LeadingZeroCount(Bits[i] >> 1));
                }
            }

            return 0;
        }

        /// <summary>Inserts a sequence into the subtree rooted at this node, rebalancing as needed.</summary>
        public static Node Insert(Node? n, ulong seq, ref bool inserted, ref int nodes)
        {
            if (n == null)
            {
                var nodeBase = (seq / NumEntries) * NumEntries;
                var newNode = new Node { Base = nodeBase, Height = 1 };
                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 = MaxHeight(n) + 1;

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

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

                return RotateLeft(n);
            }

            return n;
        }

        /// <summary>Deletes a sequence from the subtree rooted at this node, rebalancing as needed.</summary>
        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))
            {
                // Node is now empty, remove it.
                nodes--;
                if (n.Left == null)
                {
                    n = n.Right;
                }
                else if (n.Right == null)
                {
                    n = n.Left;
                }
                else
                {
                    // Both children present: insert left subtree into the leftmost position of right subtree.
                    n.Right = InsertNodePrev(n.Right, n.Left);
                    n = n.Right;
                }
            }

            if (n != null)
            {
                n.Height = MaxHeight(n) + 1;
            }

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

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

                return RotateLeft(n);
            }

            return n;
        }

        /// <summary>Inserts nn into the leftmost position of n's subtree, then rebalances.</summary>
        private static Node InsertNodePrev(Node n, Node nn)
        {
            if (n.Left == null)
            {
                n.Left = nn;
            }
            else
            {
                n.Left = InsertNodePrev(n.Left, nn);
            }

            n.Height = MaxHeight(n) + 1;

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

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

                return RotateLeft(n);
            }

            return n;
        }

        /// <summary>Left rotation.</summary>
        private static Node RotateLeft(Node n)
        {
            var r = n.Right;
            if (r != null)
            {
                n.Right = r.Left;
                r.Left = n;
                n.Height = MaxHeight(n) + 1;
                r.Height = MaxHeight(r) + 1;
            }
            else
            {
                n.Right = null;
                n.Height = MaxHeight(n) + 1;
            }

            return r ?? n;
        }

        /// <summary>Right rotation.</summary>
        private static Node RotateRight(Node n)
        {
            var l = n.Left;
            if (l != null)
            {
                n.Left = l.Right;
                l.Right = n;
                n.Height = MaxHeight(n) + 1;
                l.Height = MaxHeight(l) + 1;
            }
            else
            {
                n.Left = null;
                n.Height = MaxHeight(n) + 1;
            }

            return l ?? n;
        }

        /// <summary>Returns the balance factor (left height - right height).</summary>
        internal static int BalanceFactor(Node? n)
        {
            if (n == null)
            {
                return 0;
            }

            var lh = n.Left?.Height ?? 0;
            var rh = n.Right?.Height ?? 0;
            return lh - rh;
        }

        /// <summary>Returns the max of left and right child heights.</summary>
        internal static int MaxHeight(Node? n)
        {
            if (n == null)
            {
                return 0;
            }

            var lh = n.Left?.Height ?? 0;
            var rh = n.Right?.Height ?? 0;
            return Math.Max(lh, rh);
        }

        /// <summary>Iterates nodes in pre-order (root, left, right) for encoding.</summary>
        internal static void NodeIter(Node? n, Action<Node> f)
        {
            if (n == null)
            {
                return;
            }

            f(n);
            NodeIter(n.Left, f);
            NodeIter(n.Right, f);
        }

        /// <summary>Iterates items in ascending order. Returns false if iteration was terminated early.</summary>
        internal 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);
        }
    }
}
