//! Rate limiting middleware to prevent API abuse
//!
//! Implements sliding window rate limiting per IP address:
//! - Configurable request limits and time windows
//! - Per-IP tracking with automatic cleanup
//! - Graceful handling with 429 Too Many Requests
//! - Optional exemptions for trusted IPs

use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};

use axum::{
    extract::Request,
    http::StatusCode,
    middleware::Next,
    response::{IntoResponse, Response},
};
use tokio::sync::RwLock;
use tracing::{debug, warn};

/// Rate limiter configuration
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Maximum requests allowed per window
    pub max_requests: u32,

    /// Time window duration
    pub window_duration: Duration,

    /// IPs exempt from rate limiting
    pub exempt_ips: Vec<IpAddr>,

    /// Enable rate limiting
    pub enabled: bool,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            max_requests: 100,
            window_duration: Duration::from_secs(60), // 100 requests per minute
            exempt_ips: vec![],
            enabled: true,
        }
    }
}

impl RateLimitConfig {
    /// Create new rate limit config with custom limits
    pub fn new(max_requests: u32, window_secs: u64) -> Self {
        Self {
            max_requests,
            window_duration: Duration::from_secs(window_secs),
            exempt_ips: vec![],
            enabled: true,
        }
    }

    /// Add exempt IP address
    pub fn with_exempt_ip(mut self, ip: IpAddr) -> Self {
        self.exempt_ips.push(ip);
        self
    }

    /// Add multiple exempt IPs
    pub fn with_exempt_ips(mut self, ips: Vec<IpAddr>) -> Self {
        self.exempt_ips.extend(ips);
        self
    }

    /// Disable rate limiting
    pub fn disabled() -> Self {
        Self {
            enabled: false,
            ..Default::default()
        }
    }
}

/// Request tracking for rate limiting
#[derive(Debug, Clone)]
struct RequestWindow {
    /// Request count in current window
    count: u32,

    /// Window start time
    window_start: Instant,
}

impl RequestWindow {
    fn new() -> Self {
        Self {
            count: 0,
            window_start: Instant::now(),
        }
    }

    fn increment(&mut self) {
        self.count += 1;
    }

    fn reset(&mut self) {
        self.count = 0;
        self.window_start = Instant::now();
    }

    fn is_expired(&self, window_duration: Duration) -> bool {
        self.window_start.elapsed() > window_duration
    }
}

/// Rate limiter state
pub struct RateLimiter {
    /// Configuration
    config: RateLimitConfig,

    /// Request tracking per IP
    requests: Arc<RwLock<HashMap<IpAddr, RequestWindow>>>,
}

impl RateLimiter {
    /// Create new rate limiter with default configuration
    pub fn new() -> Self {
        Self::with_config(RateLimitConfig::default())
    }

    /// Create new rate limiter with custom configuration
    pub fn with_config(config: RateLimitConfig) -> Self {
        let limiter = Self {
            config,
            requests: Arc::new(RwLock::new(HashMap::new())),
        };

        // Start background cleanup task
        limiter.start_cleanup_task();

        limiter
    }

    /// Check if request is allowed
    pub async fn check(&self, ip: IpAddr) -> Result<(), RateLimitError> {
        // Skip if disabled
        if !self.config.enabled {
            return Ok(());
        }

        // Check if IP is exempt
        if self.config.exempt_ips.contains(&ip) {
            debug!(ip = %ip, "IP exempt from rate limiting");
            return Ok(());
        }

        let mut requests = self.requests.write().await;
        let entry = requests.entry(ip).or_insert_with(RequestWindow::new);

        // Reset window if expired
        if entry.is_expired(self.config.window_duration) {
            debug!(ip = %ip, "Rate limit window expired, resetting");
            entry.reset();
        }

        // Check limit
        if entry.count >= self.config.max_requests {
            warn!(
                ip = %ip,
                count = entry.count,
                limit = self.config.max_requests,
                "Rate limit exceeded"
            );
            return Err(RateLimitError::TooManyRequests {
                limit: self.config.max_requests,
                window_secs: self.config.window_duration.as_secs(),
            });
        }

        // Increment counter
        entry.increment();

        debug!(
            ip = %ip,
            count = entry.count,
            limit = self.config.max_requests,
            "Request allowed"
        );

        Ok(())
    }

    /// Get current request count for IP
    pub async fn get_count(&self, ip: IpAddr) -> u32 {
        let requests = self.requests.read().await;
        requests.get(&ip).map(|w| w.count).unwrap_or(0)
    }

    /// Get statistics
    pub async fn get_stats(&self) -> RateLimitStats {
        let requests = self.requests.read().await;

        let total_ips = requests.len();
        let total_requests: u32 = requests.values().map(|w| w.count).sum();
        let limited_ips = requests
            .values()
            .filter(|w| w.count >= self.config.max_requests)
            .count();

        RateLimitStats {
            total_ips,
            total_requests,
            limited_ips,
            config: self.config.clone(),
        }
    }

    /// Start background cleanup task to remove old entries
    fn start_cleanup_task(&self) {
        let requests = self.requests.clone();
        let cleanup_interval = self.config.window_duration * 2;

        tokio::spawn(async move {
            loop {
                tokio::time::sleep(cleanup_interval).await;

                let mut requests = requests.write().await;
                let before_count = requests.len();

                // Remove entries older than 2x window duration
                requests.retain(|_, window| window.window_start.elapsed() < cleanup_interval);

                let after_count = requests.len();
                let removed = before_count - after_count;

                if removed > 0 {
                    debug!(
                        removed = removed,
                        remaining = after_count,
                        "Cleaned up expired rate limit entries"
                    );
                }
            }
        });
    }
}

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

