//! Middleware system for request processing
//!
//! Provides a simple middleware chain for processing requests with rate limiting

use std::sync::{Arc, RwLock};

use throttle_machines::gcra;

use crate::protocol::errors::McpError;
use crate::server::handler::RequestContext;

/// Middleware function type
///
/// Takes a RequestContext and returns a HandlerResult
pub type MiddlewareFn =
    Arc<dyn Fn(RequestContext) -> Result<RequestContext, McpError> + Send + Sync>;

/// Middleware chain for processing requests
#[derive(Clone)]
pub struct MiddlewareChain {
    middleware: Vec<MiddlewareFn>,
}

impl MiddlewareChain {
    /// Create new empty middleware chain
    pub fn new() -> Self {
        Self {
            middleware: Vec::new(),
        }
    }

    /// Add middleware to the chain
    pub fn add(&mut self, middleware: MiddlewareFn) {
        self.middleware.push(middleware);
    }

    /// Process request through middleware chain
    ///
    /// Each middleware can:
    /// - Modify the context and pass it to next middleware
    /// - Return an error to short-circuit the chain
    pub fn process(&self, mut ctx: RequestContext) -> Result<RequestContext, McpError> {
        for middleware in &self.middleware {
            ctx = middleware(ctx)?;
        }
        Ok(ctx)
    }

    /// Get the number of middleware in the chain
    pub fn len(&self) -> usize {
        self.middleware.len()
    }

    /// Check if chain is empty
    pub fn is_empty(&self) -> bool {
        self.middleware.is_empty()
    }
}

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

/// Helper to create logging middleware
pub fn logging_middleware() -> MiddlewareFn {
    Arc::new(|ctx: RequestContext| {
        tracing::debug!(
            session_id = %ctx.session.id,
            method = %ctx.request.method,
            "Processing request"
        );
        Ok(ctx)
    })
}

/// Helper to create validation middleware
pub fn validation_middleware() -> MiddlewareFn {
    Arc::new(|ctx: RequestContext| {
        // Basic validation - ensure method is not empty
        if ctx.request.method.is_empty() {
            return Err(McpError::validation(
                "empty_method",
                "Method name cannot be empty",
            ));
        }
        Ok(ctx)
    })
}

/// Rate limiter configuration using GCRA algorithm
#[derive(Debug, Clone)]
pub struct RateLimiterConfig {
    /// Requests allowed per second
    pub requests_per_second: f64,
    /// Burst capacity (number of requests that can be made instantly)
    pub burst_capacity: usize,
}

impl Default for RateLimiterConfig {
    fn default() -> Self {
        Self {
            requests_per_second: 100.0,
            burst_capacity: 10,
        }
    }
}

impl RateLimiterConfig {
    /// Create a new rate limiter configuration
    pub fn new(requests_per_second: f64, burst_capacity: usize) -> Self {
        Self {
            requests_per_second,
            burst_capacity,
        }
    }

    /// Calculate GCRA emission interval (time between requests)
    pub fn emission_interval(&self) -> f64 {
        1.0 / self.requests_per_second
    }

    /// Calculate GCRA delay tolerance (for burst support)
    pub fn delay_tolerance(&self) -> f64 {
        self.burst_capacity as f64 * self.emission_interval()
    }
}

/// Thread-safe rate limiter using GCRA algorithm
#[derive(Clone)]
pub struct RateLimiter {
    config: RateLimiterConfig,
    /// Theoretical Arrival Time (TAT) - shared state
    tat: Arc<RwLock<f64>>,
}

impl RateLimiter {
    /// Create a new rate limiter with the given configuration
    pub fn new(config: RateLimiterConfig) -> Self {
        Self {
            config,
            tat: Arc::new(RwLock::new(0.0)),
        }
    }

    /// Create rate limiter with defaults (100 req/s, burst of 10)
    pub fn default_limiter() -> Self {
        Self::new(RateLimiterConfig::default())
    }

    /// Check if a request is allowed
    ///
    /// Returns Ok(()) if allowed, Err(retry_after_secs) if rate limited
    pub fn check(&self) -> Result<(), f64> {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs_f64())
            .unwrap_or(0.0);

        let emission_interval = self.config.emission_interval();
        let delay_tolerance = self.config.delay_tolerance();

        let mut tat_guard = self.tat.write().map_err(|_| 1.0)?; // On lock error, ask to retry after 1 second

        let result = gcra::check(*tat_guard, now, emission_interval, delay_tolerance);

