MarketAlly.AIPlugin.Extensions/MarketAlly.AIPlugin.DevOps/Security/RateLimiter.cs

using System;
using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;

namespace MarketAlly.AIPlugin.DevOps.Security
{
    public class RateLimiter
    {
        private readonly ConcurrentDictionary<string, TokenBucket> _buckets;
        private readonly AuditLogger _auditLogger;

        public RateLimiter(AuditLogger auditLogger = null)
        {
            _buckets = new ConcurrentDictionary<string, TokenBucket>();
            _auditLogger = auditLogger;
        }

        public async Task<bool> TryExecuteAsync(string clientId, int tokensRequired = 1, 
            int maxTokens = 100, TimeSpan? refillInterval = null)
        {
            var interval = refillInterval ?? TimeSpan.FromMinutes(1);
            var bucket = _buckets.GetOrAdd(clientId, _ => new TokenBucket(maxTokens, interval));

            var allowed = bucket.TryConsume(tokensRequired);

            await _auditLogger?.LogSecurityEventAsync(new SecurityAuditEvent
            {
                EventType = SecurityEventType.PermissionChecked,
                Severity = allowed ? SecuritySeverity.Low : SecuritySeverity.Medium,
                Source = nameof(RateLimiter),
                UserId = clientId,
                Details = $"Rate limit check: {(allowed ? "ALLOWED" : "DENIED")}",
                Metadata = new()
                {
                    ["clientId"] = clientId,
                    ["tokensRequired"] = tokensRequired,
                    ["tokensAvailable"] = bucket.AvailableTokens,
                    ["allowed"] = allowed
                }
            });

            return allowed;
        }

        public void ClearClient(string clientId)
        {
            _buckets.TryRemove(clientId, out _);
        }

        public void ClearAll()
        {
            _buckets.Clear();
        }
    }

    public class TokenBucket
    {
        private readonly int _maxTokens;
        private readonly TimeSpan _refillInterval;
        private readonly object _lock = new();
        private int _availableTokens;
        private DateTime _lastRefill;

        public TokenBucket(int maxTokens, TimeSpan refillInterval)
        {
            _maxTokens = maxTokens;
            _refillInterval = refillInterval;
            _availableTokens = maxTokens;
            _lastRefill = DateTime.UtcNow;
        }

        public int AvailableTokens
        {
            get
            {
                lock (_lock)
                {
                    Refill();
                    return _availableTokens;
                }
            }
        }

        public bool TryConsume(int tokens)
        {
            lock (_lock)
            {
                Refill();
                
                if (_availableTokens >= tokens)
                {
                    _availableTokens -= tokens;
                    return true;
                }
                
                return false;
            }
        }

        private void Refill()
        {
            var now = DateTime.UtcNow;
            var timeSinceLastRefill = now - _lastRefill;
            
            if (timeSinceLastRefill >= _refillInterval)
            {
                var refillCycles = (int)(timeSinceLastRefill.TotalMilliseconds / _refillInterval.TotalMilliseconds);
                var tokensToAdd = refillCycles * _maxTokens;
                
                _availableTokens = Math.Min(_maxTokens, _availableTokens + tokensToAdd);
                _lastRefill = now;
            }
        }
    }
}