rustytorch/crates/production/rtx-serving-api/src/rate_limiting.rs

//! Advanced rate limiting and token management
//!
//! Provides comprehensive rate limiting capabilities including:
//! - Per-user token quotas with sliding window tracking
//! - Organization-level rate limits with hierarchical enforcement
//! - Token cost calculation per model and operation type
//! - Dynamic rate limiting based on system load and user tier
//! - Burst allowances with exponential backoff
//! - Token usage analytics and reporting

use anyhow::Result;
use dashmap::DashMap;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::{sync::Arc, time::Duration};
use tokio::time::Instant;

/// User tier for rate limiting
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum UserTier {
    Free,
    Pro,
    Enterprise,
}

impl UserTier {
    /// Get base rate limit multiplier
    #[must_use]
    pub fn rate_multiplier(self) -> f64 {
        match self {
            Self::Free => 1.0,
            Self::Pro => 5.0,
            Self::Enterprise => 20.0,
        }
    }

    /// Get burst allowance multiplier
    #[must_use]
    pub fn burst_multiplier(self) -> f64 {
        match self {
            Self::Free => 2.0,
            Self::Pro => 3.0,
            Self::Enterprise => 5.0,
        }
    }
}

/// Token cost calculation based on model and operation
#[derive(Debug, Clone, PartialEq)]
pub struct TokenCost {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub model_cost_multiplier: f64,
    pub operation_cost_multiplier: f64,
}

impl TokenCost {
    /// Calculate total cost
    #[must_use]
    pub fn total_cost(&self) -> f64 {
        let base_cost = (self.input_tokens as f64 * 0.001) + (self.output_tokens as f64 * 0.002);
        base_cost * self.model_cost_multiplier * self.operation_cost_multiplier
    }

    /// Calculate token usage
    #[must_use]
    pub fn total_tokens(&self) -> u64 {
        self.input_tokens + self.output_tokens
    }
}

/// Sliding window for rate limiting
#[derive(Debug, Clone)]
pub struct SlidingWindow {
    window_size: Duration,
    buckets: Vec<u64>,
    bucket_duration: Duration,
    last_update: Instant,
}

impl SlidingWindow {
    /// Create new sliding window
    #[must_use]
    pub fn new(window_size: Duration, bucket_count: usize) -> Self {
        let bucket_duration = window_size / bucket_count as u32;
        Self {
            window_size,
            buckets: vec![0; bucket_count],
            bucket_duration,
            last_update: Instant::now(),
        }
    }

    /// Update window and add tokens
    pub fn add_tokens(&mut self, tokens: u64) {
        self.update_buckets();
        let current_bucket = self.buckets.len() - 1;
        self.buckets[current_bucket] += tokens;
    }

    /// Get current usage in window
    pub fn current_usage(&mut self) -> u64 {
        self.update_buckets();
        self.buckets.iter().sum()
    }

    /// Update buckets based on elapsed time
    fn update_buckets(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_update);
        let buckets_to_advance = (elapsed.as_millis() / self.bucket_duration.as_millis()) as usize;

        if buckets_to_advance > 0 {
            // Shift buckets
            let advance = buckets_to_advance.min(self.buckets.len());
            self.buckets.rotate_left(advance);

            // Clear advanced buckets
            let clear_start = self.buckets.len().saturating_sub(advance);
            for bucket in &mut self.buckets[clear_start..] {
                *bucket = 0;
            }

            self.last_update = now;
        }
    }
}

/// User rate limit state
#[derive(Debug)]
pub struct UserRateLimit {
    pub user_id: String,
    pub organization_id: Option<String>,
    pub tier: UserTier,
    pub token_window: Arc<RwLock<SlidingWindow>>,
    pub request_window: Arc<RwLock<SlidingWindow>>,
    pub last_request: Instant,
    pub consecutive_limits: u32,
    pub burst_tokens_used: u64,
    pub total_cost: f64,
}

