using System.IO;
using System.Text;
using NATS.Server.Internal.Gsl;

// Go reference: server/stree/stree.go
namespace NATS.Server.Internal.SubjectTree;

/// <summary>
/// SubjectTree is an adaptive radix trie (ART) for storing subject information on literal subjects.
/// Uses dynamic nodes, path compression and lazy expansion.
/// Go reference: server/stree/stree.go
/// </summary>
public class SubjectTree<T>
{
    internal INode? Root;
    private int _size;

    /// <summary>
    /// Returns the number of elements stored.
    /// </summary>
    public int Size => _size;

    /// <summary>
    /// Empties the tree and returns it. If called on a new tree, returns it unchanged.
    /// </summary>
    public SubjectTree<T> Empty()
    {
        Root = null;
        _size = 0;
        return this;
    }

    /// <summary>
    /// Insert a value into the tree. Returns (oldValue, existed).
    /// If the subject already existed, oldValue is the previous value and existed is true.
    /// </summary>
    public (T? OldValue, bool Existed) Insert(ReadOnlySpan<byte> subject, T value)
    {
        // Make sure we never insert anything with a noPivot byte.
        if (subject.IndexOf(Parts.NoPivot) >= 0)
        {
            return (default, false);
        }

        var (old, updated) = InsertInternal(ref Root, subject.ToArray(), value, 0);
        if (!updated)
        {
            _size++;
        }
        return (old, updated);
    }

    /// <summary>
    /// Find the value for an exact subject match.
    /// </summary>
    public (T? Value, bool Found) Find(ReadOnlySpan<byte> subject)
    {
        int si = 0;
        var n = Root;
        while (n != null)
        {
            if (n.IsLeaf)
            {
                var ln = (Leaf<T>)n;
                if (ln.Match(subject[si..]))
                {
                    return (ln.Value, true);
                }
                return (default, false);
            }

            // We are a node type here, grab meta portion.
            var bn = n.Base!;
            if (bn.Prefix.Length > 0)
            {
                var end = Math.Min(si + bn.Prefix.Length, subject.Length);
                if (!subject[si..end].SequenceEqual(bn.Prefix))
                {
                    return (default, false);
                }
                si += bn.Prefix.Length;
            }

            var childRef = n.FindChild(Parts.Pivot(subject, si));
            if (childRef != null)
            {
                n = childRef.Node;
            }
            else
            {
                return (default, false);
            }
        }
        return (default, false);
    }

    /// <summary>
    /// Delete the item for the given subject.
    /// Returns (deletedValue, wasFound).
    /// </summary>
    public (T? Value, bool Found) Delete(ReadOnlySpan<byte> subject)
    {
        if (subject.Length == 0)
        {
            return (default, false);
        }

        var (val, deleted) = DeleteInternal(ref Root, subject.ToArray(), 0);
        if (deleted)
        {
            _size--;
        }
        return (val, deleted);
    }

    /// <summary>
    /// Match against a filter subject with wildcards and invoke the callback for each matched value.
    /// </summary>
    public void Match(ReadOnlySpan<byte> filter, Action<byte[], T>? callback)
    {
        if (Root == null || filter.Length == 0 || callback == null)
        {
            return;
        }

        var parts = Parts.GenParts(filter);
        MatchInternal(Root, parts, [], (subject, val) =>
        {
            callback(subject, val);
            return true;
        });
    }

    /// <summary>
    /// Match against a filter subject with wildcards and invoke the callback for each matched value.
    /// Returning false from the callback stops matching immediately.
    /// Returns true if matching ran to completion, false if callback stopped it early.
    /// </summary>
    public bool MatchUntil(ReadOnlySpan<byte> filter, Func<byte[], T, bool>? callback)
    {
        if (Root == null || filter.Length == 0 || callback == null)
        {
            return true;
        }

        var parts = Parts.GenParts(filter);
        return MatchInternal(Root, parts, [], callback);
    }

