// Go reference: server/gsl/gsl.go
// Trie-based generic subject list with wildcard support for NATS subject matching.

namespace NATS.Server.Internal.Gsl;

/// <summary>
/// Sublist related errors.
/// </summary>
public static class GslErrors
{
    public static readonly InvalidOperationException InvalidSubject = new("gsl: invalid subject");
    public static readonly KeyNotFoundException NotFound = new("gsl: no matches found");
}

/// <summary>
/// A level represents a group of nodes and special pointers to wildcard nodes.
/// Go reference: server/gsl/gsl.go level struct
/// </summary>
internal sealed class Level<T> where T : IEquatable<T>
{
    public Dictionary<string, Node<T>> Nodes { get; } = new();
    public Node<T>? Pwc { get; set; } // partial wildcard '*'
    public Node<T>? Fwc { get; set; } // full wildcard '>'

    public int NumNodes()
    {
        var num = Nodes.Count;
        if (Pwc is not null) num++;
        if (Fwc is not null) num++;
        return num;
    }

    /// <summary>
    /// Prune an empty node from the tree.
    /// Go reference: server/gsl/gsl.go pruneNode
    /// </summary>
    public void PruneNode(Node<T> n, string token)
    {
        if (ReferenceEquals(n, Fwc))
            Fwc = null;
        else if (ReferenceEquals(n, Pwc))
            Pwc = null;
        else
            Nodes.Remove(token);
    }
}

/// <summary>
/// A node contains subscriptions and a pointer to the next level.
/// Go reference: server/gsl/gsl.go node struct
/// </summary>
internal sealed class Node<T> where T : IEquatable<T>
{
    public Level<T>? Next { get; set; }
    public Dictionary<T, string> Subs { get; } = new(); // value -> subject

    /// <summary>
    /// Returns true if the node has no subscriptions and no children.
    /// Go reference: server/gsl/gsl.go isEmpty
    /// </summary>
    public bool IsEmpty() => Subs.Count == 0 && (Next is null || Next.NumNodes() == 0);
}

/// <summary>
/// Tracks descent into levels during removal for pruning.
/// Go reference: server/gsl/gsl.go lnt struct
/// </summary>
internal readonly record struct Lnt<T>(Level<T> L, Node<T> N, string T_) where T : IEquatable<T>;

/// <summary>
/// A GenericSubjectList stores and efficiently retrieves subscriptions using a trie.
/// Supports wildcard subjects: '*' matches a single token, '>' matches one or more tokens.
/// Thread-safe via ReaderWriterLockSlim.
/// Go reference: server/gsl/gsl.go GenericSublist
/// </summary>
public class GenericSubjectList<T> where T : IEquatable<T>
{
    private const char Pwc = '*';
    private const char Fwc = '>';
    private const char Btsep = '.';

    private readonly ReaderWriterLockSlim _lock = new();
    private readonly Level<T> _root = new();
    private uint _count;

    /// <summary>
    /// Returns the number of subscriptions.
    /// Go reference: server/gsl/gsl.go Count
    /// </summary>
    public uint Count
    {
        get
        {
            _lock.EnterReadLock();
            try
            {
                return _count;
            }
            finally
            {
                _lock.ExitReadLock();
            }
        }
    }

    /// <summary>
    /// Insert adds a subscription into the sublist.
    /// Go reference: server/gsl/gsl.go Insert
    /// </summary>
    public void Insert(string subject, T value)
    {
        _lock.EnterWriteLock();
        try
        {
            var sfwc = false;
            Node<T>? n = null;
            var l = _root;

            foreach (var token in TokenizeSubject(subject))
            {
                var lt = token.Length;
                if (lt == 0 || sfwc)
                    throw GslErrors.InvalidSubject;

                if (lt > 1)
                {
                    l.Nodes.TryGetValue(token, out n);
                }
                else
                {
                    switch (token[0])
                    {
                        case Pwc:
                            n = l.Pwc;
                            break;
                        case Fwc:
                            n = l.Fwc;
                            sfwc = true;
                            break;
                        default:
                            l.Nodes.TryGetValue(token, out n);
                            break;
                    }
                }

                if (n is null)
                {
                    n = new Node<T>();
                    if (lt > 1)
                    {
                        l.Nodes[token] = n;
                    }
                    else
                    {
                        switch (token[0])
                        {
                            case Pwc:
                                l.Pwc = n;
                                break;
                            case Fwc:
                                l.Fwc = n;
                                break;
                            default:
                                l.Nodes[token] = n;
                                break;
                        }
                    }
                }

                n.Next ??= new Level<T>();
                l = n.Next;
            }

            // n should never be null here if subject was valid (non-empty)
            n!.Subs[value] = subject;
            _count++;
        }
        finally
        {
            _lock.ExitWriteLock();
        }
    }