impl UserRateLimit {
    /// Create new user rate limit
    #[must_use]
    pub fn new(user_id: String, organization_id: Option<String>, tier: UserTier) -> Self {
        let token_window = Arc::new(RwLock::new(SlidingWindow::new(Duration::from_secs(60), 60)));
        let request_window = Arc::new(RwLock::new(SlidingWindow::new(Duration::from_secs(60), 60)));

        Self {
            user_id,
            organization_id,
            tier,
            token_window,
            request_window,
            last_request: Instant::now(),
            consecutive_limits: 0,
            burst_tokens_used: 0,
            total_cost: 0.0,
        }
    }

    /// Check if request is allowed
    pub fn check_rate_limit(&mut self, cost: &TokenCost) -> Result<bool> {
        let now = Instant::now();

        // Calculate limits based on tier
        let base_token_limit = 1000_u64;
        let base_request_limit = 60_u64;

        let token_limit = (base_token_limit as f64 * self.tier.rate_multiplier()) as u64;
        let request_limit = (base_request_limit as f64 * self.tier.rate_multiplier()) as u64;

        // Check request limit
        {
            let mut request_window = self.request_window.write();
            let current_requests = request_window.current_usage();

            if current_requests >= request_limit {
                self.consecutive_limits += 1;
                return Ok(false);
            }
        }

        // Check token limit
        {
            let mut token_window = self.token_window.write();
            let current_tokens = token_window.current_usage();
            let required_tokens = cost.total_tokens();

            if current_tokens + required_tokens > token_limit {
                // Check burst allowance
                let burst_limit = (token_limit as f64 * self.tier.burst_multiplier()) as u64;
                if current_tokens + required_tokens > burst_limit {
                    self.consecutive_limits += 1;
                    return Ok(false);
                }

                self.burst_tokens_used += required_tokens;
            }
        }

        // Apply exponential backoff for consecutive limits
        if self.consecutive_limits > 0 {
            let backoff_duration = Duration::from_millis(100 * 2_u64.pow(self.consecutive_limits));
            let time_since_last = now.duration_since(self.last_request);

            if time_since_last < backoff_duration {
                return Ok(false);
            }
        }

        // Request allowed - update counters
        {
            let mut token_window = self.token_window.write();
            token_window.add_tokens(cost.total_tokens());
        }

        {
            let mut request_window = self.request_window.write();
            request_window.add_tokens(1);
        }

        self.last_request = now;
        self.consecutive_limits = 0;
        self.total_cost += cost.total_cost();

        Ok(true)
    }

    /// Get current usage statistics
    pub fn get_usage_stats(&mut self) -> UserUsageStats {
        let token_usage = {
            let mut token_window = self.token_window.write();
            token_window.current_usage()
        };

        let request_usage = {
            let mut request_window = self.request_window.write();
            request_window.current_usage()
        };

        UserUsageStats {
            user_id: self.user_id.clone(),
            tier: self.tier,
            tokens_used: token_usage,
            requests_made: request_usage,
            burst_tokens_used: self.burst_tokens_used,
            total_cost: self.total_cost,
            consecutive_limits: self.consecutive_limits,
        }
    }
}

/// Organization rate limit state
#[derive(Debug)]
pub struct OrganizationRateLimit {
    pub organization_id: String,
    pub token_window: Arc<RwLock<SlidingWindow>>,
    pub request_window: Arc<RwLock<SlidingWindow>>,
    pub total_cost: f64,
    pub user_limits: DashMap<String, Arc<RwLock<UserRateLimit>>>,
}

impl OrganizationRateLimit {
    /// Create new organization rate limit
    #[must_use]
    pub fn new(organization_id: String) -> Self {
        let token_window = Arc::new(RwLock::new(SlidingWindow::new(
            Duration::from_secs(3600),
            60,
        )));
        let request_window = Arc::new(RwLock::new(SlidingWindow::new(
            Duration::from_secs(3600),
            60,
        )));

        Self {
            organization_id,
            token_window,
            request_window,
            total_cost: 0.0,
            user_limits: DashMap::new(),
        }
    }