    /// <summary>
    /// Walk all entries in lexicographic order. The callback can return false to terminate.
    /// </summary>
    public void IterOrdered(Func<byte[], T, bool> cb)
    {
        if (Root == null) return;
        IterInternal(Root, [], ordered: true, cb);
    }

    /// <summary>
    /// Walk all entries in no guaranteed order. The callback can return false to terminate.
    /// </summary>
    public void IterFast(Func<byte[], T, bool> cb)
    {
        if (Root == null) return;
        IterInternal(Root, [], ordered: false, cb);
    }

    /// <summary>
    /// Dumps a human-readable representation of the tree.
    /// Go reference: server/stree/dump.go
    /// </summary>
    public void Dump(TextWriter writer)
    {
        Dump(writer, Root, 0);
        writer.WriteLine();
    }

    #region Internal Methods

    private void Dump(TextWriter writer, INode? node, int depth)
    {
        if (node == null)
        {
            writer.WriteLine("EMPTY");
            return;
        }

        if (node.IsLeaf)
        {
            var leaf = (Leaf<T>)node;
            writer.WriteLine($"{DumpPre(depth)} LEAF: Suffix: {QuoteBytes(leaf.Suffix)} Value: {leaf.Value}");
            return;
        }

        var bn = node.Base!;
        writer.WriteLine($"{DumpPre(depth)} {node.Kind} Prefix: {QuoteBytes(bn.Prefix)}");
        depth++;
        node.Iter(n =>
        {
            Dump(writer, n, depth);
            return true;
        });
    }

    private static string DumpPre(int depth)
    {
        if (depth == 0)
            return "-- ";

        var sb = new StringBuilder(depth * 2 + 4);
        for (var i = 0; i < depth; i++)
            sb.Append("  ");
        sb.Append("|__ ");
        return sb.ToString();
    }

    private static string QuoteBytes(byte[] data) => $"\"{Encoding.UTF8.GetString(data)}\"";

