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

// Internal node interface for the adaptive radix trie.
internal interface ISubjectTreeNode<T>
{
    bool IsLeaf { get; }
    byte[] Prefix { get; }
    void AddChild(byte key, ISubjectTreeNode<T> child);
    ISubjectTreeNode<T>? FindChild(byte key);
    void DeleteChild(byte key);
    bool IsFull { get; }
    ISubjectTreeNode<T> Grow();
    ISubjectTreeNode<T>? Shrink();
    ISubjectTreeNode<T>[] Children();
    int NumChildren { get; }
    byte[] Path { get; }
    (byte[][] remainingParts, bool matched) MatchParts(byte[][] parts);
    string Kind { get; }
}

// Base class for non-leaf nodes, holding prefix and child count.
internal abstract class SubjectTreeMeta<T> : ISubjectTreeNode<T>
{
    protected byte[] _prefix;
    protected int _size;

    protected SubjectTreeMeta(byte[] prefix)
    {
        _prefix = SubjectTreeParts.CopyBytes(prefix);
    }

    public bool IsLeaf => false;
    public byte[] Prefix => _prefix;
    public int NumChildren => _size;
    public byte[] Path => _prefix;

    public void SetPrefix(byte[] prefix)
    {
        _prefix = SubjectTreeParts.CopyBytes(prefix);
    }

    public (byte[][] remainingParts, bool matched) MatchParts(byte[][] parts)
        => SubjectTreeParts.MatchParts(parts, _prefix);

    public abstract void AddChild(byte key, ISubjectTreeNode<T> child);
    public abstract ISubjectTreeNode<T>? FindChild(byte key);
    public abstract void DeleteChild(byte key);
    public abstract bool IsFull { get; }
    public abstract ISubjectTreeNode<T> Grow();
    public abstract ISubjectTreeNode<T>? Shrink();
    public abstract ISubjectTreeNode<T>[] Children();
    public abstract string Kind { get; }
}

// Leaf node storing the terminal value plus a suffix byte[].
internal sealed class SubjectTreeLeaf<T> : ISubjectTreeNode<T>
{
    public T Value;
    public byte[] Suffix;

    public SubjectTreeLeaf(byte[] suffix, T value)
    {
        Suffix = SubjectTreeParts.CopyBytes(suffix);
        Value = value;
    }

    public bool IsLeaf => true;
    public byte[] Prefix => Array.Empty<byte>();
    public int NumChildren => 0;
    public byte[] Path => Suffix;
    public string Kind => "LEAF";

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

    public void SetSuffix(byte[] suffix)
        => Suffix = SubjectTreeParts.CopyBytes(suffix);

    public bool IsFull => true;

    public (byte[][] remainingParts, bool matched) MatchParts(byte[][] parts)
        => SubjectTreeParts.MatchParts(parts, Suffix);

    // Leaf nodes do not support child operations.
    public void AddChild(byte key, ISubjectTreeNode<T> child)
        => throw new InvalidOperationException("AddChild called on leaf");

    public ISubjectTreeNode<T>? FindChild(byte key)
        => throw new InvalidOperationException("FindChild called on leaf");

    public void DeleteChild(byte key)
        => throw new InvalidOperationException("DeleteChild called on leaf");

    public ISubjectTreeNode<T> Grow()
        => throw new InvalidOperationException("Grow called on leaf");

    public ISubjectTreeNode<T>? Shrink()
        => throw new InvalidOperationException("Shrink called on leaf");

    public ISubjectTreeNode<T>[] Children()
        => Array.Empty<ISubjectTreeNode<T>>();
}

// Node with up to 4 children (keys + children arrays, unsorted).
internal sealed class SubjectTreeNode4<T> : SubjectTreeMeta<T>
{
    private readonly byte[] _keys = new byte[4];
    private readonly ISubjectTreeNode<T>?[] _children = new ISubjectTreeNode<T>?[4];

    public SubjectTreeNode4(byte[] prefix) : base(prefix) { }

    public override string Kind => "NODE4";

    public override void AddChild(byte key, ISubjectTreeNode<T> child)
    {
        if (_size >= 4) throw new InvalidOperationException("node4 full!");
        _keys[_size] = key;
        _children[_size] = child;
        _size++;
    }

    public override ISubjectTreeNode<T>? FindChild(byte key)
    {
        for (var i = 0; i < _size; i++)
        {
            if (_keys[i] == key) return _children[i];
        }
        return null;
    }

    public override void DeleteChild(byte key)
    {
        for (var i = 0; i < _size; i++)
        {
            if (_keys[i] == key)
            {
                var last = _size - 1;
                if (i < last)
                {
                    _keys[i] = _keys[last];
                    _children[i] = _children[last];
                }
                _keys[last] = 0;
                _children[last] = null;
                _size--;
                return;
            }
        }
    }

