// Reference: golang/nats-server/server/avl/seqset.go
// Go uses an AVL tree with bitmask nodes (2048 sequences per node).
// .NET port uses a sorted list of (Start, End) ranges — simpler, still O(log n)
// via binary search for the common case of mostly-sequential deletions in JetStream.
//
// Range compression: adding sequences 1, 2, 3 stores as a single range [1, 3].
// This is significantly more memory efficient than HashSet<ulong> for contiguous
// deletion runs (e.g. TTL expiry, bulk removes), which are the dominant pattern.
//
// Not thread-safe — callers must hold their own lock (MsgBlock uses ReaderWriterLockSlim).

namespace NATS.Server.JetStream.Storage;

/// <summary>
/// A memory-efficient sparse set for storing unsigned sequence numbers, using
/// range compression to merge contiguous sequences into (Start, End) intervals.
///
/// Analogous to Go's <c>avl.SequenceSet</c> but implemented with a sorted list
/// of ranges for simplicity. Binary search gives O(log n) Contains/Add on the
/// number of distinct ranges (not the total count of sequences).
///
/// Reference: golang/nats-server/server/avl/seqset.go (SequenceSet struct).
/// </summary>
internal sealed class SequenceSet : IEnumerable<ulong>
{
    // Sorted list of non-overlapping, non-adjacent ranges in ascending order.
    // Invariant: for all i, ranges[i].End + 1 < ranges[i+1].Start (strict gap between consecutive ranges).
    private readonly List<(ulong Start, ulong End)> _ranges = [];

    /// <summary>Total number of sequences across all ranges.</summary>
    public int Count
    {
        get
        {
            var total = 0;
            foreach (var (start, end) in _ranges)
                total += (int)(end - start + 1);
            return total;
        }
    }

    /// <summary>True when the set contains no sequences.</summary>
    public bool IsEmpty => _ranges.Count == 0;

    /// <summary>
    /// Adds <paramref name="seq"/> to the set.
    /// Merges adjacent or overlapping ranges automatically.
    /// Returns <c>true</c> if the sequence was not already present.
    /// Reference: golang/nats-server/server/avl/seqset.go:44 (Insert).
    /// </summary>
    public bool Add(ulong seq)
    {
        // Strategy: find the position where seq belongs (binary search by Start),
        // then check if seq is already in the previous range or the range at that position,
        // and determine left/right adjacency for merging.

        // Binary search: find the first index i where _ranges[i].Start > seq.
        // Everything before i has Start <= seq. The range that might contain seq
        // is at i-1 (if it exists and its End >= seq).
        var lo = 0;
        var hi = _ranges.Count; // exclusive upper bound
        while (lo < hi)
        {
            var mid = lo + (hi - lo) / 2;
            if (_ranges[mid].Start <= seq)
                lo = mid + 1;
            else
                hi = mid;
        }
        // lo == first index where Start > seq, so lo-1 is the last range with Start <= seq.
        var leftIdx = lo - 1; // may be -1 if no range has Start <= seq

        // Check if seq is inside the range at leftIdx.
        if (leftIdx >= 0 && _ranges[leftIdx].End >= seq)
            return false; // already present

        // At this point seq is not in any range. Determine insertion context.
        // leftIdx: the range immediately to the left of seq (Start <= seq, End < seq)
        // rightIdx = lo: the range immediately to the right of seq (Start > seq)
        var rightIdx = lo;

        // Check adjacency with left neighbor (End + 1 == seq, i.e. seq extends the right edge).
        // Safe: if End == ulong.MaxValue the range would already contain seq.
        var leftAdjacent = leftIdx >= 0 && _ranges[leftIdx].End + 1 == seq;

        // Check adjacency with right neighbor (Start - 1 == seq, i.e. seq extends the left edge).
        // Need to guard against Start == 0 (underflow), but if Start == 0 then Start > seq is impossible
        // since seq >= 0 and Start <= seq would have been caught by leftIdx check.
        var rightAdjacent = rightIdx < _ranges.Count && _ranges[rightIdx].Start > 0
                             && _ranges[rightIdx].Start - 1 == seq;

        if (leftAdjacent && rightAdjacent)
        {
            // seq bridges the gap between the two neighbors → merge into one range.
            var newStart = _ranges[leftIdx].Start;
            var newEnd = _ranges[rightIdx].End;
            _ranges.RemoveAt(rightIdx);      // remove right neighbor first (higher index)
            _ranges[leftIdx] = (newStart, newEnd);
        }
        else if (leftAdjacent)
        {
            // Extend the left neighbor's right edge to include seq.
            _ranges[leftIdx] = (_ranges[leftIdx].Start, seq);
        }
        else if (rightAdjacent)
        {
            // Extend the right neighbor's left edge to include seq.
            _ranges[rightIdx] = (seq, _ranges[rightIdx].End);
        }
        else
        {
            // No adjacency — insert a new single-element range at the correct position.
            _ranges.Insert(rightIdx, (seq, seq));
        }

        return true;
    }

