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

use axum::{extract::ConnectInfo, http::Request, response::Response};
use tokio::time;
use tower::Layer;
use tower_service::Service;
use tracing::{debug, warn};

/// Rate limiting configuration
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    /// Maximum requests per window
    pub max_requests: u32,
    /// Time window duration
    pub window: Duration,
    /// Whether to apply rate limiting per IP address
    pub per_ip: bool,
    /// Whether to apply global rate limiting
    pub global: bool,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            max_requests: 100,
            window: Duration::from_secs(60), // 1 minute
            per_ip: true,
            global: false,
        }
    }
}

impl RateLimitConfig {
    /// Create a permissive rate limit config (for development)
    pub fn permissive() -> Self {
        Self {
            max_requests: 1000,
            window: Duration::from_secs(60),
            per_ip: true,
            global: false,
        }
    }

    /// Create a strict rate limit config (for production)
    pub fn strict() -> Self {
        Self {
            max_requests: 50,
            window: Duration::from_secs(60),
            per_ip: true,
            global: true,
        }
    }

    /// Create an authentication-specific rate limit config
    pub fn auth() -> Self {
        Self {
            max_requests: 10,
            window: Duration::from_secs(300), // 5 minutes
            per_ip: true,
            global: false,
        }
    }
}

/// Rate limiter implementation
#[derive(Debug)]
pub struct RateLimiter {
    config: RateLimitConfig,
    per_ip_buckets: Arc<Mutex<HashMap<String, RequestBucket>>>,
    global_bucket: Arc<Mutex<RequestBucket>>,
}

/// Request bucket for tracking requests within a time window
#[derive(Debug)]
struct RequestBucket {
    requests: Vec<Instant>,
    last_cleanup: Instant,
}

impl RequestBucket {
    fn new() -> Self {
        Self {
            requests: Vec::new(),
            last_cleanup: Instant::now(),
        }
    }

    /// Check if request is allowed and add it to the bucket
    fn try_request(&mut self, max_requests: u32, window: Duration) -> bool {
        let now = Instant::now();

        // Clean up old requests periodically
        if now.duration_since(self.last_cleanup) > window / 2 {
            self.cleanup(window);
            self.last_cleanup = now;
        }

        // Count recent requests
        let recent_requests = self
            .requests
            .iter()
            .filter(|&&req_time| now.duration_since(req_time) <= window)
            .count();

        if recent_requests >= max_requests as usize {
            debug!(
                "Rate limit exceeded: {} requests in window",
                recent_requests
            );
            return false;
        }

        // Add this request to the bucket
        self.requests.push(now);
        true
    }

    /// Remove requests older than the window
    fn cleanup(&mut self, window: Duration) {
        let now = Instant::now();
        self.requests
            .retain(|&req_time| now.duration_since(req_time) <= window);
    }
}

impl RateLimiter {
    /// Create a new rate limiter
    pub fn new(config: RateLimitConfig) -> Self {
        Self {
            config,
            per_ip_buckets: Arc::new(Mutex::new(HashMap::new())),
            global_bucket: Arc::new(Mutex::new(RequestBucket::new())),
        }
    }

    /// Check if request is allowed
    pub fn check_request(&self, client_ip: Option<&str>) -> bool {
        // Check global rate limit
        if self.config.global {
            let mut global_bucket = self.global_bucket.lock().unwrap();
            if !global_bucket.try_request(self.config.max_requests, self.config.window) {
                warn!("Global rate limit exceeded");
                return false;
            }
        }

        // Check per-IP rate limit
        if self.config.per_ip {
            if let Some(ip) = client_ip {
                let mut buckets = self.per_ip_buckets.lock().unwrap();
                let bucket = buckets
                    .entry(ip.to_string())
                    .or_insert_with(RequestBucket::new);

                if !bucket.try_request(self.config.max_requests, self.config.window) {
                    warn!("Rate limit exceeded for IP: {}", ip);
                    return false;
                }
            }
        }

        true
    }

    /// Start a cleanup task to remove old buckets
    pub fn start_cleanup_task(self: Arc<Self>) {
        let cleanup_interval = self.config.window * 2;

        tokio::spawn(async move {
            let mut interval = time::interval(cleanup_interval);

            loop {
                interval.tick().await;
                self.cleanup_buckets().await;
                debug!("Rate limiter cleanup completed");
            }
        });
    }

    async fn cleanup_buckets(&self) {
        // Clean up per-IP buckets
        let mut removed_ips = Vec::new();
        {
            let mut buckets = self.per_ip_buckets.lock().unwrap();
            buckets.retain(|ip, bucket| {
                bucket.cleanup(self.config.window);
                let keep = !bucket.requests.is_empty();
                if !keep {
                    removed_ips.push(ip.clone());
                }
                keep
            });
        }
        for ip in removed_ips {
            debug!("Removing empty rate limit bucket for IP: {}", ip);
        }

        // Clean up global bucket
        {
            let mut global_bucket = self.global_bucket.lock().unwrap();
            global_bucket.cleanup(self.config.window);
        }
    }

    /// Get current statistics
    pub fn get_stats(&self) -> RateLimitStats {
        let buckets = self.per_ip_buckets.lock().unwrap();
        let global_bucket = self.global_bucket.lock().unwrap();

        let active_ips = buckets.len();
        let total_recent_requests: usize = buckets
            .values()
            .map(|bucket| {
                let now = Instant::now();
                bucket
                    .requests
                    .iter()
                    .filter(|&&req_time| now.duration_since(req_time) <= self.config.window)
                    .count()
            })
            .sum();

        let global_recent_requests = {
            let now = Instant::now();
            global_bucket
                .requests
                .iter()
                .filter(|&&req_time| now.duration_since(req_time) <= self.config.window)
                .count()
        };

        RateLimitStats {
            active_ips,
            total_recent_requests,
            global_recent_requests,
            max_requests_per_window: self.config.max_requests,
            window_seconds: self.config.window.as_secs(),
        }
    }
}