    public override bool IsFull => _size >= 4;

    public override ISubjectTreeNode<T> Grow()
    {
        var nn = new SubjectTreeNode10<T>(_prefix);
        for (var i = 0; i < 4; i++)
            nn.AddChild(_keys[i], _children[i]!);
        return nn;
    }

    public override ISubjectTreeNode<T>? Shrink()
    {
        if (_size == 1) return _children[0];
        return null;
    }

    public override ISubjectTreeNode<T>[] Children()
    {
        var result = new ISubjectTreeNode<T>[_size];
        for (var i = 0; i < _size; i++)
            result[i] = _children[i]!;
        return result;
    }

    // Internal access for tests.
    internal byte GetKey(int index) => _keys[index];
    internal ISubjectTreeNode<T>? GetChild(int index) => _children[index];
}

// Node with up to 10 children (for numeric token segments).
internal sealed class SubjectTreeNode10<T> : SubjectTreeMeta<T>
{
    private readonly byte[] _keys = new byte[10];
    private readonly ISubjectTreeNode<T>?[] _children = new ISubjectTreeNode<T>?[10];

    public SubjectTreeNode10(byte[] prefix) : base(prefix) { }

    public override string Kind => "NODE10";

    public override void AddChild(byte key, ISubjectTreeNode<T> child)
    {
        if (_size >= 10) throw new InvalidOperationException("node10 full!");
        _keys[_size] = key;
        _children[_size] = child;
        _size++;
    }

    public override ISubjectTreeNode<T>? FindChild(byte key)
    {
        for (var i = 0; i < _size; i++)
        {
            if (_keys[i] == key) return _children[i];
        }
        return null;
    }

    public override void DeleteChild(byte key)
    {
        for (var i = 0; i < _size; i++)
        {
            if (_keys[i] == key)
            {
                var last = _size - 1;
                if (i < last)
                {
                    _keys[i] = _keys[last];
                    _children[i] = _children[last];
                }
                _keys[last] = 0;
                _children[last] = null;
                _size--;
                return;
            }
        }
    }

    public override bool IsFull => _size >= 10;

    public override ISubjectTreeNode<T> Grow()
    {
        var nn = new SubjectTreeNode16<T>(_prefix);
        for (var i = 0; i < _size; i++)
            nn.AddChild(_keys[i], _children[i]!);
        return nn;
    }

    public override ISubjectTreeNode<T>? Shrink()
    {
        if (_size > 4) return null;
        var nn = new SubjectTreeNode4<T>(Array.Empty<byte>());
        for (var i = 0; i < _size; i++)
            nn.AddChild(_keys[i], _children[i]!);
        return nn;
    }

    public override ISubjectTreeNode<T>[] Children()
    {
        var result = new ISubjectTreeNode<T>[_size];
        for (var i = 0; i < _size; i++)
            result[i] = _children[i]!;
        return result;
    }
}

// Node with up to 16 children.
internal sealed class SubjectTreeNode16<T> : SubjectTreeMeta<T>
{
    private readonly byte[] _keys = new byte[16];
    private readonly ISubjectTreeNode<T>?[] _children = new ISubjectTreeNode<T>?[16];

    public SubjectTreeNode16(byte[] prefix) : base(prefix) { }

    public override string Kind => "NODE16";

    public override void AddChild(byte key, ISubjectTreeNode<T> child)
    {
        if (_size >= 16) throw new InvalidOperationException("node16 full!");
        _keys[_size] = key;
        _children[_size] = child;
        _size++;
    }

    public override ISubjectTreeNode<T>? FindChild(byte key)
    {
        for (var i = 0; i < _size; i++)
        {
            if (_keys[i] == key) return _children[i];
        }
        return null;
    }

    public override void DeleteChild(byte key)
    {
        for (var i = 0; i < _size; i++)
        {
            if (_keys[i] == key)
            {
                var last = _size - 1;
                if (i < last)
                {
                    _keys[i] = _keys[last];
                    _children[i] = _children[last];
                }
                _keys[last] = 0;
                _children[last] = null;
                _size--;
                return;
            }
        }
    }

    public override bool IsFull => _size >= 16;

    public override ISubjectTreeNode<T> Grow()
    {
        var nn = new SubjectTreeNode48<T>(_prefix);
        for (var i = 0; i < _size; i++)
            nn.AddChild(_keys[i], _children[i]!);
        return nn;
    }

