// Go reference: server/stree/node.go, leaf.go, node4.go, node10.go, node16.go, node48.go, node256.go
namespace NATS.Server.Internal.SubjectTree;

/// <summary>
/// Internal node interface for the Adaptive Radix Tree.
/// </summary>
internal interface INode
{
    /// <summary>
    /// Gets whether this node is a terminal subject node that directly stores a subscription value.
    /// </summary>
    bool IsLeaf { get; }

    /// <summary>
    /// Gets structural metadata for branch nodes, including compressed path prefix and child count.
    /// </summary>
    NodeMeta? Base { get; }

    /// <summary>
    /// Sets the compressed path fragment represented by this node.
    /// </summary>
    /// <param name="pre">Subject bytes shared by all descendants below this node.</param>
    void SetPrefix(ReadOnlySpan<byte> pre);

    /// <summary>
    /// Adds a child edge for the next subject byte in the adaptive radix tree.
    /// </summary>
    /// <param name="c">Subject byte used to route lookups to the child node.</param>
    /// <param name="n">Child node that owns the remaining subject suffix for this edge.</param>
    void AddChild(byte c, INode n);
    /// <summary>
    /// Returns the child node for the given key byte, or null if not found.
    /// The returned wrapper allows in-place replacement of the child reference.
    /// </summary>
    /// <param name="c">Subject byte to look up in the node's child index.</param>
    ChildRef? FindChild(byte c);

    /// <summary>
    /// Removes the child edge for the provided subject byte.
    /// </summary>
    /// <param name="c">Subject byte whose child mapping should be removed.</param>
    void DeleteChild(byte c);

    /// <summary>
    /// Gets whether this node has reached its capacity and must grow to the next node shape.
    /// </summary>
    bool IsFull { get; }

    /// <summary>
    /// Expands this node to a larger branching factor to accept more distinct subject bytes.
    /// </summary>
    INode Grow();

    /// <summary>
    /// Attempts to shrink this node to a smaller branching representation when sparse.
    /// </summary>
    INode? Shrink();

    /// <summary>
    /// Matches a subject split into tokens against this node's compressed path fragment.
    /// </summary>
    /// <param name="parts">Remaining subject tokens to match from this node downward.</param>
    /// <returns>The remaining tokens after consuming this node, and whether the fragment matched.</returns>
    (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts);

    /// <summary>
    /// Gets a short node kind name used by diagnostics and debugging tools.
    /// </summary>
    string Kind { get; }

    /// <summary>
    /// Iterates child nodes until the callback returns <see langword="false" />.
    /// </summary>
    /// <param name="f">Callback invoked for each child node in this branch.</param>
    void Iter(Func<INode, bool> f);

    /// <summary>
    /// Returns the current child nodes for traversal or inspection.
    /// </summary>
    INode?[] Children();

    /// <summary>
    /// Gets the number of active child edges in this node.
    /// </summary>
    ushort NumChildren { get; }

    /// <summary>
    /// Gets the compressed path bytes represented by this node.
    /// </summary>
    byte[] Path();
}

/// <summary>
/// Wrapper that allows in-place replacement of a child reference in a node.
/// This is analogous to Go's *node pointer.
/// </summary>
internal sealed class ChildRef(Func<INode?> getter, Action<INode?> setter)
{
    /// <summary>
    /// Gets or replaces the child node reference stored at a specific branch slot.
    /// </summary>
    public INode? Node
    {
        get => getter();
        set => setter(value);
    }
}

/// <summary>
/// Base metadata for internal (non-leaf) nodes.
/// </summary>
internal sealed class NodeMeta
{
    /// <summary>
    /// Gets or sets the compressed subject prefix shared by descendants of this branch node.
    /// </summary>
    public byte[] Prefix { get; set; } = [];

    /// <summary>
    /// Gets or sets the number of child edges currently populated for this branch node.
    /// </summary>
    public ushort Size { get; set; }
}

#region Leaf Node

/// <summary>
/// Leaf node holding a value and suffix.
/// Go reference: server/stree/leaf.go
/// </summary>
internal sealed class Leaf<T> : INode
{
    public T Value;
    public byte[] Suffix;

