// Copyright 2023-2025 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.

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

/// <summary>
/// An adaptive radix trie (ART) for storing subject information on literal NATS subjects.
/// Uses dynamic nodes (4/10/16/48/256 children), path compression, and lazy expansion.
/// Supports exact lookup, wildcard matching ('*' and '>'), and ordered/fast iteration.
/// Not thread-safe.
/// </summary>
public sealed class SubjectTree<T>
{
    internal ISubjectTreeNode<T>? _root;
    private int _size;

    /// <summary>Returns the number of entries stored in the tree.</summary>
    public int Size() => _size;

    /// <summary>Returns true if the tree has no entries.</summary>
    public bool Empty() => _size == 0;

    /// <summary>Clears all entries from the tree.</summary>
    public SubjectTree<T> Reset()
    {
        _root = null;
        _size = 0;
        return this;
    }

    /// <summary>
    /// Inserts a value into the tree under the given subject.
    /// If the subject already exists, returns the old value with updated=true.
    /// Subjects containing byte 127 (the noPivot sentinel) are rejected silently.
    /// </summary>
    public (T? oldVal, bool updated) Insert(ReadOnlySpan<byte> subject, T value)
    {
        if (subject.IndexOf(SubjectTreeParts.NoPivot) >= 0)
            return (default, false);

        var subjectBytes = subject.ToArray();
        var (old, updated) = DoInsert(ref _root, subjectBytes, value, 0);
        if (!updated)
            _size++;
        return (old, updated);
    }

    /// <summary>
    /// Finds the value stored at the given literal subject.
    /// Returns (value, true) if found, (default, false) otherwise.
    /// </summary>
    public (T? val, bool found) Find(ReadOnlySpan<byte> subject)
    {
        var si = 0;
        var n = _root;
        var subjectBytes = subject.ToArray();

        while (n != null)
        {
            if (n.IsLeaf)
            {
                var ln = (SubjectTreeLeaf<T>)n;
                return ln.Match(subjectBytes.AsSpan(si))
                    ? (ln.Value, true)
                    : (default, false);
            }

            var prefix = n.Prefix;
            if (prefix.Length > 0)
            {
                var end = Math.Min(si + prefix.Length, subjectBytes.Length);
                if (!subjectBytes.AsSpan(si, end - si).SequenceEqual(prefix.AsSpan(0, end - si)))
                    return (default, false);
                si += prefix.Length;
            }

            var next = n.FindChild(SubjectTreeParts.Pivot(subjectBytes, si));
            if (next == null) return (default, false);
            n = next;
        }

        return (default, false);
    }

    /// <summary>
    /// Deletes the entry at the given literal subject.
    /// Returns (value, true) if deleted, (default, false) if not found.
    /// </summary>
    public (T? val, bool found) Delete(ReadOnlySpan<byte> subject)
    {
        if (_root == null || subject.IsEmpty) return (default, false);

        var subjectBytes = subject.ToArray();
        var (val, deleted) = DoDelete(ref _root, subjectBytes, 0);
        if (deleted) _size--;
        return (val, deleted);
    }

    /// <summary>
    /// Matches all stored subjects against a filter that may contain wildcards ('*' and '>').
    /// Invokes fn for each match. Return false from the callback to stop early.
    /// </summary>
    public void Match(ReadOnlySpan<byte> filter, Func<byte[], T, bool> fn)
    {
        if (_root == null || filter.IsEmpty || fn == null) return;
        var parts = SubjectTreeParts.GenParts(filter.ToArray());
        MatchNode(_root, parts, Array.Empty<byte>(), fn);
    }

    /// <summary>
    /// Like Match but returns false if the callback stopped iteration early.
    /// Returns true if matching ran to completion.
    /// </summary>
    public bool MatchUntil(ReadOnlySpan<byte> filter, Func<byte[], T, bool> fn)
    {
        if (_root == null || filter.IsEmpty || fn == null) return true;
        var parts = SubjectTreeParts.GenParts(filter.ToArray());
        return MatchNode(_root, parts, Array.Empty<byte>(), fn);
    }

