namespace NATS.Server.Auth;

/// <summary>
/// Tracks reply subjects that a client is temporarily allowed to publish to.
/// Reference: Go client.go resp struct, setResponsePermissionIfNeeded.
/// </summary>
public sealed class ResponseTracker
{
    private readonly int _maxMsgs;      // 0 = unlimited
    private readonly TimeSpan _expires; // TimeSpan.Zero = no TTL
    private readonly Dictionary<string, (DateTime RegisteredAt, int Count)> _replies = new(StringComparer.Ordinal);
    private readonly object _lock = new();

    public ResponseTracker(int maxMsgs, TimeSpan expires)
    {
        _maxMsgs = maxMsgs;
        _expires = expires;
    }

    public int Count
    {
        get { lock (_lock) return _replies.Count; }
    }

    public void RegisterReply(string replySubject)
    {
        lock (_lock)
        {
            _replies[replySubject] = (DateTime.UtcNow, 0);
        }
    }

    public bool IsReplyAllowed(string subject)
    {
        lock (_lock)
        {
            if (!_replies.TryGetValue(subject, out var entry))
                return false;

            if (_expires > TimeSpan.Zero && DateTime.UtcNow - entry.RegisteredAt > _expires)
            {
                _replies.Remove(subject);
                return false;
            }

            var newCount = entry.Count + 1;
            if (_maxMsgs > 0 && newCount > _maxMsgs)
            {
                _replies.Remove(subject);
                return false;
            }

            _replies[subject] = (entry.RegisteredAt, newCount);
            return true;
        }
    }

    public void Prune()
    {
        lock (_lock)
        {
            if (_expires <= TimeSpan.Zero && _maxMsgs <= 0)
                return;

            var now = DateTime.UtcNow;
            var toRemove = new List<string>();
            foreach (var (key, entry) in _replies)
            {
                if (_expires > TimeSpan.Zero && now - entry.RegisteredAt > _expires)
                    toRemove.Add(key);
                else if (_maxMsgs > 0 && entry.Count >= _maxMsgs)
                    toRemove.Add(key);
            }
            foreach (var key in toRemove)
                _replies.Remove(key);
        }
    }
}