    /// <summary>
    /// Internal recursive insert.
    /// Go reference: server/stree/stree.go:insert
    /// </summary>
    private (T? OldValue, bool Updated) InsertInternal(ref INode? nodeRef, byte[] subject, T value, int si)
    {
        var n = nodeRef;
        if (n == null)
        {
            nodeRef = new Leaf<T>(subject[si..], value);
            return (default, false);
        }

        if (n.IsLeaf)
        {
            var ln = (Leaf<T>)n;
            if (ln.Match(subject.AsSpan(si)))
            {
                // Replace with new value.
                var old = ln.Value;
                ln.Value = value;
                return (old, true);
            }

            // Here we need to split this leaf.
            int cpi = Parts.CommonPrefixLen(ln.Suffix, subject.AsSpan(si));
            var nn = new Node4(subject.AsSpan(si, cpi));
            ln.SetSuffix(ln.Suffix.AsSpan(cpi));
            si += cpi;

            // Make sure we have different pivot, normally this will be the case unless we have overflowing prefixes.
            byte p = Parts.Pivot(ln.Suffix, 0);
            if (cpi > 0 && si < subject.Length && p == subject[si])
            {
                // We need to split the original leaf. Recursively call into insert.
                InsertInternal(ref nodeRef, subject, value, si);
                // Now add the updated version of nodeRef as a child to the new node4.
                nn.AddChild(p, nodeRef!);
            }
            else
            {
                // Can just add this new leaf as a sibling.
                var nl = new Leaf<T>(subject.AsSpan(si), value);
                nn.AddChild(Parts.Pivot(nl.Suffix, 0), nl);
                // Add back original.
                nn.AddChild(Parts.Pivot(ln.Suffix, 0), ln);
            }

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

        // Non-leaf nodes.
        var bn = n.Base!;
        if (bn.Prefix.Length > 0)
        {
            int cpi = Parts.CommonPrefixLen(bn.Prefix, subject.AsSpan(si));
            int pli = bn.Prefix.Length;
            if (cpi >= pli)
            {
                // Move past this node.
                si += pli;
                var childRef = n.FindChild(Parts.Pivot(subject, si));
                if (childRef != null)
                {
                    var childNode = childRef.Node;
                    var result = InsertInternal(ref childNode, subject, value, si);
                    childRef.Node = childNode;
                    return result;
                }
                if (n.IsFull)
                {
                    n = n.Grow();
                    nodeRef = n;
                }
                n.AddChild(Parts.Pivot(subject, si), new Leaf<T>(subject.AsSpan(si), value));
                return (default, false);
            }
            else
            {
                // We did not match the prefix completely here.
                var prefix = subject.AsSpan(si, cpi);
                si += prefix.Length;
                // We will insert a new node4 and attach our current node below after adjusting prefix.
                var nn = new Node4(prefix);
                // Shift the prefix for our original node.
                n.SetPrefix(bn.Prefix.AsSpan(cpi));
                nn.AddChild(Parts.Pivot(bn.Prefix, 0), n);
                // Add in our new leaf.
                nn.AddChild(Parts.Pivot(subject.AsSpan(si), 0), new Leaf<T>(subject.AsSpan(si), value));
                // Update our node reference.
                nodeRef = nn;
            }
        }
        else
        {
            var childRef = n.FindChild(Parts.Pivot(subject, si));
            if (childRef != null)
            {
                var childNode = childRef.Node;
                var result = InsertInternal(ref childNode, subject, value, si);
                childRef.Node = childNode;
                return result;
            }
            // No prefix and no matched child, so add in new leafnode as needed.
            if (n.IsFull)
            {
                n = n.Grow();
                nodeRef = n;
            }
            n.AddChild(Parts.Pivot(subject, si), new Leaf<T>(subject.AsSpan(si), value));
        }

        return (default, false);
    }

    /// <summary>
    /// Internal recursive delete with compaction.
    /// Go reference: server/stree/stree.go:delete
    /// </summary>
    private (T? Value, bool Deleted) DeleteInternal(ref INode? nodeRef, byte[] subject, int si)
    {
        if (nodeRef == null || subject.Length == 0)
        {
            return (default, false);
        }

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

        // Not a leaf node.
        var bn = n.Base!;
        if (bn.Prefix.Length > 0)
        {
            // subject could be shorter and would panic on bad index.
            if (subject.Length < si + bn.Prefix.Length)
            {
                return (default, false);
            }
            if (!subject.AsSpan(si, bn.Prefix.Length).SequenceEqual(bn.Prefix))
            {
                return (default, false);
            }
            si += bn.Prefix.Length;
        }

        var p = Parts.Pivot(subject, si);
        var childRef = n.FindChild(p);
        if (childRef == null)
        {
            return (default, false);
        }

        var nn = childRef.Node;
        if (nn != null && nn.IsLeaf)
        {
            var ln = (Leaf<T>)nn;
            if (ln.Match(subject.AsSpan(si)))
            {
                n.DeleteChild(p);

                var sn = n.Shrink();
                if (sn != null)
                {
                    // Make sure to copy prefix so we force a copy below.
                    var pre = bn.Prefix.ToArray();

                    // Need to fix up prefixes/suffixes.
                    if (sn.IsLeaf)
                    {
                        var shrunkLeaf = (Leaf<T>)sn;
                        // Prepend old prefix to leaf suffix.
                        var newSuffix = new byte[pre.Length + shrunkLeaf.Suffix.Length];
                        pre.CopyTo(newSuffix, 0);
                        shrunkLeaf.Suffix.CopyTo(newSuffix, pre.Length);
                        shrunkLeaf.Suffix = newSuffix;
                    }
                    else
                    {
                        // We are a node here, we need to add in the old prefix.
                        if (pre.Length > 0)
                        {
                            var bsn = sn.Base!;
                            var newPrefix = new byte[pre.Length + bsn.Prefix.Length];
                            pre.CopyTo(newPrefix, 0);
                            bsn.Prefix.CopyTo(newPrefix, pre.Length);
                            sn.SetPrefix(newPrefix);
                        }
                    }
                    nodeRef = sn;
                }

                return (ln.Value, true);
            }
            return (default, false);
        }

        // Recurse into child node.
        var childNode = childRef.Node;
        var result = DeleteInternal(ref childNode, subject, si);
        childRef.Node = childNode;
        return result;
    }

    /// <summary>
    /// Internal recursive match.
    /// Go reference: server/stree/stree.go:match
    /// </summary>
    internal bool MatchInternal(INode? n, ReadOnlyMemory<byte>[] parts, byte[] pre, Func<byte[], T, bool> cb)
    {
        // Capture if we are sitting on a terminal fwc.
        bool hasFWC = false;
        if (parts.Length > 0 && parts[^1].Length > 0 && parts[^1].Span[0] == Parts.Fwc)
        {
            hasFWC = true;
        }

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

            // We have matched here. If we are a leaf and have exhausted all parts or have a FWC, fire callback.
            if (n.IsLeaf)
            {
                if (nparts.Length == 0 || (hasFWC && nparts.Length == 1))
                {
                    var ln = (Leaf<T>)n;
                    var subject = Concat(pre, ln.Suffix);
                    if (!cb(subject, ln.Value))
                    {
                        return false;
                    }
                }
                return true;
            }

            // We have normal nodes here. Append our prefix.
            var bn = n.Base!;
            if (bn.Prefix.Length > 0)
            {
                pre = Concat(pre, bn.Prefix);
            }

            // Check our remaining parts.
            if (nparts.Length == 0 && !hasFWC)
            {
                // We are a node with no parts left and we are not looking at a fwc.
                bool hasTermPWC = false;
                if (parts.Length > 0 && parts[^1].Length == 1 && parts[^1].Span[0] == Parts.Pwc)
                {
                    nparts = parts[^1..];
                    hasTermPWC = true;
                }

                foreach (var cn in n.Children())
                {
                    if (cn == null) continue;

                    if (cn.IsLeaf)
                    {
                        var ln = (Leaf<T>)cn;
                        if (ln.Suffix.Length == 0)
                        {
                            var subject = Concat(pre, ln.Suffix);
                            if (!cb(subject, ln.Value))
                            {
                                return false;
                            }
                        }
                        else if (hasTermPWC && ln.Suffix.AsSpan().IndexOf(Parts.Tsep) < 0)
                        {
                            var subject = Concat(pre, ln.Suffix);
                            if (!cb(subject, ln.Value))
                            {
                                return false;
                            }
                        }
                    }
                    else if (hasTermPWC)
                    {
                        if (!MatchInternal(cn, nparts, pre, cb))
                        {
                            return false;
                        }
                    }
                }
                return true;
            }

            // If we are sitting on a terminal fwc, put back and continue.
            if (hasFWC && nparts.Length == 0)
            {
                nparts = parts[^1..];
            }

            // Here we are a node type with a partial match.
            // Check if the first part is a wildcard.
            var fp = nparts[0];
            var pvt = Parts.Pivot(fp.Span, 0);

            if (fp.Length == 1 && (pvt == Parts.Pwc || pvt == Parts.Fwc))
            {
                // We need to iterate over all children here for the current node
                // to see if we match further down.
                foreach (var cn in n.Children())
                {
                    if (cn != null)
                    {
                        if (!MatchInternal(cn, nparts, pre, cb))
                        {
                            return false;
                        }
                    }
                }
                return true;
            }

            // Here we have normal traversal, so find the next child.
            var next = n.FindChild(pvt);
            if (next == null)
            {
                return true;
            }
            n = next.Node;
            parts = nparts;
        }
        return true;
    }