    /// <summary>
    /// Walks all entries in lexicographical order.
    /// Return false from the callback to stop early.
    /// </summary>
    public void IterOrdered(Func<byte[], T, bool> fn)
    {
        if (_root == null || fn == null) return;
        IterNode(_root, Array.Empty<byte>(), ordered: true, fn);
    }

    /// <summary>
    /// Walks all entries in storage order (no ordering guarantee).
    /// Return false from the callback to stop early.
    /// </summary>
    public void IterFast(Func<byte[], T, bool> fn)
    {
        if (_root == null || fn == null) return;
        IterNode(_root, Array.Empty<byte>(), ordered: false, fn);
    }

    // -------------------------------------------------------------------------
    // Internal recursive insert
    // -------------------------------------------------------------------------

    private static (T? old, bool updated) DoInsert(ref ISubjectTreeNode<T>? np, byte[] subject, T value, int si)
    {
        if (np == null)
        {
            np = new SubjectTreeLeaf<T>(subject, value);
            return (default, false);
        }

        if (np.IsLeaf)
        {
            var ln = (SubjectTreeLeaf<T>)np;
            if (ln.Match(subject.AsSpan(si)))
            {
                var oldVal = ln.Value;
                ln.Value = value;
                return (oldVal, true);
            }

            // Split the leaf: compute common prefix between existing suffix and new subject tail.
            var cpi = SubjectTreeParts.CommonPrefixLen(ln.Suffix, subject.AsSpan(si));
            var nn = new SubjectTreeNode4<T>(subject[si..(si + cpi)]);
            ln.SetSuffix(ln.Suffix[cpi..]);
            si += cpi;

            var p = SubjectTreeParts.Pivot(ln.Suffix, 0);
            if (cpi > 0 && si < subject.Length && p == subject[si])
            {
                // Same pivot after the split — recurse to separate further.
                DoInsert(ref np, subject, value, si);
                nn.AddChild(p, np!);
            }
            else
            {
                var nl = new SubjectTreeLeaf<T>(subject[si..], value);
                nn.AddChild(SubjectTreeParts.Pivot(nl.Suffix, 0), nl);
                nn.AddChild(SubjectTreeParts.Pivot(ln.Suffix, 0), ln);
            }

            np = nn;
            return (default, false);
        }

        // Non-leaf node.
        var prefix = np.Prefix;
        if (prefix.Length > 0)
        {
            var cpi = SubjectTreeParts.CommonPrefixLen(prefix, subject.AsSpan(si));
            if (cpi >= prefix.Length)
            {
                // Full prefix match: move past this node.
                si += prefix.Length;
                var pivotByte = SubjectTreeParts.Pivot(subject, si);
                var existingChild = np.FindChild(pivotByte);
                if (existingChild != null)
                {
                    var before = existingChild;
                    var (old, upd) = DoInsert(ref existingChild, subject, value, si);
                    // Only re-register if the child reference changed identity (grew or split).
                    if (!ReferenceEquals(before, existingChild))
                    {
                        np.DeleteChild(pivotByte);
                        np.AddChild(pivotByte, existingChild!);
                    }
                    return (old, upd);
                }

                if (np.IsFull)
                    np = np.Grow();

                np.AddChild(SubjectTreeParts.Pivot(subject, si), new SubjectTreeLeaf<T>(subject[si..], value));
                return (default, false);
            }
            else
            {
                // Partial prefix match — insert a new node4 above the current node.
                var newPrefix = subject[si..(si + cpi)];
                si += cpi;
                var splitNode = new SubjectTreeNode4<T>(newPrefix);

                ((SubjectTreeMeta<T>)np).SetPrefix(prefix[cpi..]);
                // Use np.Prefix (updated) to get the correct pivot for the demoted node.
                splitNode.AddChild(SubjectTreeParts.Pivot(np.Prefix, 0), np);
                splitNode.AddChild(
                    SubjectTreeParts.Pivot(subject.AsSpan(si), 0),
                    new SubjectTreeLeaf<T>(subject[si..], value));
                np = splitNode;
            }
        }
        else
        {
            // No prefix on this node.
            var pivotByte = SubjectTreeParts.Pivot(subject, si);
            var existingChild = np.FindChild(pivotByte);
            if (existingChild != null)
            {
                var before = existingChild;
                var (old, upd) = DoInsert(ref existingChild, subject, value, si);
                if (!ReferenceEquals(before, existingChild))
                {
                    np.DeleteChild(pivotByte);
                    np.AddChild(pivotByte, existingChild!);
                }
                return (old, upd);
            }

            if (np.IsFull)
                np = np.Grow();

            np.AddChild(SubjectTreeParts.Pivot(subject, si), new SubjectTreeLeaf<T>(subject[si..], value));
        }

        return (default, false);
    }

