// rate_limiter.rs
// Advanced rate limiting with authentication integration
// Critical for production security and DDoS protection

use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{error, info, warn};
use uuid::Uuid;

/// Advanced rate limiter with per-user, per-IP, and global limits
#[derive(Clone)]
pub struct RateLimiter {
    config: RateLimitConfig,
    user_buckets: Arc<RwLock<HashMap<Uuid, TokenBucket>>>,
    ip_buckets: Arc<RwLock<HashMap<IpAddr, TokenBucket>>>,
    global_bucket: Arc<RwLock<TokenBucket>>,
}

/// Rate limiting configuration
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
    // Per-user limits
    pub user_requests_per_minute: u32,
    pub user_burst_capacity: u32,

    // Per-IP limits
    pub ip_requests_per_minute: u32,
    pub ip_burst_capacity: u32,

    // Global limits
    pub global_requests_per_minute: u32,
    pub global_burst_capacity: u32,

    // Authentication failures
    pub auth_failures_per_minute: u32,
    pub auth_failure_penalty_minutes: u32,

    // Trading specific limits
    pub orders_per_minute: u32,
    pub order_burst_capacity: u32,

    // Cleanup intervals
    pub cleanup_interval_minutes: u32,
}

impl Default for RateLimitConfig {
    fn default() -> Self {
        Self {
            // Conservative production defaults
            user_requests_per_minute: 1000, // 1000 requests per minute per user
            user_burst_capacity: 100,       // Allow bursts up to 100 requests

            ip_requests_per_minute: 2000, // 2000 requests per minute per IP
            ip_burst_capacity: 200,       // Allow bursts up to 200 requests

            global_requests_per_minute: 50000, // 50k requests per minute globally
            global_burst_capacity: 5000,       // Allow bursts up to 5k requests

            auth_failures_per_minute: 5, // Max 5 auth failures per minute
            auth_failure_penalty_minutes: 15, // 15 minute penalty for repeated failures

            orders_per_minute: 600,   // 600 orders per minute (10/second)
            order_burst_capacity: 60, // Allow bursts up to 60 orders

            cleanup_interval_minutes: 60, // Cleanup every hour
        }
    }
}

/// Token bucket for rate limiting
#[derive(Debug, Clone)]
struct TokenBucket {
    tokens: f64,
    capacity: f64,
    refill_rate: f64, // tokens per second
    last_refill: Instant,
    penalty_until: Option<Instant>,
}

impl TokenBucket {
    fn new(capacity: u32, refill_per_minute: u32) -> Self {
        let capacity = capacity as f64;
        Self {
            tokens: capacity,
            capacity,
            refill_rate: refill_per_minute as f64 / 60.0, // Convert to per second
            last_refill: Instant::now(),
            penalty_until: None,
        }
    }

    fn try_consume(&mut self, tokens: f64) -> bool {
        // Check if under penalty
        if let Some(penalty_until) = self.penalty_until {
            if Instant::now() < penalty_until {
                return false; // Still under penalty
            } else {
                self.penalty_until = None; // Penalty expired
            }
        }

        // Refill tokens based on time passed
        let now = Instant::now();
        let time_passed = now.duration_since(self.last_refill).as_secs_f64();
        let new_tokens = time_passed * self.refill_rate;
        self.tokens = (self.tokens + new_tokens).min(self.capacity);
        self.last_refill = now;

        // Try to consume tokens
        if self.tokens >= tokens {
            self.tokens -= tokens;
            true
        } else {
            false
        }
    }

    fn apply_penalty(&mut self, duration: Duration) {
        self.penalty_until = Some(Instant::now() + duration);
        self.tokens = 0.0; // Drain all tokens
    }

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

/// Rate limit check result
#[derive(Debug, Clone)]
pub enum RateLimitResult {
    Allowed,
    UserLimitExceeded,
    IpLimitExceeded,
    GlobalLimitExceeded,
    AuthFailurePenalty,
    OrderLimitExceeded,
}

/// Rate limit check context
#[derive(Debug, Clone)]
pub struct RateLimitContext {
    pub user_id: Option<Uuid>,
    pub ip_addr: IpAddr,
    pub request_type: RequestType,
    pub tokens_requested: f64,
}

#[derive(Debug, Clone)]
pub enum RequestType {
    General,
    Authentication,
    Trading,
    OrderPlacement,
    MarketData,
    Risk,
    ML,
    Config,
}

impl RateLimiter {
    /// Create new rate limiter with configuration
    pub fn new(config: RateLimitConfig) -> Self {
        let global_bucket = TokenBucket::new(
            config.global_burst_capacity,
            config.global_requests_per_minute,
        );

        Self {
            config,
            user_buckets: Arc::new(RwLock::new(HashMap::new())),
            ip_buckets: Arc::new(RwLock::new(HashMap::new())),
            global_bucket: Arc::new(RwLock::new(global_bucket)),
        }
    }