    /// Check organization-level limits
    pub fn check_organization_limit(&mut self, cost: &TokenCost) -> Result<bool> {
        // Organization limits (much higher than individual)
        let org_token_limit = 1_000_000_u64;
        let org_request_limit = 10_000_u64;

        // Check request limit
        {
            let mut request_window = self.request_window.write();
            let current_requests = request_window.current_usage();

            if current_requests >= org_request_limit {
                return Ok(false);
            }
        }

        // Check token limit
        {
            let mut token_window = self.token_window.write();
            let current_tokens = token_window.current_usage();

            if current_tokens + cost.total_tokens() > org_token_limit {
                return Ok(false);
            }

            // Update usage
            token_window.add_tokens(cost.total_tokens());
        }

        {
            let mut request_window = self.request_window.write();
            request_window.add_tokens(1);
        }

        self.total_cost += cost.total_cost();

        Ok(true)
    }
}

/// Rate limiting manager
pub struct RateLimitManager {
    user_limits: DashMap<String, Arc<RwLock<UserRateLimit>>>,
    organization_limits: DashMap<String, Arc<RwLock<OrganizationRateLimit>>>,
    system_load_factor: Arc<RwLock<f64>>,
}

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

impl RateLimitManager {
    /// Create new rate limit manager
    #[must_use]
    pub fn new() -> Self {
        Self {
            user_limits: DashMap::new(),
            organization_limits: DashMap::new(),
            system_load_factor: Arc::new(RwLock::new(1.0)),
        }
    }

    /// Check rate limits for a request
    pub async fn check_rate_limit(
        &self,
        user_id: &str,
        organization_id: Option<&str>,
        tier: UserTier,
        cost: &TokenCost,
    ) -> Result<bool> {
        // Apply system load factor
        let load_factor = *self.system_load_factor.read();
        let adjusted_cost = TokenCost {
            input_tokens: (cost.input_tokens as f64 * load_factor) as u64,
            output_tokens: (cost.output_tokens as f64 * load_factor) as u64,
            model_cost_multiplier: cost.model_cost_multiplier,
            operation_cost_multiplier: cost.operation_cost_multiplier,
        };

        // Check organization limits first if applicable
        if let Some(org_id) = organization_id {
            let org_limit = self
                .organization_limits
                .entry(org_id.to_string())
                .or_insert_with(|| {
                    Arc::new(RwLock::new(OrganizationRateLimit::new(org_id.to_string())))
                });

            let mut org_limit = org_limit.write();
            if !org_limit.check_organization_limit(&adjusted_cost)? {
                return Ok(false);
            }
        }

        // Check user limits
        let user_limit = self
            .user_limits
            .entry(user_id.to_string())
            .or_insert_with(|| {
                Arc::new(RwLock::new(UserRateLimit::new(
                    user_id.to_string(),
                    organization_id.map(String::from),
                    tier,
                )))
            });

        let mut user_limit = user_limit.write();
        user_limit.check_rate_limit(&adjusted_cost)
    }

    /// Update system load factor
    pub fn update_system_load(&self, load_factor: f64) {
        *self.system_load_factor.write() = load_factor.clamp(0.1, 10.0);
    }

    /// Get usage statistics for a user
    #[must_use]
    pub fn get_user_usage(&self, user_id: &str) -> Option<UserUsageStats> {
        self.user_limits.get(user_id).map(|limit| {
            let mut limit = limit.write();
            limit.get_usage_stats()
        })
    }

    /// Get organization usage statistics
    #[must_use]
    pub fn get_organization_usage(&self, organization_id: &str) -> Option<OrganizationUsageStats> {
        self.organization_limits.get(organization_id).map(|limit| {
            let limit = limit.read();

            let token_usage = {
                let mut token_window = limit.token_window.write();
                token_window.current_usage()
            };

            let request_usage = {
                let mut request_window = limit.request_window.write();
                request_window.current_usage()
            };

            OrganizationUsageStats {
                organization_id: limit.organization_id.clone(),
                tokens_used: token_usage,
                requests_made: request_usage,
                total_cost: limit.total_cost,
                user_count: limit.user_limits.len(),
            }
        })
    }