    /// <summary>
    /// Removes <paramref name="seq"/> from the set.
    /// Splits ranges if necessary.
    /// Returns <c>true</c> if the sequence was present.
    /// Reference: golang/nats-server/server/avl/seqset.go:80 (Delete).
    /// </summary>
    public bool Remove(ulong seq)
    {
        // Binary search for the range that contains seq.
        var lo = 0;
        var hi = _ranges.Count - 1;
        while (lo <= hi)
        {
            var mid = lo + (hi - lo) / 2;
            var (rs, re) = _ranges[mid];
            if (seq < rs)
                hi = mid - 1;
            else if (seq > re)
                lo = mid + 1;
            else
            {
                // Found the range [rs, re] that contains seq.
                if (rs == re)
                {
                    // Single-element range → remove entirely.
                    _ranges.RemoveAt(mid);
                }
                else if (seq == rs)
                {
                    // Trim left edge.
                    _ranges[mid] = (rs + 1, re);
                }
                else if (seq == re)
                {
                    // Trim right edge.
                    _ranges[mid] = (rs, re - 1);
                }
                else
                {
                    // Split: [rs, seq-1] and [seq+1, re].
                    _ranges[mid] = (seq + 1, re);
                    _ranges.Insert(mid, (rs, seq - 1));
                }
                return true;
            }
        }
        return false;
    }

    /// <summary>
    /// Returns <c>true</c> if <paramref name="seq"/> is a member of the set.
    /// Binary search: O(log R) where R is the number of distinct ranges.
    /// Reference: golang/nats-server/server/avl/seqset.go:52 (Exists).
    /// </summary>
    public bool Contains(ulong seq)
    {
        var lo = 0;
        var hi = _ranges.Count - 1;
        while (lo <= hi)
        {
            var mid = lo + (hi - lo) / 2;
            var (rs, re) = _ranges[mid];
            if (seq < rs)
                hi = mid - 1;
            else if (seq > re)
                lo = mid + 1;
            else
                return true;
        }
        return false;
    }

    /// <summary>
    /// Removes all sequences from the set.
    /// Reference: golang/nats-server/server/avl/seqset.go:107 (Empty).
    /// </summary>
    public void Clear() => _ranges.Clear();

    /// <summary>
    /// Copies all sequences in this set into a new <see cref="HashSet{T}"/>.
    /// Used when callers require an <see cref="IReadOnlySet{T}"/> snapshot.
    /// </summary>
    public HashSet<ulong> ToHashSet()
    {
        var set = new HashSet<ulong>(Count);
        foreach (var seq in this)
            set.Add(seq);
        return set;
    }

    /// <summary>
    /// Returns the number of distinct compressed ranges stored internally.
    /// </summary>
    internal int RangeCount => _ranges.Count;

    /// <summary>
    /// Enumerates all sequences in ascending order.
    /// Reference: golang/nats-server/server/avl/seqset.go:122 (Range).
    /// </summary>
    public IEnumerator<ulong> GetEnumerator()
    {
        foreach (var (start, end) in _ranges)
            for (var seq = start; seq <= end; seq++)
            {
                yield return seq;
                if (seq == ulong.MaxValue) yield break; // prevent overflow
            }
    }

    System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        => GetEnumerator();
}