    /// Check if request should be rate limited
    pub async fn check_rate_limit(&self, context: &RateLimitContext) -> RateLimitResult {
        let tokens = context.tokens_requested;

        // Check global limit first (fastest check)
        {
            let mut global_bucket = self.global_bucket.write().await;
            if !global_bucket.try_consume(tokens) {
                warn!("Global rate limit exceeded from IP: {}", context.ip_addr);
                return RateLimitResult::GlobalLimitExceeded;
            }
        }

        // Check user-based limits BEFORE IP limits (if authenticated)
        // This ensures AuthFailurePenalty is returned before IpLimitExceeded
        // when both buckets are penalized
        if let Some(user_id) = context.user_id {
            let mut user_buckets = self.user_buckets.write().await;

            let (capacity, rate) = match context.request_type {
                RequestType::OrderPlacement | RequestType::Trading => (
                    self.config.order_burst_capacity,
                    self.config.orders_per_minute,
                ),
                RequestType::Authentication => (
                    self.config.auth_failures_per_minute,
                    self.config.auth_failures_per_minute,
                ),
                _ => (
                    self.config.user_burst_capacity,
                    self.config.user_requests_per_minute,
                ),
            };

            let user_bucket = user_buckets
                .entry(user_id)
                .or_insert_with(|| TokenBucket::new(capacity, rate));

            if !user_bucket.try_consume(tokens) {
                match context.request_type {
                    RequestType::Authentication => {
                        warn!(
                            "Authentication failure rate limit exceeded for user: {}",
                            user_id
                        );
                        return RateLimitResult::AuthFailurePenalty;
                    },
                    RequestType::OrderPlacement | RequestType::Trading => {
                        warn!("Order rate limit exceeded for user: {}", user_id);
                        return RateLimitResult::OrderLimitExceeded;
                    },
                    _ => {
                        warn!("User rate limit exceeded for: {}", user_id);
                        return RateLimitResult::UserLimitExceeded;
                    },
                }
            }
        }

        // Check IP-based limits after user limits
        {
            let mut ip_buckets = self.ip_buckets.write().await;
            let ip_bucket = ip_buckets.entry(context.ip_addr).or_insert_with(|| {
                TokenBucket::new(
                    self.config.ip_burst_capacity,
                    self.config.ip_requests_per_minute,
                )
            });

            if !ip_bucket.try_consume(tokens) {
                warn!("IP rate limit exceeded for: {}", context.ip_addr);
                return RateLimitResult::IpLimitExceeded;
            }
        }

        RateLimitResult::Allowed
    }

    /// Apply penalty for authentication failures
    pub async fn apply_auth_failure_penalty(&self, user_id: Uuid, ip_addr: IpAddr) {
        let penalty_duration =
            Duration::from_secs((self.config.auth_failure_penalty_minutes as u64) * 60);

        // Apply penalty to user bucket
        {
            let mut user_buckets = self.user_buckets.write().await;
            if let Some(user_bucket) = user_buckets.get_mut(&user_id) {
                user_bucket.apply_penalty(penalty_duration);
            }
        }

        // Apply penalty to IP bucket
        {
            let mut ip_buckets = self.ip_buckets.write().await;
            if let Some(ip_bucket) = ip_buckets.get_mut(&ip_addr) {
                ip_bucket.apply_penalty(penalty_duration);
            }
        }

        error!(
            "Applied authentication failure penalty: user={}, ip={}, duration={}min",
            user_id, ip_addr, self.config.auth_failure_penalty_minutes
        );
    }