/// Rate limiter statistics
#[derive(Debug, Clone)]
pub struct RateLimitStats {
    pub total_ips: usize,
    pub total_requests: u32,
    pub limited_ips: usize,
    pub config: RateLimitConfig,
}

/// Rate limiting middleware
pub async fn rate_limit_middleware(
    limiter: Arc<RateLimiter>,
    request: Request,
    next: Next,
) -> Result<Response, RateLimitError> {
    // Extract IP address from request
    let ip = extract_ip_address(&request);

    // Check rate limit
    limiter.check(ip).await?;

    // Continue to next middleware/handler
    Ok(next.run(request).await)
}

/// Extract IP address from request
fn extract_ip_address(req: &Request) -> IpAddr {
    use std::str::FromStr;

    // Try X-Forwarded-For first (proxy/load balancer)
    if let Some(forwarded) = req
        .headers()
        .get("X-Forwarded-For")
        .and_then(|h| h.to_str().ok())
    {
        if let Some(first_ip) = forwarded.split(',').next() {
            if let Ok(ip) = IpAddr::from_str(first_ip.trim()) {
                return ip;
            }
        }
    }

    // Try X-Real-IP next
    if let Some(real_ip) = req.headers().get("X-Real-IP").and_then(|h| h.to_str().ok()) {
        if let Ok(ip) = IpAddr::from_str(real_ip.trim()) {
            return ip;
        }
    }

    // Default to localhost
    IpAddr::from_str("127.0.0.1").unwrap()
}

/// Rate limit error
#[derive(Debug)]
pub enum RateLimitError {
    TooManyRequests { limit: u32, window_secs: u64 },
}

impl IntoResponse for RateLimitError {
    fn into_response(self) -> Response {
        match self {
            RateLimitError::TooManyRequests { limit, window_secs } => {
                let message = format!(
                    "Rate limit exceeded. Maximum {} requests per {} seconds.",
                    limit, window_secs
                );
                (StatusCode::TOO_MANY_REQUESTS, message).into_response()
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use super::*;

    #[tokio::test]
    async fn test_rate_limiter_basic() {
        let config = RateLimitConfig::new(5, 60);
        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::from_str("192.168.1.1").unwrap();

        // First 5 requests should succeed
        for _ in 0..5 {
            assert!(limiter.check(ip).await.is_ok());
        }

        // 6th request should fail
        assert!(limiter.check(ip).await.is_err());
    }

    #[tokio::test]
    async fn test_rate_limiter_window_reset() {
        let config = RateLimitConfig::new(2, 1); // 2 requests per second
        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::from_str("192.168.1.2").unwrap();

        // Use up the limit
        assert!(limiter.check(ip).await.is_ok());
        assert!(limiter.check(ip).await.is_ok());
        assert!(limiter.check(ip).await.is_err());

        // Wait for window to reset
        tokio::time::sleep(Duration::from_secs(2)).await;

        // Should work again
        assert!(limiter.check(ip).await.is_ok());
    }

    #[tokio::test]
    async fn test_rate_limiter_exempt_ip() {
        let exempt_ip = IpAddr::from_str("10.0.0.1").unwrap();
        let config = RateLimitConfig::new(2, 60).with_exempt_ip(exempt_ip);
        let limiter = RateLimiter::with_config(config);

        // Exempt IP can make unlimited requests
        for _ in 0..10 {
            assert!(limiter.check(exempt_ip).await.is_ok());
        }
    }

    #[tokio::test]
    async fn test_rate_limiter_disabled() {
        let config = RateLimitConfig::disabled();
        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::from_str("192.168.1.3").unwrap();

        // Should allow unlimited requests when disabled
        for _ in 0..100 {
            assert!(limiter.check(ip).await.is_ok());
        }
    }

    #[tokio::test]
    async fn test_get_count() {
        let config = RateLimitConfig::new(10, 60);
        let limiter = RateLimiter::with_config(config);
        let ip = IpAddr::from_str("192.168.1.4").unwrap();

        assert_eq!(limiter.get_count(ip).await, 0);

        limiter.check(ip).await.unwrap();
        assert_eq!(limiter.get_count(ip).await, 1);

        limiter.check(ip).await.unwrap();
        limiter.check(ip).await.unwrap();
        assert_eq!(limiter.get_count(ip).await, 3);
    }

    #[tokio::test]
    async fn test_get_stats() {
        let config = RateLimitConfig::new(5, 60);
        let limiter = RateLimiter::with_config(config);

        let ip1 = IpAddr::from_str("192.168.1.5").unwrap();
        let ip2 = IpAddr::from_str("192.168.1.6").unwrap();

        for _ in 0..3 {
            limiter.check(ip1).await.unwrap();
        }

        for _ in 0..6 {
            let _ = limiter.check(ip2).await; // Will fail after 5
        }

        let stats = limiter.get_stats().await;
        assert_eq!(stats.total_ips, 2);
        assert_eq!(stats.total_requests, 8); // 3 + 5
        assert_eq!(stats.limited_ips, 1); // ip2 hit limit
    }
}