    /// <summary>
    /// Remove will remove a subscription.
    /// Go reference: server/gsl/gsl.go Remove
    /// </summary>
    public void Remove(string subject, T value)
    {
        _lock.EnterWriteLock();
        try
        {
            RemoveInternal(subject, value);
        }
        finally
        {
            _lock.ExitWriteLock();
        }
    }

    /// <summary>
    /// Match will match all entries to the literal subject and invoke the callback for each.
    /// Go reference: server/gsl/gsl.go Match
    /// </summary>
    public void Match(string subject, Action<T> callback)
    {
        MatchInternal(subject, callback, doLock: true);
    }

    /// <summary>
    /// MatchBytes will match all entries to the literal subject (as bytes) and invoke the callback for each.
    /// Go reference: server/gsl/gsl.go MatchBytes
    /// </summary>
    public void MatchBytes(ReadOnlySpan<byte> subject, Action<T> callback)
    {
        // Convert bytes to string then delegate
        var subjectStr = System.Text.Encoding.UTF8.GetString(subject);
        MatchInternal(subjectStr, callback, doLock: true);
    }

    /// <summary>
    /// HasInterest will return whether or not there is any interest in the subject.
    /// Go reference: server/gsl/gsl.go HasInterest
    /// </summary>
    public bool HasInterest(string subject)
    {
        return HasInterestInternal(subject, doLock: true, np: null);
    }

    /// <summary>
    /// NumInterest will return the number of subs interested in the subject.
    /// Go reference: server/gsl/gsl.go NumInterest
    /// </summary>
    public int NumInterest(string subject)
    {
        var np = new int[1]; // use array to pass by reference
        HasInterestInternal(subject, doLock: true, np: np);
        return np[0];
    }

    /// <summary>
    /// HasInterestStartingIn is a helper for subject tree intersection.
    /// Go reference: server/gsl/gsl.go HasInterestStartingIn
    /// </summary>
    public bool HasInterestStartingIn(string subject)
    {
        _lock.EnterReadLock();
        try
        {
            Span<string> tokenBuffer = new string[64];
            var tokens = TokenizeSubjectIntoSpan(subject, tokenBuffer);
            return HasInterestStartingInLevel(_root, tokens);
        }
        finally
        {
            _lock.ExitReadLock();
        }
    }

    /// <summary>
    /// Returns the maximum number of levels in the trie. Used for testing.
    /// Go reference: server/gsl/gsl.go numLevels
    /// </summary>
    internal int NumLevels()
    {
        return VisitLevel(_root, 0);
    }

    // --- Private implementation ---

    /// <summary>
    /// Go reference: server/gsl/gsl.go match
    /// </summary>
    private void MatchInternal(string subject, Action<T> callback, bool doLock)
    {
        Span<string> tokenBuffer = new string[32];
        var tokens = TokenizeSubjectForMatch(subject, tokenBuffer);
        if (tokens.Length == 0)
            return;

        if (doLock) _lock.EnterReadLock();
        try
        {
            MatchLevel(_root, tokens, callback);
        }
        finally
        {
            if (doLock) _lock.ExitReadLock();
        }
    }

    /// <summary>
    /// Go reference: server/gsl/gsl.go hasInterest
    /// </summary>
    private bool HasInterestInternal(string subject, bool doLock, int[]? np)
    {
        Span<string> tokenBuffer = new string[32];
        var tokens = TokenizeSubjectForMatch(subject, tokenBuffer);
        if (tokens.Length == 0)
            return false;

        if (doLock) _lock.EnterReadLock();
        try
        {
            return MatchLevelForAny(_root, tokens, np);
        }
        finally
        {
            if (doLock) _lock.ExitReadLock();
        }
    }