    /// <summary>
    /// Initializes a terminal subject-tree node that stores a value for an exact suffix match.
    /// </summary>
    /// <param name="suffix">Remaining subject bytes that must match to resolve this leaf.</param>
    /// <param name="value">Subscription payload or state associated with the matched subject.</param>
    public Leaf(ReadOnlySpan<byte> suffix, T value)
    {
        Value = value;
        Suffix = Parts.CopyBytes(suffix);
    }

    /// <inheritdoc />
    public bool IsLeaf => true;
    /// <inheritdoc />
    public NodeMeta? Base => null;
    /// <inheritdoc />
    public bool IsFull => true;
    /// <inheritdoc />
    public ushort NumChildren => 0;
    /// <inheritdoc />
    public string Kind => "LEAF";

    /// <summary>
    /// Checks whether the provided subject bytes exactly match this leaf suffix.
    /// </summary>
    /// <param name="subject">Subject bytes remaining after traversing parent branch prefixes.</param>
    /// <returns><see langword="true" /> when the subject resolves to this exact leaf.</returns>
    public bool Match(ReadOnlySpan<byte> subject) => subject.SequenceEqual(Suffix);

    /// <summary>
    /// Replaces the stored suffix when leaf content is split or merged during tree updates.
    /// </summary>
    /// <param name="suffix">New exact-match suffix bytes for this leaf.</param>
    public void SetSuffix(ReadOnlySpan<byte> suffix) => Suffix = Parts.CopyBytes(suffix);

    /// <inheritdoc />
    public byte[] Path() => Suffix;

    /// <inheritdoc />
    public INode?[] Children() => [];

    /// <inheritdoc />
    public void Iter(Func<INode, bool> f) { }

    /// <inheritdoc />
    public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
        => Parts.MatchPartsAgainstFragment(parts, Suffix);

    // These should not be called on a leaf.
    /// <inheritdoc />
    public void SetPrefix(ReadOnlySpan<byte> pre) => throw new InvalidOperationException("setPrefix called on leaf");
    /// <inheritdoc />
    public void AddChild(byte c, INode n) => throw new InvalidOperationException("addChild called on leaf");
    /// <inheritdoc />
    public ChildRef? FindChild(byte c) => throw new InvalidOperationException("findChild called on leaf");
    /// <inheritdoc />
    public INode Grow() => throw new InvalidOperationException("grow called on leaf");
    /// <inheritdoc />
    public void DeleteChild(byte c) => throw new InvalidOperationException("deleteChild called on leaf");
    /// <inheritdoc />
    public INode? Shrink() => throw new InvalidOperationException("shrink called on leaf");
}

#endregion

#region Node4

/// <summary>
/// Node with up to 4 children.
/// Go reference: server/stree/node4.go
/// </summary>
internal sealed class Node4 : INode
{
    private readonly INode?[] _child = new INode?[4];
    private readonly byte[] _key = new byte[4];
    internal readonly NodeMeta Meta = new();

    /// <summary>
    /// Initializes a small branch node for up to four subject-byte fan-out edges.
    /// </summary>
    /// <param name="prefix">Compressed subject prefix represented by this branch.</param>
    public Node4(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

    /// <inheritdoc />
    public bool IsLeaf => false;
    /// <inheritdoc />
    public NodeMeta? Base => Meta;
    /// <inheritdoc />
    public ushort NumChildren => Meta.Size;
    /// <inheritdoc />
    public bool IsFull => Meta.Size >= 4;
    /// <inheritdoc />
    public string Kind => "NODE4";
    /// <inheritdoc />
    public byte[] Path() => Meta.Prefix;

    /// <inheritdoc />
    public void SetPrefix(ReadOnlySpan<byte> pre)
    {
        Meta.Prefix = pre.ToArray();
    }

    /// <inheritdoc />
    public void AddChild(byte c, INode n)
    {
        if (Meta.Size >= 4) throw new InvalidOperationException("node4 full!");
        _key[Meta.Size] = c;
        _child[Meta.Size] = n;
        Meta.Size++;
    }

    /// <inheritdoc />
    public ChildRef? FindChild(byte c)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (_key[i] == c)
            {
                var idx = i;
                return new ChildRef(() => _child[idx], v => _child[idx] = v);
            }
        }
        return null;
    }

    /// <inheritdoc />
    public void DeleteChild(byte c)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (_key[i] == c)
            {
                var last = Meta.Size - 1;
                if (i < last)
                {
                    _key[i] = _key[last];
                    _child[i] = _child[last];
                    _key[last] = 0;
                    _child[last] = null;
                }
                else
                {
                    _key[i] = 0;
                    _child[i] = null;
                }
                Meta.Size--;
                return;
            }
        }
    }

    /// <inheritdoc />
    public INode Grow()
    {
        var nn = new Node10(Meta.Prefix);
        for (int i = 0; i < 4; i++)
        {
            nn.AddChild(_key[i], _child[i]!);
        }
        return nn;
    }

    /// <inheritdoc />
    public INode? Shrink()
    {
        if (Meta.Size == 1) return _child[0];
        return null;
    }

    /// <inheritdoc />
    public void Iter(Func<INode, bool> f)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (!f(_child[i]!)) return;
        }
    }

    /// <inheritdoc />
    public INode?[] Children()
    {
        var result = new INode?[Meta.Size];
        Array.Copy(_child, result, Meta.Size);
        return result;
    }

    /// <inheritdoc />
    public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
        => Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}