/// Rate limit statistics
#[derive(Debug, serde::Serialize)]
pub struct RateLimitStats {
    pub active_ips: usize,
    pub total_recent_requests: usize,
    pub global_recent_requests: usize,
    pub max_requests_per_window: u32,
    pub window_seconds: u64,
}

/// Rate limiting middleware layer
#[derive(Clone)]
pub struct RateLimitLayer {
    rate_limiter: Arc<RateLimiter>,
}

impl RateLimitLayer {
    /// Create a new rate limit layer
    pub fn new(config: RateLimitConfig) -> Self {
        let rate_limiter = Arc::new(RateLimiter::new(config));

        // Start cleanup task
        rate_limiter.clone().start_cleanup_task();

        Self { rate_limiter }
    }

    /// Get the rate limiter instance
    pub fn rate_limiter(&self) -> Arc<RateLimiter> {
        self.rate_limiter.clone()
    }
}

impl<S> Layer<S> for RateLimitLayer {
    type Service = RateLimitService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        RateLimitService {
            inner,
            rate_limiter: self.rate_limiter.clone(),
        }
    }
}

/// Rate limiting service
#[derive(Clone)]
pub struct RateLimitService<S> {
    inner: S,
    rate_limiter: Arc<RateLimiter>,
}

impl<S> Service<Request<axum::body::Body>> for RateLimitService<S>
where
    S: Service<Request<axum::body::Body>, Response = Response> + Clone + Send + 'static,
    S::Future: Send + 'static,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>> + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
    >;

    fn poll_ready(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Result<(), <Self as Service<Request<axum::body::Body>>>::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(
        &mut self,
        request: Request<axum::body::Body>,
    ) -> <Self as Service<Request<axum::body::Body>>>::Future {
        // Extract client IP
        let client_ip = request
            .extensions()
            .get::<ConnectInfo<SocketAddr>>()
            .map(|connect_info| connect_info.0.ip().to_string());

        // Check rate limit
        let allowed = self.rate_limiter.check_request(client_ip.as_deref());

        let mut inner = self.inner.clone();
        let rate_limiter = self.rate_limiter.clone();

        Box::pin(async move {
            if !allowed {
                // Create rate limited response
                let response = Response::builder()
                    .status(429) // Too Many Requests
                    .header(
                        "X-RateLimit-Limit",
                        rate_limiter.config.max_requests.to_string(),
                    )
                    .header(
                        "X-RateLimit-Window",
                        rate_limiter.config.window.as_secs().to_string(),
                    )
                    .header(
                        "Retry-After",
                        rate_limiter.config.window.as_secs().to_string(),
                    )
                    .body(axum::body::Body::from("Rate limit exceeded"))
                    .unwrap();

                return Ok(response);
            }

            // Process the request normally
            let response = inner.call(request).await?;

            // Add rate limit headers to response
            let response = response.map(|body| {
                // Note: In a real implementation, you might want to add headers here
                // This is a simplified version
                body
            });

            Ok(response)
        })
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;

    #[tokio::test]
    async fn test_rate_limiter_basic() {
        let config = RateLimitConfig {
            max_requests: 2,
            window: Duration::from_secs(1),
            per_ip: true,
            global: false,
        };

        let limiter = RateLimiter::new(config);

        // First two requests should be allowed
        assert!(limiter.check_request(Some("127.0.0.1")));
        assert!(limiter.check_request(Some("127.0.0.1")));

        // Third request should be denied
        assert!(!limiter.check_request(Some("127.0.0.1")));

        // Different IP should be allowed
        assert!(limiter.check_request(Some("192.168.1.1")));
    }

    #[tokio::test]
    async fn test_rate_limiter_window() {
        let config = RateLimitConfig {
            max_requests: 1,
            window: Duration::from_millis(100),
            per_ip: true,
            global: false,
        };

        let limiter = RateLimiter::new(config);

        // First request should be allowed
        assert!(limiter.check_request(Some("127.0.0.1")));

        // Second request should be denied
        assert!(!limiter.check_request(Some("127.0.0.1")));

        // Wait for window to pass
        tokio::time::sleep(Duration::from_millis(150)).await;

        // Should be allowed again
        assert!(limiter.check_request(Some("127.0.0.1")));
    }

    #[tokio::test]
    async fn test_rate_limiter_global() {
        let config = RateLimitConfig {
            max_requests: 2,
            window: Duration::from_secs(1),
            per_ip: false,
            global: true,
        };

        let limiter = RateLimiter::new(config);

        // First two requests from different IPs should be allowed
        assert!(limiter.check_request(Some("127.0.0.1")));
        assert!(limiter.check_request(Some("192.168.1.1")));

        // Third request from yet another IP should be denied (global limit)
        assert!(!limiter.check_request(Some("10.0.0.1")));
    }

    #[test]
    fn test_rate_limit_configs() {
        let permissive = RateLimitConfig::permissive();
        assert_eq!(permissive.max_requests, 1000);

        let strict = RateLimitConfig::strict();
        assert_eq!(strict.max_requests, 50);
        assert!(strict.global);

        let auth = RateLimitConfig::auth();
        assert_eq!(auth.max_requests, 10);
        assert_eq!(auth.window, Duration::from_secs(300));
    }
}