    /// <summary>
    /// Internal iter function to walk nodes.
    /// Go reference: server/stree/stree.go:iter
    /// </summary>
    internal bool IterInternal(INode n, byte[] pre, bool ordered, Func<byte[], T, bool> cb)
    {
        if (n.IsLeaf)
        {
            var ln = (Leaf<T>)n;
            return cb(Concat(pre, ln.Suffix), ln.Value);
        }

        // We are normal node here.
        var bn = n.Base!;
        if (bn.Prefix.Length > 0)
        {
            pre = Concat(pre, bn.Prefix);
        }

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

        // Collect non-null children and sort by path for lexicographic order.
        var children = n.Children().Where(c => c != null).ToList();
        children.Sort((a, b) =>
        {
            var pa = a!.Path();
            var pb = b!.Path();
            return pa.AsSpan().SequenceCompareTo(pb);
        });

        foreach (var cn in children)
        {
            if (!IterInternal(cn!, pre, true, cb))
            {
                return false;
            }
        }
        return true;
    }

    /// <summary>
    /// Helper to concatenate two byte arrays.
    /// </summary>
    private static byte[] Concat(byte[] a, byte[] b)
    {
        if (a.Length == 0) return 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;
    }

    #endregion
}

/// <summary>
/// Static helper methods for SubjectTree operations.
/// </summary>
public static class SubjectTreeHelper
{
    /// <summary>
    /// Iterates the smaller of the two provided subject trees and looks for matching entries in the other.
    /// Go reference: server/stree/stree.go:LazyIntersect
    /// </summary>
    public static void LazyIntersect<TL, TR>(SubjectTree<TL>? tl, SubjectTree<TR>? tr, Action<byte[], TL, TR> cb)
    {
        if (tl == null || tr == null || tl.Root == null || tr.Root == null)
        {
            return;
        }

        if (tl.Size <= tr.Size)
        {
            tl.IterFast((key, v1) =>
            {
                var (v2, ok) = tr.Find(key);
                if (ok)
                {
                    cb(key, v1, v2!);
                }
                return true;
            });
        }
        else
        {
            tr.IterFast((key, v2) =>
            {
                var (v1, ok) = tl.Find(key);
                if (ok)
                {
                    cb(key, v1!, v2);
                }
                return true;
            });
        }
    }