    // -------------------------------------------------------------------------
    // Internal recursive delete
    // -------------------------------------------------------------------------

    private static (T? val, bool deleted) DoDelete(ref ISubjectTreeNode<T>? np, byte[] subject, int si)
    {
        if (np == null || subject.Length == 0) return (default, false);

        var n = np;
        if (n.IsLeaf)
        {
            var ln = (SubjectTreeLeaf<T>)n;
            if (ln.Match(subject.AsSpan(si)))
            {
                np = null;
                return (ln.Value, true);
            }
            return (default, false);
        }

        // Check prefix.
        var prefix = n.Prefix;
        if (prefix.Length > 0)
        {
            if (subject.Length < si + prefix.Length)
                return (default, false);
            if (!subject.AsSpan(si, prefix.Length).SequenceEqual(prefix))
                return (default, false);
            si += prefix.Length;
        }

        var p = SubjectTreeParts.Pivot(subject, si);
        var childNode = n.FindChild(p);
        if (childNode == null) return (default, false);

        if (childNode.IsLeaf)
        {
            var childLeaf = (SubjectTreeLeaf<T>)childNode;
            if (childLeaf.Match(subject.AsSpan(si)))
            {
                n.DeleteChild(p);
                TryShrink(ref np!, prefix);
                return (childLeaf.Value, true);
            }
            return (default, false);
        }

        // Recurse into non-leaf child.
        var (val, deleted) = DoDelete(ref childNode, subject, si);
        if (deleted)
        {
            if (childNode == null)
            {
                // Child was nulled out — remove slot and try to shrink.
                n.DeleteChild(p);
                TryShrink(ref np!, prefix);
            }
            else
            {
                // Child changed identity — re-register.
                n.DeleteChild(p);
                n.AddChild(p, childNode);
            }
        }
        return (val, deleted);
    }

    private static void TryShrink(ref ISubjectTreeNode<T> np, byte[] parentPrefix)
    {
        var shrunk = np.Shrink();
        if (shrunk == null) return;

        if (shrunk.IsLeaf)
        {
            var shrunkLeaf = (SubjectTreeLeaf<T>)shrunk;
            if (parentPrefix.Length > 0)
                shrunkLeaf.Suffix = [.. parentPrefix, .. shrunkLeaf.Suffix];
        }
        else if (parentPrefix.Length > 0)
        {
            ((SubjectTreeMeta<T>)shrunk).SetPrefix([.. parentPrefix, .. shrunk.Prefix]);
        }
        np = shrunk;
    }

    // -------------------------------------------------------------------------
    // Internal recursive wildcard match
    // -------------------------------------------------------------------------

