// 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
{
    bool IsLeaf { get; }
    NodeMeta? Base { get; }
    void SetPrefix(ReadOnlySpan<byte> pre);
    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>
    ChildRef? FindChild(byte c);
    void DeleteChild(byte c);
    bool IsFull { get; }
    INode Grow();
    INode? Shrink();
    (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts);
    string Kind { get; }
    void Iter(Func<INode, bool> f);
    INode?[] Children();
    ushort NumChildren { get; }
    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)
{
    public INode? Node
    {
        get => getter();
        set => setter(value);
    }
}

/// <summary>
/// Base metadata for internal (non-leaf) nodes.
/// </summary>
internal sealed class NodeMeta
{
    public byte[] Prefix { get; set; } = [];
    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;

    public Leaf(ReadOnlySpan<byte> suffix, T value)
    {
        Value = value;
        Suffix = Parts.CopyBytes(suffix);
    }

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

    public bool Match(ReadOnlySpan<byte> subject) => subject.SequenceEqual(Suffix);

    public void SetSuffix(ReadOnlySpan<byte> suffix) => Suffix = Parts.CopyBytes(suffix);

    public byte[] Path() => Suffix;

    public INode?[] Children() => [];

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

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

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

    public Node4(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

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

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

    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++;
    }

    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;
    }

    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;
            }
        }
    }

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

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

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

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

    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();

    public Node10(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

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

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

    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++;
    }

    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;
    }

    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;
            }
        }
    }

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

    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;
    }

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

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

    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();

    public Node16(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

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

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

    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++;
    }

    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;
    }

    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;
            }
        }
    }

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

    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;
    }

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

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

    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();

    public Node48(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

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

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

    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++;
    }

    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);
    }

    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--;
    }

    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;
    }

    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;
    }

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

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

    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();

    public Node256(ReadOnlySpan<byte> prefix)
    {
        SetPrefix(prefix);
    }

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

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

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

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

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

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

    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;
    }

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

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

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

#endregion