#endregion

#region Node10

/// <summary>
/// Node with up to 10 children. Optimized for numeric subject tokens (0-9).
/// Go reference: server/stree/node10.go
/// </summary>
internal sealed class Node10 : INode
{
    private readonly INode?[] _child = new INode?[10];
    private readonly byte[] _key = new byte[10];
    internal readonly NodeMeta Meta = new();

    /// <summary>
    /// Initializes a branch node tuned for numeric token fan-out, common in ordered stream subjects.
    /// </summary>
    /// <param name="prefix">Compressed subject prefix represented by this branch.</param>
    public Node10(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

    /// <inheritdoc />
    public bool IsLeaf => false;
    /// <inheritdoc />
    public NodeMeta? Base => Meta;
    /// <inheritdoc />
    public ushort NumChildren => Meta.Size;
    /// <inheritdoc />
    public bool IsFull => Meta.Size >= 10;
    /// <inheritdoc />
    public string Kind => "NODE10";
    /// <inheritdoc />
    public byte[] Path() => Meta.Prefix;

    /// <inheritdoc />
    public void SetPrefix(ReadOnlySpan<byte> pre)
    {
        Meta.Prefix = pre.ToArray();
    }

    /// <inheritdoc />
    public void AddChild(byte c, INode n)
    {
        if (Meta.Size >= 10) throw new InvalidOperationException("node10 full!");
        _key[Meta.Size] = c;
        _child[Meta.Size] = n;
        Meta.Size++;
    }

    /// <inheritdoc />
    public ChildRef? FindChild(byte c)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (_key[i] == c)
            {
                var idx = i;
                return new ChildRef(() => _child[idx], v => _child[idx] = v);
            }
        }
        return null;
    }

    /// <inheritdoc />
    public void DeleteChild(byte c)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (_key[i] == c)
            {
                var last = Meta.Size - 1;
                if (i < last)
                {
                    _key[i] = _key[last];
                    _child[i] = _child[last];
                    _key[last] = 0;
                    _child[last] = null;
                }
                else
                {
                    _key[i] = 0;
                    _child[i] = null;
                }
                Meta.Size--;
                return;
            }
        }
    }

    /// <inheritdoc />
    public INode Grow()
    {
        var nn = new Node16(Meta.Prefix);
        for (int i = 0; i < 10; i++)
        {
            nn.AddChild(_key[i], _child[i]!);
        }
        return nn;
    }

    /// <inheritdoc />
    public INode? Shrink()
    {
        if (Meta.Size > 4) return null;
        var nn = new Node4([]);
        for (int i = 0; i < Meta.Size; i++)
        {
            nn.AddChild(_key[i], _child[i]!);
        }
        return nn;
    }

    /// <inheritdoc />
    public void Iter(Func<INode, bool> f)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (!f(_child[i]!)) return;
        }
    }

    /// <inheritdoc />
    public INode?[] Children()
    {
        var result = new INode?[Meta.Size];
        Array.Copy(_child, result, Meta.Size);
        return result;
    }

    /// <inheritdoc />
    public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
        => Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}

#endregion

#region Node16

/// <summary>
/// Node with up to 16 children.
/// Go reference: server/stree/node16.go
/// </summary>
internal sealed class Node16 : INode
{
    private readonly INode?[] _child = new INode?[16];
    private readonly byte[] _key = new byte[16];
    internal readonly NodeMeta Meta = new();