    private static bool MatchNode(ISubjectTreeNode<T> n, byte[][] parts, byte[] pre, Func<byte[], T, bool> fn)
    {
        var hasFwc = parts.Length > 0 && parts[^1].Length == 1 && parts[^1][0] == SubjectTreeParts.Fwc;

        while (n != null!)
        {
            var (nparts, matched) = n.MatchParts(parts);
            if (!matched) return true;

            if (n.IsLeaf)
            {
                if (nparts.Length == 0 || (hasFwc && nparts.Length == 1))
                {
                    var ln = (SubjectTreeLeaf<T>)n;
                    if (!fn(ConcatBytes(pre, ln.Suffix), ln.Value)) return false;
                }
                return true;
            }

            // Append this node's prefix to the running accumulator.
            var prefix = n.Prefix;
            if (prefix.Length > 0)
                pre = ConcatBytes(pre, prefix);

            if (nparts.Length == 0 && !hasFwc)
            {
                // No parts remaining and no fwc — look for terminal matches.
                var hasTermPwc = parts.Length > 0 && parts[^1].Length == 1 && parts[^1][0] == SubjectTreeParts.Pwc;
                var termParts = hasTermPwc ? parts[^1..] : Array.Empty<byte[]>();

                foreach (var cn in n.Children())
                {
                    if (cn == null!) continue;
                    if (cn.IsLeaf)
                    {
                        var ln = (SubjectTreeLeaf<T>)cn;
                        if (ln.Suffix.Length == 0)
                        {
                            if (!fn(ConcatBytes(pre, ln.Suffix), ln.Value)) return false;
                        }
                        else if (hasTermPwc && Array.IndexOf(ln.Suffix, SubjectTreeParts.TSep) < 0)
                        {
                            if (!fn(ConcatBytes(pre, ln.Suffix), ln.Value)) return false;
                        }
                    }
                    else if (hasTermPwc)
                    {
                        if (!MatchNode(cn, termParts, pre, fn)) return false;
                    }
                }
                return true;
            }

            // Re-put the terminal fwc if nparts was exhausted by matching.
            if (hasFwc && nparts.Length == 0)
                nparts = parts[^1..];

            var fp = nparts[0];
            var pByte = SubjectTreeParts.Pivot(fp, 0);

            if (fp.Length == 1 && (pByte == SubjectTreeParts.Pwc || pByte == SubjectTreeParts.Fwc))
            {
                // Wildcard part — iterate all children.
                foreach (var cn in n.Children())
                {
                    if (cn != null!)
                    {
                        if (!MatchNode(cn, nparts, pre, fn)) return false;
                    }
                }
                return true;
            }

            // Literal part — find specific child and loop.
            var nextNode = n.FindChild(pByte);
            if (nextNode == null) return true;
            n = nextNode;
            parts = nparts;
        }

        return true;
    }

    // -------------------------------------------------------------------------
    // Internal iteration
    // -------------------------------------------------------------------------

    private static bool IterNode(ISubjectTreeNode<T> n, byte[] pre, bool ordered, Func<byte[], T, bool> fn)
    {
        if (n.IsLeaf)
        {
            var ln = (SubjectTreeLeaf<T>)n;
            return fn(ConcatBytes(pre, ln.Suffix), ln.Value);
        }

        pre = ConcatBytes(pre, n.Prefix);

        if (!ordered)
        {
            foreach (var cn in n.Children())
            {
                if (cn == null!) continue;
                if (!IterNode(cn, pre, false, fn)) return false;
            }
            return true;
        }

        // Ordered: sort children by their path bytes lexicographically.
        var children = n.Children().Where(c => c != null!).ToArray();
        Array.Sort(children, static (a, b) => a.Path.AsSpan().SequenceCompareTo(b.Path.AsSpan()));
        foreach (var cn in children)
        {
            if (!IterNode(cn, pre, true, fn)) return false;
        }
        return true;
    }

    // -------------------------------------------------------------------------
    // Byte array helpers
    // -------------------------------------------------------------------------

    internal static byte[] ConcatBytes(byte[] a, byte[] b)
    {
        if (a.Length == 0) return b.Length == 0 ? Array.Empty<byte>() : b;
        if (b.Length == 0) return a;
        var result = new byte[a.Length + b.Length];
        a.CopyTo(result, 0);
        b.CopyTo(result, a.Length);
        return result;
    }
}