    /// Clear expired rate limit entries
    pub async fn cleanup_expired_entries(&self) {
        // Remove users that haven't made requests in the last hour
        let cutoff = Instant::now() - Duration::from_secs(3600);

        self.user_limits.retain(|_, limit| {
            let limit = limit.read();
            limit.last_request > cutoff
        });
    }
}

/// User usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserUsageStats {
    pub user_id: String,
    pub tier: UserTier,
    pub tokens_used: u64,
    pub requests_made: u64,
    pub burst_tokens_used: u64,
    pub total_cost: f64,
    pub consecutive_limits: u32,
}

/// Organization usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OrganizationUsageStats {
    pub organization_id: String,
    pub tokens_used: u64,
    pub requests_made: u64,
    pub total_cost: f64,
    pub user_count: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;

    #[test]
    fn test_sliding_window_basic() {
        let mut window = SlidingWindow::new(Duration::from_millis(100), 10);

        window.add_tokens(10);
        assert_eq!(window.current_usage(), 10);

        window.add_tokens(5);
        assert_eq!(window.current_usage(), 15);
    }

    #[test]
    fn test_sliding_window_expiry() {
        let mut window = SlidingWindow::new(Duration::from_millis(50), 5);

        window.add_tokens(10);
        assert_eq!(window.current_usage(), 10);

        // Wait for window to expire
        thread::sleep(Duration::from_millis(60));
        assert_eq!(window.current_usage(), 0);
    }

    #[test]
    fn test_token_cost_calculation() {
        let cost = TokenCost {
            input_tokens: 100,
            output_tokens: 50,
            model_cost_multiplier: 2.0,
            operation_cost_multiplier: 1.5,
        };

        assert_eq!(cost.total_tokens(), 150);
        // (100*0.001 + 50*0.002) * 2.0 * 1.5 = 0.6
        assert!((cost.total_cost() - 0.6).abs() < 1e-9);
    }

    #[test]
    fn test_user_tier_multipliers() {
        assert_eq!(UserTier::Free.rate_multiplier(), 1.0);
        assert_eq!(UserTier::Pro.rate_multiplier(), 5.0);
        assert_eq!(UserTier::Enterprise.rate_multiplier(), 20.0);
    }

    #[tokio::test]
    async fn test_rate_limit_manager() {
        let manager = RateLimitManager::new();
        let cost = TokenCost {
            input_tokens: 10,
            output_tokens: 5,
            model_cost_multiplier: 1.0,
            operation_cost_multiplier: 1.0,
        };

        // Should allow first request
        assert!(
            manager
                .check_rate_limit("user1", None, UserTier::Free, &cost)
                .await
                .unwrap()
        );

        // Should still allow reasonable usage
        for _ in 0..50 {
            assert!(
                manager
                    .check_rate_limit("user1", None, UserTier::Free, &cost)
                    .await
                    .unwrap()
            );
        }
    }

    #[tokio::test]
    async fn test_organization_limits() {
        let manager = RateLimitManager::new();
        let cost = TokenCost {
            input_tokens: 10,
            output_tokens: 5,
            model_cost_multiplier: 1.0,
            operation_cost_multiplier: 1.0,
        };

        // Should allow requests under organization limit
        assert!(
            manager
                .check_rate_limit("user1", Some("org1"), UserTier::Free, &cost)
                .await
                .unwrap()
        );
    }

    #[test]
    fn test_user_rate_limit_creation() {
        let user_limit =
            UserRateLimit::new("user1".to_string(), Some("org1".to_string()), UserTier::Pro);

        assert_eq!(user_limit.user_id, "user1");
        assert_eq!(user_limit.organization_id, Some("org1".to_string()));
        assert_eq!(user_limit.tier, UserTier::Pro);
    }

    #[test]
    fn test_organization_rate_limit_creation() {
        let org_limit = OrganizationRateLimit::new("org1".to_string());
        assert_eq!(org_limit.organization_id, "org1");
        assert_eq!(org_limit.total_cost, 0.0);
    }
}