    /// Get current rate limit status for monitoring
    pub async fn get_rate_limit_status(&self) -> RateLimitStatus {
        let global_bucket = self.global_bucket.read().await;
        let user_buckets = self.user_buckets.read().await;
        let ip_buckets = self.ip_buckets.read().await;

        RateLimitStatus {
            global_tokens_available: global_bucket.tokens,
            global_capacity: global_bucket.capacity,
            active_users: user_buckets.len(),
            active_ips: ip_buckets.len(),
            penalized_users: user_buckets
                .values()
                .filter(|bucket| bucket.penalty_until.is_some())
                .count(),
            penalized_ips: ip_buckets
                .values()
                .filter(|bucket| bucket.penalty_until.is_some())
                .count(),
        }
    }

    /// Start background cleanup task
    pub async fn start_cleanup_task(self: Arc<Self>) {
        let cleanup_interval =
            Duration::from_secs((self.config.cleanup_interval_minutes as u64) * 60);
        let max_age = cleanup_interval * 2; // Keep buckets for 2x cleanup interval

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

            loop {
                interval.tick().await;

                // Cleanup expired user buckets
                {
                    let mut user_buckets = self.user_buckets.write().await;
                    let before_count = user_buckets.len();
                    user_buckets.retain(|_, bucket| !bucket.is_expired(max_age));
                    let after_count = user_buckets.len();

                    if before_count > after_count {
                        info!(
                            "Cleaned up {} expired user rate limit buckets",
                            before_count - after_count
                        );
                    }
                }

                // Cleanup expired IP buckets
                {
                    let mut ip_buckets = self.ip_buckets.write().await;
                    let before_count = ip_buckets.len();
                    ip_buckets.retain(|_, bucket| !bucket.is_expired(max_age));
                    let after_count = ip_buckets.len();

                    if before_count > after_count {
                        info!(
                            "Cleaned up {} expired IP rate limit buckets",
                            before_count - after_count
                        );
                    }
                }

                // Log current status
                let status = self.get_rate_limit_status().await;
                info!("Rate limiter status: global_tokens={:.1}/{:.1}, users={}, ips={}, penalties={}+{}", 
                      status.global_tokens_available, status.global_capacity,
                      status.active_users, status.active_ips,
                      status.penalized_users, status.penalized_ips);
            }
        });
    }
}

/// Current rate limiter status for monitoring
#[derive(Debug, Clone)]
pub struct RateLimitStatus {
    pub global_tokens_available: f64,
    pub global_capacity: f64,
    pub active_users: usize,
    pub active_ips: usize,
    pub penalized_users: usize,
    pub penalized_ips: usize,
}

/// Rate limiter middleware for tonic
use std::task::{Context, Poll};
use tonic::{Request, Response, Status};
use tower::{Layer, Service};

#[derive(Clone)]
pub struct RateLimitLayer {
    rate_limiter: Arc<RateLimiter>,
}

impl RateLimitLayer {
    pub fn new(rate_limiter: Arc<RateLimiter>) -> Self {
        Self { rate_limiter }
    }
}

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

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

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