    public override ISubjectTreeNode<T>? Shrink()
    {
        if (_size > 10) return null;
        var nn = new SubjectTreeNode10<T>(Array.Empty<byte>());
        for (var i = 0; i < _size; i++)
            nn.AddChild(_keys[i], _children[i]!);
        return nn;
    }

    public override ISubjectTreeNode<T>[] Children()
    {
        var result = new ISubjectTreeNode<T>[_size];
        for (var i = 0; i < _size; i++)
            result[i] = _children[i]!;
        return result;
    }
}

// Node with up to 48 children, using a 256-byte key index (1-indexed, 0 means empty).
internal sealed class SubjectTreeNode48<T> : SubjectTreeMeta<T>
{
    // _keyIndex[byte] = 1-based index into _children; 0 means no entry.
    private readonly byte[] _keyIndex = new byte[256];
    private readonly ISubjectTreeNode<T>?[] _children = new ISubjectTreeNode<T>?[48];

    public SubjectTreeNode48(byte[] prefix) : base(prefix) { }

    public override string Kind => "NODE48";

    public override void AddChild(byte key, ISubjectTreeNode<T> child)
    {
        if (_size >= 48) throw new InvalidOperationException("node48 full!");
        _children[_size] = child;
        _keyIndex[key] = (byte)(_size + 1); // 1-indexed
        _size++;
    }

    public override ISubjectTreeNode<T>? FindChild(byte key)
    {
        var i = _keyIndex[key];
        if (i == 0) return null;
        return _children[i - 1];
    }

    public override void DeleteChild(byte key)
    {
        var i = _keyIndex[key];
        if (i == 0) return;
        i--; // Convert from 1-indexed
        var last = _size - 1;
        if (i < last)
        {
            _children[i] = _children[last];
            // Find which key index points to 'last' and redirect it to 'i'.
            for (var ic = 0; ic < 256; ic++)
            {
                if (_keyIndex[ic] == last + 1)
                {
                    _keyIndex[ic] = (byte)(i + 1);
                    break;
                }
            }
        }
        _children[last] = null;
        _keyIndex[key] = 0;
        _size--;
    }

    public override bool IsFull => _size >= 48;

    public override ISubjectTreeNode<T> Grow()
    {
        var nn = new SubjectTreeNode256<T>(_prefix);
        for (var c = 0; c < 256; c++)
        {
            var i = _keyIndex[c];
            if (i > 0)
                nn.AddChild((byte)c, _children[i - 1]!);
        }
        return nn;
    }

    public override ISubjectTreeNode<T>? Shrink()
    {
        if (_size > 16) return null;
        var nn = new SubjectTreeNode16<T>(Array.Empty<byte>());
        for (var c = 0; c < 256; c++)
        {
            var i = _keyIndex[c];
            if (i > 0)
                nn.AddChild((byte)c, _children[i - 1]!);
        }
        return nn;
    }

    public override ISubjectTreeNode<T>[] Children()
    {
        var result = new ISubjectTreeNode<T>[_size];
        var idx = 0;
        for (var i = 0; i < _size; i++)
        {
            if (_children[i] != null)
                result[idx++] = _children[i]!;
        }
        return result[..idx];
    }

    // Internal access for tests.
    internal byte GetKeyIndex(int key) => _keyIndex[key];
    internal ISubjectTreeNode<T>? GetChildAt(int index) => _children[index];
}

// Node with 256 children, indexed directly by byte value.
internal sealed class SubjectTreeNode256<T> : SubjectTreeMeta<T>
{
    private readonly ISubjectTreeNode<T>?[] _children = new ISubjectTreeNode<T>?[256];

    public SubjectTreeNode256(byte[] prefix) : base(prefix) { }

    public override string Kind => "NODE256";

    public override void AddChild(byte key, ISubjectTreeNode<T> child)
    {
        _children[key] = child;
        _size++;
    }

    public override ISubjectTreeNode<T>? FindChild(byte key)
        => _children[key];

    public override void DeleteChild(byte key)
    {
        if (_children[key] != null)
        {
            _children[key] = null;
            _size--;
        }
    }

    public override bool IsFull => false;

    public override ISubjectTreeNode<T> Grow()
        => throw new InvalidOperationException("Grow cannot be called on node256");

    public override ISubjectTreeNode<T>? Shrink()
    {
        if (_size > 48) return null;
        var nn = new SubjectTreeNode48<T>(Array.Empty<byte>());
        for (var c = 0; c < 256; c++)
        {
            if (_children[c] != null)
                nn.AddChild((byte)c, _children[c]!);
        }
        return nn;
    }

    public override ISubjectTreeNode<T>[] Children()
        => _children.Where(c => c != null).Select(c => c!).ToArray();
}
