//! Rate Limiter - Token bucket algorithm for federation API throttling
//!
//! Supports per-key (IP / API key / global) rate limiting using the token bucket
//! algorithm, which allows short bursts while enforcing a long-term average rate.

use std::collections::HashMap;
use std::sync::{Mutex, RwLock};
use std::time::Instant;

use super::error::{FederationError, Result};

/// Rate limiter configuration
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Whether rate limiting is enabled
    pub enabled: bool,
    /// Maximum tokens in the bucket (burst capacity)
    pub max_tokens: u64,
    /// Tokens added per second (sustained rate)
    pub refill_rate_per_sec: f64,
    /// Default per-key configuration (for IP / API key limiting)
    pub per_key_max_tokens: u64,
    /// Default per-key refill rate
    pub per_key_refill_rate_per_sec: f64,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            max_tokens: 100,
            refill_rate_per_sec: 20.0,
            per_key_max_tokens: 30,
            per_key_refill_rate_per_sec: 10.0,
        }
    }
}

/// A single token bucket
struct TokenBucket {
    tokens: f64,
    max_tokens: u64,
    refill_rate: f64, // tokens per second
    last_refill: Instant,
}

impl TokenBucket {
    fn new(max_tokens: u64, refill_rate: f64) -> Self {
        Self {
            tokens: max_tokens as f64,
            max_tokens,
            refill_rate,
            last_refill: Instant::now(),
        }
    }

    /// Try to consume one token. Returns true if allowed.
    fn try_acquire(&mut self) -> bool {
        self.refill();
        if self.tokens >= 1.0 {
            self.tokens -= 1.0;
            true
        } else {
            false
        }
    }

    /// Refill tokens based on elapsed time
    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
        self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.max_tokens as f64);
        self.last_refill = now;
    }

    /// How many milliseconds until the next token is available
    fn retry_after_ms(&self) -> u64 {
        if self.tokens >= 1.0 {
            return 0;
        }
        let deficit = 1.0 - self.tokens;
        let wait_secs = deficit / self.refill_rate;
        (wait_secs * 1000.0).ceil() as u64
    }
}

/// Rate limiter with global and per-key buckets
pub struct RateLimiter {
    config: RateLimitConfig,
    /// Global bucket
    global: Mutex<TokenBucket>,
    /// Per-key buckets (IP, API key, etc.)
    per_key: RwLock<HashMap<String, Mutex<TokenBucket>>>,
}

impl RateLimiter {
    pub fn new(config: RateLimitConfig) -> Self {
        let global = Mutex::new(TokenBucket::new(
            config.max_tokens,
            config.refill_rate_per_sec,
        ));
        Self {
            config,
            global,
            per_key: RwLock::new(HashMap::new()),
        }
    }

    pub fn with_defaults() -> Self {
        Self::new(RateLimitConfig::default())
    }

    /// Check global rate limit. Returns Ok(()) if the request is allowed.
    pub fn check_global(&self) -> Result<()> {
        if !self.config.enabled {
            return Ok(());
        }

        let mut bucket = self.global.lock()
            .map_err(|_| FederationError::Internal("rate limiter lock poisoned".to_string()))?;

        if bucket.try_acquire() {
            Ok(())
        } else {
            Err(FederationError::RateLimitExceeded)
        }
    }

    /// Check per-key rate limit (e.g. per IP or per API key).
    pub fn check_key(&self, key: &str) -> Result<()> {
        if !self.config.enabled {
            return Ok(());
        }

        // Fast path: bucket exists
        {
            let buckets = self.per_key.read()
                .map_err(|_| FederationError::Internal("rate limiter lock poisoned".to_string()))?;
            if let Some(bucket_mutex) = buckets.get(key) {
                let mut bucket = bucket_mutex.lock()
                    .map_err(|_| FederationError::Internal("rate limiter lock poisoned".to_string()))?;
                return if bucket.try_acquire() {
                    Ok(())
                } else {
                    Err(FederationError::RateLimitExceeded)
                };
            }
        }

        // Slow path: create bucket
        let mut buckets = self.per_key.write()
            .map_err(|_| FederationError::Internal("rate limiter lock poisoned".to_string()))?;

        let bucket_mutex = buckets.entry(key.to_string()).or_insert_with(|| {
            Mutex::new(TokenBucket::new(
                self.config.per_key_max_tokens,
                self.config.per_key_refill_rate_per_sec,
            ))
        });

        let mut bucket = bucket_mutex.lock()
            .map_err(|_| FederationError::Internal("rate limiter lock poisoned".to_string()))?;

        if bucket.try_acquire() {
            Ok(())
        } else {
            Err(FederationError::RateLimitExceeded)
        }
    }

    /// Check both global and per-key limits
    pub fn check(&self, key: &str) -> Result<()> {
        self.check_global()?;
        self.check_key(key)
    }

    /// Get the number of tracked keys
    pub fn tracked_key_count(&self) -> usize {
        self.per_key.read().map(|m| m.len()).unwrap_or(0)
    }

    /// Remove stale per-key buckets that are fully refilled (housekeeping)
    pub fn cleanup_stale_keys(&self) {
        if let Ok(mut buckets) = self.per_key.write() {
            buckets.retain(|_, bucket_mutex| {
                if let Ok(bucket) = bucket_mutex.lock() {
                    // Keep if bucket has been used recently (not full)
                    bucket.tokens < bucket.max_tokens as f64
                } else {
                    true // keep on error
                }
            });
        }
    }
}

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

    #[test]
    fn test_global_rate_limit() {
        let limiter = RateLimiter::new(RateLimitConfig {
            enabled: true,
            max_tokens: 3,
            refill_rate_per_sec: 0.0, // no refill — exhausts quickly
            per_key_max_tokens: 10,
            per_key_refill_rate_per_sec: 0.0,
        });

        assert!(limiter.check_global().is_ok());
        assert!(limiter.check_global().is_ok());
        assert!(limiter.check_global().is_ok());
        assert!(limiter.check_global().is_err()); // exhausted
    }

    #[test]
    fn test_per_key_rate_limit() {
        let limiter = RateLimiter::new(RateLimitConfig {
            enabled: true,
            max_tokens: 1000,
            refill_rate_per_sec: 1000.0,
            per_key_max_tokens: 2,
            per_key_refill_rate_per_sec: 0.0,
        });

        assert!(limiter.check_key("user1").is_ok());
        assert!(limiter.check_key("user1").is_ok());
        assert!(limiter.check_key("user1").is_err()); // user1 exhausted

        // user2 still has tokens
        assert!(limiter.check_key("user2").is_ok());
    }

    #[test]
    fn test_disabled_limiter() {
        let limiter = RateLimiter::new(RateLimitConfig {
            enabled: false,
            ..Default::default()
        });

        for _ in 0..1000 {
            assert!(limiter.check("any_key").is_ok());
        }
    }

    #[test]
    fn test_tracked_key_count() {
        let limiter = RateLimiter::with_defaults();
        limiter.check_key("a").ok();
        limiter.check_key("b").ok();
        limiter.check_key("c").ok();
        assert_eq!(limiter.tracked_key_count(), 3);
    }
}