// Implement NamedService to maintain service name through middleware
// NamedService is used by tonic for gRPC service reflection
impl<S> tonic::server::NamedService for RateLimitService<S>
where
    S: tonic::server::NamedService,
{
    const NAME: &'static str = S::NAME;
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RateLimitService<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
    S::Future: Send + 'static,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>> + From<Status>,
    ReqBody: Send + 'static,
    ResBody: http_body::Body<Data = bytes::Bytes> + Send + 'static,
    ResBody::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    type Response = Response<ResBody>;
    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 Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
        let rate_limiter = self.rate_limiter.clone();
        let mut inner = self.inner.clone();

        Box::pin(async move {
            // Extract IP address and user info from request metadata
            // Parse IP address from headers, fallback to localhost
            // SAFETY: "127.0.0.1" is a valid IP address constant
            const LOCALHOST: std::net::IpAddr =
                std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1));
            let ip_addr = request
                .metadata()
                .get("x-forwarded-for")
                .or_else(|| request.metadata().get("x-real-ip"))
                .and_then(|value| value.to_str().ok())
                .and_then(|addr_str| addr_str.parse().ok())
                .unwrap_or(LOCALHOST);

            let user_id = request
                .metadata()
                .get("user-id")
                .and_then(|value| value.to_str().ok())
                .and_then(|uuid_str| uuid_str.parse().ok());

            // Determine request type based on method path stored in extensions
            let request_type =
                if let Some(method_name) = request.extensions().get::<hyper::Method>() {
                    match method_name.as_str() {
                        path if path.contains("authenticate") => RequestType::Authentication,
                        path if path.contains("place_order") || path.contains("cancel_order") => {
                            RequestType::OrderPlacement
                        },
                        path if path.contains("trading") => RequestType::Trading,
                        path if path.contains("market_data") => RequestType::MarketData,
                        path if path.contains("risk") => RequestType::Risk,
                        path if path.contains("ml") => RequestType::ML,
                        path if path.contains("config") => RequestType::Config,
                        _ => RequestType::General,
                    }
                } else {
                    // Default request type when path cannot be determined
                    RequestType::General
                };

            let context = RateLimitContext {
                user_id,
                ip_addr,
                request_type,
                tokens_requested: 1.0, // Default to 1 token per request
            };

            // Check rate limits
            match rate_limiter.check_rate_limit(&context).await {
                RateLimitResult::Allowed => {
                    // Request allowed, continue to service
                    inner.call(request).await
                },
                RateLimitResult::UserLimitExceeded => Err(Status::resource_exhausted(
                    "User rate limit exceeded. Please slow down.",
                )
                .into()),
                RateLimitResult::IpLimitExceeded => Err(Status::resource_exhausted(
                    "IP rate limit exceeded. Please slow down.",
                )
                .into()),
                RateLimitResult::GlobalLimitExceeded => Err(Status::resource_exhausted(
                    "Service temporarily overloaded. Please retry later.",
                )
                .into()),
                RateLimitResult::AuthFailurePenalty => Err(Status::permission_denied(
                    "Too many authentication failures. Access temporarily restricted.",
                )
                .into()),
                RateLimitResult::OrderLimitExceeded => Err(Status::resource_exhausted(
                    "Order rate limit exceeded. Please slow down trading activity.",
                )
                .into()),
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[tokio::test]
    async fn test_rate_limiter_basic() {
        let config = RateLimitConfig {
            user_requests_per_minute: 10,
            user_burst_capacity: 5,
            ip_requests_per_minute: 20,
            ip_burst_capacity: 10,
            global_requests_per_minute: 100,
            global_burst_capacity: 50,
            ..Default::default()
        };

        let rate_limiter = RateLimiter::new(config);
        let user_id = Uuid::new_v4();
        let ip_addr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));

        let context = RateLimitContext {
            user_id: Some(user_id),
            ip_addr,
            request_type: RequestType::General,
            tokens_requested: 1.0,
        };

        // Should allow initial requests
        for _ in 0..5 {
            let result = rate_limiter.check_rate_limit(&context).await;
            assert!(matches!(result, RateLimitResult::Allowed));
        }

        // Should exceed user burst capacity
        let result = rate_limiter.check_rate_limit(&context).await;
        assert!(matches!(result, RateLimitResult::UserLimitExceeded));
    }

    #[tokio::test]
    async fn test_auth_failure_penalty() {
        let config = RateLimitConfig {
            auth_failures_per_minute: 3,
            auth_failure_penalty_minutes: 1,
            ..Default::default()
        };

        let rate_limiter = RateLimiter::new(config);
        let user_id = Uuid::new_v4();
        let ip_addr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2));

        let context = RateLimitContext {
            user_id: Some(user_id),
            ip_addr,
            request_type: RequestType::Authentication,
            tokens_requested: 1.0,
        };

        // First request should succeed - this creates the buckets
        let result = rate_limiter.check_rate_limit(&context).await;
        assert!(matches!(result, RateLimitResult::Allowed));

        // Now apply penalty to the existing buckets
        rate_limiter
            .apply_auth_failure_penalty(user_id, ip_addr)
            .await;

        // Second request should be blocked due to penalty
        let result = rate_limiter.check_rate_limit(&context).await;
        assert!(
            matches!(result, RateLimitResult::AuthFailurePenalty),
            "Expected AuthFailurePenalty but got {:?}",
            result
        );
    }
}
