mecha10-core 0.8.0

//! Rate Limiting
//!
//! Per-topic and per-node rate limiting to prevent message flooding and ensure QoS.
//!
//! # Features
//!
//! - Token bucket algorithm for smooth rate limiting
//! - Per-topic limits
//! - Per-node limits
//! - Burst capacity
//! - Configurable policies
//! - Metrics integration
//! - Low overhead with async
//!
//! # Example
//!
//! ```rust
//! use mecha10_core::prelude::*;
//! use mecha10_core::rate_limit::{RateLimiter, RateLimit};
//!
//! # async fn example() -> Result<()> {
//! // Create rate limiter: 10 messages/sec with burst of 20
//! let limiter = RateLimiter::new(10.0, 20);
//!
//! // Check if message is allowed
//! if limiter.check_and_consume().await {
//!     // Allowed - publish message
//!     println!("Message allowed");
//! } else {
//!     // Rate limited - drop or queue
//!     println!("Rate limited!");
//! }
//! # Ok(())
//! # }
//! ```

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::RwLock;

pub struct RateLimiter {
    tokens: Arc<RwLock<f64>>,
    capacity: f64,
    refill_rate: f64,
    last_refill: Arc<RwLock<Instant>>,
}

impl RateLimiter {
    /// Create a new rate limiter
    ///
    /// # Arguments
    ///
    /// * `rate` - Messages per second
    /// * `burst` - Maximum burst size
    ///
    /// # Example
    ///
    /// ```rust
    /// use mecha10_core::rate_limit::RateLimiter;
    ///
    /// // 100 msg/sec with burst of 200
    /// let limiter = RateLimiter::new(100.0, 200);
    /// ```
    pub fn new(rate: f64, burst: usize) -> Self {
        Self {
            tokens: Arc::new(RwLock::new(burst as f64)),
            capacity: burst as f64,
            refill_rate: rate,
            last_refill: Arc::new(RwLock::new(Instant::now())),
        }
    }

    /// Check if request is allowed (does not consume token)
    pub async fn check(&self) -> bool {
        self.refill().await;
        let tokens = self.tokens.read().await;
        *tokens >= 1.0
    }

    /// Check and consume a token if available
    pub async fn check_and_consume(&self) -> bool {
        self.refill().await;

        let mut tokens = self.tokens.write().await;
        if *tokens >= 1.0 {
            *tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Get current token count
    pub async fn available_tokens(&self) -> f64 {
        self.refill().await;
        *self.tokens.read().await
    }

    /// Refill tokens based on elapsed time
    async fn refill(&self) {
        let now = Instant::now();
        let mut last_refill = self.last_refill.write().await;
        let elapsed = now.duration_since(*last_refill).as_secs_f64();

        if elapsed > 0.0 {
            let tokens_to_add = elapsed * self.refill_rate;
            let mut tokens = self.tokens.write().await;
            *tokens = (*tokens + tokens_to_add).min(self.capacity);
            *last_refill = now;
        }
    }

    /// Reset the rate limiter (fill bucket)
    pub async fn reset(&self) {
        let mut tokens = self.tokens.write().await;
        *tokens = self.capacity;
        let mut last_refill = self.last_refill.write().await;
        *last_refill = Instant::now();
    }
}

#[derive(Debug, Clone)]
pub struct RateLimit {
    pub rate: f64,
    pub burst: usize,
}

impl RateLimit {
    /// Create a new rate limit policy
    pub fn new(rate: f64, burst: usize) -> Self {
        Self { rate, burst }
    }

    /// Create rate limit from messages per second (burst = rate * 2)
    pub fn per_second(rate: f64) -> Self {
        Self {
            rate,
            burst: (rate * 2.0) as usize,
        }
    }

    /// Create rate limit from messages per minute
    pub fn per_minute(rate: f64) -> Self {
        let per_sec = rate / 60.0;
        Self {
            rate: per_sec,
            burst: (per_sec * 2.0) as usize,
        }
    }

    /// No limit
    pub fn unlimited() -> Self {
        Self {
            rate: f64::MAX,
            burst: usize::MAX,
        }
    }
}

pub struct MultiKeyRateLimiter {
    limiters: Arc<RwLock<HashMap<String, RateLimiter>>>,
    default_limit: RateLimit,
}

impl MultiKeyRateLimiter {
    /// Create a new multi-key rate limiter with default unlimited
    pub fn new() -> Self {
        Self {
            limiters: Arc::new(RwLock::new(HashMap::new())),
            default_limit: RateLimit::unlimited(),
        }
    }

    /// Create with a default limit for all keys
    pub fn with_default(default: RateLimit) -> Self {
        Self {
            limiters: Arc::new(RwLock::new(HashMap::new())),
            default_limit: default,
        }
    }

    /// Set rate limit for a specific key
    pub async fn set_limit(&self, key: impl Into<String>, limit: RateLimit) {
        let key = key.into();
        let limiter = RateLimiter::new(limit.rate, limit.burst);

        let mut limiters = self.limiters.write().await;
        limiters.insert(key, limiter);
    }

    /// Check if request for key is allowed
    pub async fn check(&self, key: &str) -> bool {
        let limiters = self.limiters.read().await;

        if let Some(limiter) = limiters.get(key) {
            limiter.check().await
        } else if self.default_limit.rate == f64::MAX {
            true
        } else {
            drop(limiters);
            self.set_limit(key, self.default_limit.clone()).await;
            true
        }
    }

    /// Check and consume token for key
    pub async fn check_and_consume(&self, key: &str) -> bool {
        let should_create = {
            let limiters = self.limiters.read().await;
            !limiters.contains_key(key) && self.default_limit.rate != f64::MAX
        };

        if should_create {
            self.set_limit(key, self.default_limit.clone()).await;
        }

        let limiters = self.limiters.read().await;

        if let Some(limiter) = limiters.get(key) {
            limiter.check_and_consume().await
        } else {
            true
        }
    }

    /// Get available tokens for a key
    pub async fn available_tokens(&self, key: &str) -> Option<f64> {
        let limiters = self.limiters.read().await;
        if let Some(limiter) = limiters.get(key) {
            Some(limiter.available_tokens().await)
        } else {
            None
        }
    }

    /// Clear all limiters
    pub async fn clear(&self) {
        let mut limiters = self.limiters.write().await;
        limiters.clear();
    }
}

impl Default for MultiKeyRateLimiter {
    fn default() -> Self {
        Self::new()
    }
}

pub struct RateLimiterManager {
    node_limiters: Arc<MultiKeyRateLimiter>,
}

impl RateLimiterManager {
    /// Create a new rate limiter manager
    pub fn new() -> Self {
        Self {
            node_limiters: Arc::new(MultiKeyRateLimiter::new()),
        }
    }

    /// Check if node can publish
    pub async fn check_node(&self, node_id: &str) -> bool {
        self.node_limiters.check_and_consume(node_id).await
    }
}

impl Default for RateLimiterManager {
    fn default() -> Self {
        Self::new()
    }
}