    /// <summary>
    /// Tokenize a subject for match/hasInterest. Returns empty span for invalid subjects
    /// (empty tokens or trailing separator).
    /// Go reference: server/gsl/gsl.go match (tokenization section)
    /// </summary>
    private static ReadOnlySpan<string> TokenizeSubjectForMatch(string subject, Span<string> buffer)
    {
        var count = 0;
        var start = 0;
        for (var i = 0; i < subject.Length; i++)
        {
            if (subject[i] == Btsep)
            {
                if (i - start == 0)
                    return ReadOnlySpan<string>.Empty; // empty token
                if (count >= buffer.Length)
                    return ReadOnlySpan<string>.Empty;
                buffer[count++] = subject[start..i];
                start = i + 1;
            }
        }

        if (start >= subject.Length)
            return ReadOnlySpan<string>.Empty; // trailing separator

        if (count >= buffer.Length)
            return ReadOnlySpan<string>.Empty;
        buffer[count++] = subject[start..];
        return buffer[..count];
    }

    /// <summary>
    /// Tokenize a subject into a span (does not validate empty tokens).
    /// Go reference: server/gsl/gsl.go tokenizeSubjectIntoSlice
    /// </summary>
    private static ReadOnlySpan<string> TokenizeSubjectIntoSpan(string subject, Span<string> buffer)
    {
        var count = 0;
        var start = 0;
        for (var i = 0; i < subject.Length; i++)
        {
            if (subject[i] == Btsep)
            {
                if (count >= buffer.Length) break;
                buffer[count++] = subject[start..i];
                start = i + 1;
            }
        }

        if (count < buffer.Length)
            buffer[count++] = subject[start..];
        return buffer[..count];
    }

    /// <summary>
    /// Recursively descend into the trie to match subscriptions.
    /// Go reference: server/gsl/gsl.go matchLevel
    /// </summary>
    private static void MatchLevel(Level<T>? l, ReadOnlySpan<string> toks, Action<T> cb)
    {
        Node<T>? pwc = null;
        Node<T>? n = null;
        for (var i = 0; i < toks.Length; i++)
        {
            if (l is null) return;

            if (l.Fwc is not null)
                CallbacksForResults(l.Fwc, cb);

            pwc = l.Pwc;
            if (pwc is not null)
                MatchLevel(pwc.Next, toks[(i + 1)..], cb);

            l.Nodes.TryGetValue(toks[i], out n);
            l = n?.Next;
        }

        if (n is not null)
            CallbacksForResults(n, cb);
        if (pwc is not null)
            CallbacksForResults(pwc, cb);
    }

    /// <summary>
    /// Recursively check if any subscription matches (optimization over full Match).
    /// Go reference: server/gsl/gsl.go matchLevelForAny
    /// </summary>
    private static bool MatchLevelForAny(Level<T>? l, ReadOnlySpan<string> toks, int[]? np)
    {
        Node<T>? pwc = null;
        Node<T>? n = null;
        for (var i = 0; i < toks.Length; i++)
        {
            if (l is null) return false;

            if (l.Fwc is not null)
            {
                if (np is not null)
                    np[0] += l.Fwc.Subs.Count;
                return true;
            }

            pwc = l.Pwc;
            if (pwc is not null)
            {
                if (MatchLevelForAny(pwc.Next, toks[(i + 1)..], np))
                    return true;
            }

            l.Nodes.TryGetValue(toks[i], out n);
            l = n?.Next;
        }

        if (n is not null)
        {
            if (np is not null)
                np[0] += n.Subs.Count;
            if (n.Subs.Count > 0)
                return true;
        }

        if (pwc is not null)
        {
            if (np is not null)
                np[0] += pwc.Subs.Count;
            return pwc.Subs.Count > 0;
        }

        return false;
    }

    /// <summary>
    /// Invoke callback for each subscription in a node.
    /// Go reference: server/gsl/gsl.go callbacksForResults
    /// </summary>
    private static void CallbacksForResults(Node<T> n, Action<T> cb)
    {
        foreach (var sub in n.Subs.Keys)
            cb(sub);
    }

    /// <summary>
    /// Internal remove with lock already held.
    /// Go reference: server/gsl/gsl.go remove
    /// </summary>
    private void RemoveInternal(string subject, T value)
    {
        var sfwc = false;
        Node<T>? n = null;
        Level<T>? l = _root;

        // Track levels for pruning
        Span<Lnt<T>> levelsBuffer = new Lnt<T>[32];
        var levelCount = 0;

        foreach (var token in TokenizeSubject(subject))
        {
            var lt = token.Length;
            if (lt == 0 || sfwc)
                throw GslErrors.InvalidSubject;

            if (l is null)
                throw GslErrors.NotFound;

            if (lt > 1)
            {
                l.Nodes.TryGetValue(token, out n);
            }
            else
            {
                switch (token[0])
                {
                    case Pwc:
                        n = l.Pwc;
                        break;
                    case Fwc:
                        n = l.Fwc;
                        sfwc = true;
                        break;
                    default:
                        l.Nodes.TryGetValue(token, out n);
                        break;
                }
            }

            if (n is not null)
            {
                levelsBuffer[levelCount++] = new Lnt<T>(l, n, token);
                l = n.Next;
            }
            else
            {
                l = null;
            }
        }

        if (!RemoveFromNode(n, value))
            throw GslErrors.NotFound;

        _count--;

        // Prune empty nodes
        for (var i = levelCount - 1; i >= 0; i--)
        {
            var lnt = levelsBuffer[i];
            if (lnt.N.IsEmpty())
                lnt.L.PruneNode(lnt.N, lnt.T_);
        }
    }