    /// <summary>
    /// Initializes a medium branch node for moderate subject fan-out without index indirection.
    /// </summary>
    /// <param name="prefix">Compressed subject prefix represented by this branch.</param>
    public Node16(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

    /// <inheritdoc />
    public bool IsLeaf => false;
    /// <inheritdoc />
    public NodeMeta? Base => Meta;
    /// <inheritdoc />
    public ushort NumChildren => Meta.Size;
    /// <inheritdoc />
    public bool IsFull => Meta.Size >= 16;
    /// <inheritdoc />
    public string Kind => "NODE16";
    /// <inheritdoc />
    public byte[] Path() => Meta.Prefix;

    /// <inheritdoc />
    public void SetPrefix(ReadOnlySpan<byte> pre)
    {
        Meta.Prefix = pre.ToArray();
    }

    /// <inheritdoc />
    public void AddChild(byte c, INode n)
    {
        if (Meta.Size >= 16) throw new InvalidOperationException("node16 full!");
        _key[Meta.Size] = c;
        _child[Meta.Size] = n;
        Meta.Size++;
    }

    /// <inheritdoc />
    public ChildRef? FindChild(byte c)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (_key[i] == c)
            {
                var idx = i;
                return new ChildRef(() => _child[idx], v => _child[idx] = v);
            }
        }
        return null;
    }

    /// <inheritdoc />
    public void DeleteChild(byte c)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (_key[i] == c)
            {
                var last = Meta.Size - 1;
                if (i < last)
                {
                    _key[i] = _key[last];
                    _child[i] = _child[last];
                    _key[last] = 0;
                    _child[last] = null;
                }
                else
                {
                    _key[i] = 0;
                    _child[i] = null;
                }
                Meta.Size--;
                return;
            }
        }
    }

    /// <inheritdoc />
    public INode Grow()
    {
        var nn = new Node48(Meta.Prefix);
        for (int i = 0; i < 16; i++)
        {
            nn.AddChild(_key[i], _child[i]!);
        }
        return nn;
    }

    /// <inheritdoc />
    public INode? Shrink()
    {
        if (Meta.Size > 10) return null;
        var nn = new Node10([]);
        for (int i = 0; i < Meta.Size; i++)
        {
            nn.AddChild(_key[i], _child[i]!);
        }
        return nn;
    }

    /// <inheritdoc />
    public void Iter(Func<INode, bool> f)
    {
        for (int i = 0; i < Meta.Size; i++)
        {
            if (!f(_child[i]!)) return;
        }
    }

    /// <inheritdoc />
    public INode?[] Children()
    {
        var result = new INode?[Meta.Size];
        Array.Copy(_child, result, Meta.Size);
        return result;
    }

    /// <inheritdoc />
    public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
        => Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}

#endregion

#region Node48

/// <summary>
/// Node with up to 48 children. Uses a 256-byte index array (1-indexed) to map keys to child slots.
/// Go reference: server/stree/node48.go
/// </summary>
internal sealed class Node48 : INode
{
    internal readonly INode?[] Child = new INode?[48];
    internal readonly byte[] Key = new byte[256]; // 1-indexed: 0 means no entry
    internal readonly NodeMeta Meta = new();

