// Go: consumer.go (rateLimitBps config, rate limiting in consumer delivery)
namespace NATS.Server.JetStream.Consumers;

/// <summary>
/// Token bucket rate limiter for consumer message delivery.
/// Tokens refill at a configurable rate (bytes per second).
/// Each message consumes tokens equal to its payload size.
/// Go reference: consumer.go rate limiting via rateLimitBps config.
/// </summary>
public sealed class TokenBucketRateLimiter
{
    private readonly object _lock = new();
    private double _tokens;
    private double _maxTokens;
    private double _refillRate; // tokens per millisecond
    private DateTime _lastRefill;

    /// <summary>
    /// Creates a rate limiter with the specified rate in bytes per second.
    /// </summary>
    /// <param name="bytesPerSecond">Maximum bytes per second. 0 = unlimited.</param>
    /// <param name="burstSize">Maximum burst size in bytes. Defaults to 2x rate.</param>
    public TokenBucketRateLimiter(long bytesPerSecond, long burstSize = 0)
    {
        BytesPerSecond = bytesPerSecond;
        _refillRate = bytesPerSecond / 1000.0; // tokens per ms
        _maxTokens = burstSize > 0 ? burstSize : bytesPerSecond * 2;
        _tokens = _maxTokens; // Start full
        _lastRefill = DateTime.UtcNow;
    }

    /// <summary>Configured rate in bytes per second.</summary>
    public long BytesPerSecond { get; private set; }

    /// <summary>Current available tokens (approximate).</summary>
    public double AvailableTokens
    {
        get
        {
            lock (_lock)
            {
                Refill();
                return _tokens;
            }
        }
    }

    /// <summary>
    /// Tries to consume the specified number of tokens (bytes).
    /// Returns true if tokens were available (message can be sent).
    /// Returns false if not enough tokens (caller should wait).
    /// </summary>
    public bool TryConsume(long bytes)
    {
        if (BytesPerSecond <= 0) return true; // Unlimited

        lock (_lock)
        {
            Refill();
            if (_tokens >= bytes)
            {
                _tokens -= bytes;
                return true;
            }
            return false;
        }
    }

    /// <summary>
    /// Returns the estimated wait time until enough tokens are available.
    /// </summary>
    public TimeSpan EstimateWait(long bytes)
    {
        if (BytesPerSecond <= 0) return TimeSpan.Zero;

        lock (_lock)
        {
            Refill();
            if (_tokens >= bytes) return TimeSpan.Zero;

            var deficit = bytes - _tokens;
            var waitMs = deficit / _refillRate;
            return TimeSpan.FromMilliseconds(waitMs);
        }
    }

    /// <summary>
    /// Waits until enough tokens are available, then consumes them.
    /// </summary>
    public async ValueTask WaitForTokensAsync(long bytes, CancellationToken ct = default)
    {
        if (BytesPerSecond <= 0) return;

        while (!ct.IsCancellationRequested)
        {
            if (TryConsume(bytes)) return;

            var wait = EstimateWait(bytes);
            if (wait > TimeSpan.Zero)
                await Task.Delay(wait, ct).ConfigureAwait(false);
        }

        ct.ThrowIfCancellationRequested();
    }

    /// <summary>
    /// Updates the rate dynamically.
    /// Go reference: consumer.go — rate can change on config update.
    /// </summary>
    public void UpdateRate(long bytesPerSecond, long burstSize = 0)
    {
        lock (_lock)
        {
            BytesPerSecond = bytesPerSecond;
            _refillRate = bytesPerSecond / 1000.0;
            _maxTokens = burstSize > 0 ? burstSize : bytesPerSecond * 2;
            _tokens = Math.Min(_tokens, _maxTokens);
        }
    }

    private void Refill()
    {
        var now = DateTime.UtcNow;
        var elapsed = (now - _lastRefill).TotalMilliseconds;
        if (elapsed <= 0) return;

        _tokens = Math.Min(_maxTokens, _tokens + elapsed * _refillRate);
        _lastRefill = now;
    }
}