    /// <summary>
    /// Remove the value from the given node.
    /// Go reference: server/gsl/gsl.go removeFromNode
    /// </summary>
    private static bool RemoveFromNode(Node<T>? n, T value)
    {
        if (n is null) return false;
        return n.Subs.Remove(value);
    }

    /// <summary>
    /// Recursively check if there is interest starting at a prefix.
    /// Go reference: server/gsl/gsl.go hasInterestStartingIn
    /// </summary>
    private static bool HasInterestStartingInLevel(Level<T>? l, ReadOnlySpan<string> tokens)
    {
        if (l is null) return false;
        if (tokens.Length == 0) return true;

        var token = tokens[0];
        if (l.Fwc is not null) return true;

        var found = false;
        if (l.Pwc is not null)
            found = HasInterestStartingInLevel(l.Pwc.Next, tokens[1..]);
        if (!found && l.Nodes.TryGetValue(token, out var n))
            found = HasInterestStartingInLevel(n.Next, tokens[1..]);
        return found;
    }

    /// <summary>
    /// Visit levels recursively to compute max depth.
    /// Go reference: server/gsl/gsl.go visitLevel
    /// </summary>
    private static int VisitLevel(Level<T>? l, int depth)
    {
        if (l is null || l.NumNodes() == 0)
            return depth;

        depth++;
        var maxDepth = depth;

        foreach (var n in l.Nodes.Values)
        {
            var newDepth = VisitLevel(n.Next, depth);
            if (newDepth > maxDepth)
                maxDepth = newDepth;
        }

        if (l.Pwc is not null)
        {
            var pwcDepth = VisitLevel(l.Pwc.Next, depth);
            if (pwcDepth > maxDepth)
                maxDepth = pwcDepth;
        }

        if (l.Fwc is not null)
        {
            var fwcDepth = VisitLevel(l.Fwc.Next, depth);
            if (fwcDepth > maxDepth)
                maxDepth = fwcDepth;
        }

        return maxDepth;
    }

    /// <summary>
    /// Tokenize a subject by splitting on '.'. Returns an enumerable of tokens.
    /// Used by Insert and Remove.
    /// </summary>
    private static SplitEnumerable TokenizeSubject(string subject)
    {
        return new SplitEnumerable(subject);
    }

    /// <summary>
    /// A stack-friendly subject tokenizer that splits on '.'.
    /// </summary>
    private readonly ref struct SplitEnumerable
    {
        private readonly string _subject;

        public SplitEnumerable(string subject) => _subject = subject;

        public SplitEnumerator GetEnumerator() => new(_subject);
    }

    private ref struct SplitEnumerator
    {
        private readonly string _subject;
        private int _start;
        private bool _done;

        public SplitEnumerator(string subject)
        {
            _subject = subject;
            _start = 0;
            _done = false;
            Current = default!;
        }

        public string Current { get; private set; }

        public bool MoveNext()
        {
            if (_done) return false;

            var idx = _subject.IndexOf(Btsep, _start);
            if (idx >= 0)
            {
                Current = _subject[_start..idx];
                _start = idx + 1;
                return true;
            }

            Current = _subject[_start..];
            _done = true;
            return true;
        }
    }
}

/// <summary>
/// Empty marker used by <see cref="SimpleSubjectList"/> to mirror Go's `struct{}` payload.
/// Go reference: server/gsl/gsl.go SimpleSublist
/// </summary>
public readonly record struct SimpleSublistValue;

/// <summary>
/// SimpleSubjectList is an alias for GenericSubjectList that uses an empty marker payload,
/// useful for tracking interest only.
/// Go reference: server/gsl/gsl.go SimpleSublist
/// </summary>
public class SimpleSubjectList : GenericSubjectList<SimpleSublistValue>;