    /// <summary>
    /// Initializes a high fan-out branch node that trades memory for faster byte-key lookups.
    /// </summary>
    /// <param name="prefix">Compressed subject prefix represented by this branch.</param>
    public Node48(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

    /// <inheritdoc />
    public bool IsLeaf => false;
    /// <inheritdoc />
    public NodeMeta? Base => Meta;
    /// <inheritdoc />
    public ushort NumChildren => Meta.Size;
    /// <inheritdoc />
    public bool IsFull => Meta.Size >= 48;
    /// <inheritdoc />
    public string Kind => "NODE48";
    /// <inheritdoc />
    public byte[] Path() => Meta.Prefix;

    /// <inheritdoc />
    public void SetPrefix(ReadOnlySpan<byte> pre)
    {
        Meta.Prefix = pre.ToArray();
    }

    /// <inheritdoc />
    public void AddChild(byte c, INode n)
    {
        if (Meta.Size >= 48) throw new InvalidOperationException("node48 full!");
        Child[Meta.Size] = n;
        Key[c] = (byte)(Meta.Size + 1); // 1-indexed
        Meta.Size++;
    }

    /// <inheritdoc />
    public ChildRef? FindChild(byte c)
    {
        var i = Key[c];
        if (i == 0) return null;
        var idx = i - 1;
        return new ChildRef(() => Child[idx], v => Child[idx] = v);
    }

    /// <inheritdoc />
    public void DeleteChild(byte c)
    {
        var i = Key[c];
        if (i == 0) return;
        i--; // Adjust for 1-indexing
        var last = (byte)(Meta.Size - 1);
        if (i < last)
        {
            Child[i] = Child[last];
            for (int ic = 0; ic < 256; ic++)
            {
                if (Key[ic] == last + 1)
                {
                    Key[ic] = (byte)(i + 1);
                    break;
                }
            }
        }
        Child[last] = null;
        Key[c] = 0;
        Meta.Size--;
    }

    /// <inheritdoc />
    public INode Grow()
    {
        var nn = new Node256(Meta.Prefix);
        for (int c = 0; c < 256; c++)
        {
            var i = Key[c];
            if (i > 0)
            {
                nn.AddChild((byte)c, Child[i - 1]!);
            }
        }
        return nn;
    }

    /// <inheritdoc />
    public INode? Shrink()
    {
        if (Meta.Size > 16) return null;
        var nn = new Node16([]);
        for (int c = 0; c < 256; c++)
        {
            var i = Key[c];
            if (i > 0)
            {
                nn.AddChild((byte)c, Child[i - 1]!);
            }
        }
        return nn;
    }

    /// <inheritdoc />
    public void Iter(Func<INode, bool> f)
    {
        foreach (var c in Child)
        {
            if (c != null && !f(c)) return;
        }
    }

    /// <inheritdoc />
    public INode?[] Children()
    {
        var result = new INode?[Meta.Size];
        Array.Copy(Child, result, Meta.Size);
        return result;
    }

    /// <inheritdoc />
    public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
        => Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}

#endregion

#region Node256

/// <summary>
/// Node with up to 256 children. Direct array indexed by byte value.
/// Go reference: server/stree/node256.go
/// </summary>
internal sealed class Node256 : INode
{
    internal readonly INode?[] Child = new INode?[256];
    internal readonly NodeMeta Meta = new();

    /// <summary>
    /// Initializes the maximum fan-out branch node with direct byte-to-child indexing.
    /// </summary>
    /// <param name="prefix">Compressed subject prefix represented by this branch.</param>
    public Node256(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

    /// <inheritdoc />
    public bool IsLeaf => false;
    /// <inheritdoc />
    public NodeMeta? Base => Meta;
    /// <inheritdoc />
    public ushort NumChildren => Meta.Size;
    /// <inheritdoc />
    public bool IsFull => false; // node256 is never full
    /// <inheritdoc />
    public string Kind => "NODE256";
    /// <inheritdoc />
    public byte[] Path() => Meta.Prefix;

    /// <inheritdoc />
    public void SetPrefix(ReadOnlySpan<byte> pre)
    {
        Meta.Prefix = pre.ToArray();
    }

    /// <inheritdoc />
    public void AddChild(byte c, INode n)
    {
        Child[c] = n;
        Meta.Size++;
    }

    /// <inheritdoc />
    public ChildRef? FindChild(byte c)
    {
        if (Child[c] == null) return null;
        return new ChildRef(() => Child[c], v => Child[c] = v);
    }

    /// <inheritdoc />
    public void DeleteChild(byte c)
    {
        if (Child[c] != null)
        {
            Child[c] = null;
            Meta.Size--;
        }
    }

    /// <inheritdoc />
    public INode Grow() => throw new InvalidOperationException("grow can not be called on node256");

    /// <inheritdoc />
    public INode? Shrink()
    {
        if (Meta.Size > 48) return null;
        var nn = new Node48([]);
        for (int c = 0; c < 256; c++)
        {
            if (Child[c] != null)
            {
                nn.AddChild((byte)c, Child[c]!);
            }
        }
        return nn;
    }

    /// <inheritdoc />
    public void Iter(Func<INode, bool> f)
    {
        for (int i = 0; i < 256; i++)
        {
            if (Child[i] != null)
            {
                if (!f(Child[i]!)) return;
            }
        }
    }

    /// <inheritdoc />
    public INode?[] Children()
    {
        // Return the full 256 array, same as Go
        return (INode?[])Child.Clone();
    }

    /// <inheritdoc />
    public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
        => Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}

#endregion