    /// <summary>
    /// Matches all items in the subject tree that have interest in the given sublist.
    /// The callback is invoked at most once per matching subject.
    /// Go reference: server/stree/stree.go IntersectGSL
    /// </summary>
    public static void IntersectGSL<T, SL>(
        SubjectTree<T>? tree,
        GenericSubjectList<SL>? sublist,
        Action<byte[], T> cb) where SL : IEquatable<SL>
    {
        if (tree?.Root == null || sublist == null)
            return;

        IntersectGslInternal(tree.Root, [], sublist, cb);
    }

    private static void IntersectGslInternal<T, SL>(
        INode node,
        byte[] pre,
        GenericSubjectList<SL> sublist,
        Action<byte[], T> cb) where SL : IEquatable<SL>
    {
        if (node.IsLeaf)
        {
            var leaf = (Leaf<T>)node;
            var subject = Concat(pre, leaf.Suffix);
            if (sublist.HasInterest(BytesToString(subject)))
                cb(subject, leaf.Value);
            return;
        }

        var bn = node.Base!;
        pre = Concat(pre, bn.Prefix);

        foreach (var child in node.Children())
        {
            if (child == null)
                continue;

            var subject = Concat(pre, child.Path());
            if (!HasInterestForTokens(sublist, subject, pre.Length))
                continue;

            IntersectGslInternal(child, pre, sublist, cb);
        }
    }

    private static bool HasInterestForTokens<SL>(
        GenericSubjectList<SL> sublist,
        byte[] subject,
        int since) where SL : IEquatable<SL>
    {
        for (var i = since; i < subject.Length; i++)
        {
            if (subject[i] != Parts.Tsep)
                continue;

            if (!sublist.HasInterestStartingIn(BytesToString(subject.AsSpan(0, i))))
                return false;
        }

        return true;
    }

    private static string BytesToString(ReadOnlySpan<byte> data)
    {
        if (data.Length == 0)
            return string.Empty;
        return Encoding.UTF8.GetString(data);
    }

    private static byte[] Concat(byte[] a, byte[] b)
    {
        if (a.Length == 0) return 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;
    }
}