        if result.allowed {
            *tat_guard = result.new_tat;
            Ok(())
        } else {
            Err(result.retry_after)
        }
    }

    /// Get current capacity remaining (approximate)
    pub fn remaining_capacity(&self) -> usize {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs_f64())
            .unwrap_or(0.0);

        let delay_tolerance = self.config.delay_tolerance();
        let emission_interval = self.config.emission_interval();

        let tat = self.tat.read().map(|t| *t).unwrap_or(0.0);

        let result = gcra::peek(tat, now, delay_tolerance);

        if result.allowed {
            // Estimate remaining based on how much tolerance is left
            let remaining_tolerance = delay_tolerance - (result.new_tat - now).max(0.0);
            (remaining_tolerance / emission_interval).floor() as usize + 1
        } else {
            0
        }
    }
}

/// Helper to create rate limiter middleware
///
/// Creates middleware that enforces rate limits using GCRA algorithm.
/// Returns `McpError::rate_limited()` when limit is exceeded.
pub fn rate_limiter_middleware(limiter: Arc<RateLimiter>) -> MiddlewareFn {
    Arc::new(move |ctx: RequestContext| match limiter.check() {
        Ok(()) => {
            tracing::trace!(
                session_id = %ctx.session.id,
                method = %ctx.request.method,
                "Request allowed by rate limiter"
            );
            Ok(ctx)
        }
        Err(retry_after) => {
            tracing::warn!(
                session_id = %ctx.session.id,
                method = %ctx.request.method,
                retry_after_secs = %retry_after,
                "Request rate limited"
            );
            Err(McpError::rate_limited(format!(
                "Rate limit exceeded. Retry after {:.2} seconds",
                retry_after
            )))
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::types::JsonRpcRequest;
    use crate::server::session::Session;
    use serde_json::Value;

    #[test]
    fn test_empty_chain() {
        let chain = MiddlewareChain::new();
        assert!(chain.is_empty());
        assert_eq!(chain.len(), 0);
    }

    #[test]
    fn test_add_middleware() {
        let mut chain = MiddlewareChain::new();

        let mw = Arc::new(|ctx: RequestContext| Ok(ctx));
        chain.add(mw);

        assert!(!chain.is_empty());
        assert_eq!(chain.len(), 1);
    }

    #[test]
    fn test_process_middleware() {
        let mut chain = MiddlewareChain::new();

        // Add middleware that modifies session state
        let mw = Arc::new(|ctx: RequestContext| {
            ctx.session.set_state("processed", Value::Bool(true));
            Ok(ctx)
        });
        chain.add(mw);

        let session = Session::new();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "test".to_string(),
            params: None,
        };
        let ctx = RequestContext::new(session, request);

        let result = chain.process(ctx).unwrap();
        assert_eq!(
            result.session.get_state("processed"),
            Some(Value::Bool(true))
        );
    }

    #[test]
    fn test_multiple_middleware() {
        let mut chain = MiddlewareChain::new();

        // Add multiple middleware
        let mw1 = Arc::new(|ctx: RequestContext| {
            ctx.session.set_state("step1", Value::Bool(true));
            Ok(ctx)
        });
        let mw2 = Arc::new(|ctx: RequestContext| {
            ctx.session.set_state("step2", Value::Bool(true));
            Ok(ctx)
        });
        chain.add(mw1);
        chain.add(mw2);

        let session = Session::new();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "test".to_string(),
            params: None,
        };
        let ctx = RequestContext::new(session, request);

        let result = chain.process(ctx).unwrap();
        assert_eq!(result.session.get_state("step1"), Some(Value::Bool(true)));
        assert_eq!(result.session.get_state("step2"), Some(Value::Bool(true)));
    }

    #[test]
    fn test_middleware_error() {
        let mut chain = MiddlewareChain::new();

        // Add middleware that returns error
        let mw =
            Arc::new(|_ctx: RequestContext| Err(McpError::validation("test_error", "Test error")));
        chain.add(mw);

        let session = Session::new();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "test".to_string(),
            params: None,
        };
        let ctx = RequestContext::new(session, request);

        let result = chain.process(ctx);
        assert!(result.is_err());
    }

    #[test]
    fn test_validation_middleware() {
        let mw = validation_middleware();

        // Valid request
        let session = Session::new();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "test".to_string(),
            params: None,
        };
        let ctx = RequestContext::new(session, request);
        assert!(mw(ctx).is_ok());

        // Invalid request (empty method)
        let session = Session::new();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::Number(1.into())),
            method: "".to_string(),
            params: None,
        };
        let ctx = RequestContext::new(session, request);
        assert!(mw(ctx).is_err());
    }
}
