mockforge-http 0.3.226

//! Middleware/utilities to apply latency and failure injection per operation.
use axum::body::Body;
use axum::extract::State;
use axum::http::{Request, StatusCode};
use axum::middleware::Next;
use axum::response::Response;

use crate::latency_profiles::LatencyProfiles;
use mockforge_chaos::core_failure_injection::FailureInjector;
use mockforge_chaos::core_traffic_shaping::TrafficShaper;

/// Metadata for the current OpenAPI operation
#[derive(Clone)]
pub struct OperationMeta {
    /// OpenAPI operation ID
    pub id: String,
    /// Tags associated with this operation
    pub tags: Vec<String>,
    /// API path pattern
    pub path: String,
}

/// Shared state for operation middleware
#[derive(Clone)]
pub struct Shared {
    /// Latency profiles for request simulation
    pub profiles: LatencyProfiles,
    /// Optional failure injector for chaos engineering
    pub failure_injector: Option<FailureInjector>,
    /// Optional traffic shaper for bandwidth/loss simulation
    pub traffic_shaper: Option<TrafficShaper>,
    /// Whether traffic shaping is enabled
    pub traffic_shaping_enabled: bool,
}

/// Middleware to add shared state to request extensions
pub async fn add_shared_extension(
    State(shared): State<Shared>,
    mut req: Request<Body>,
    next: Next,
) -> Response {
    req.extensions_mut().insert(shared);
    next.run(req).await
}

/// Middleware to apply fault injection before processing request
pub async fn fault_then_next(req: Request<Body>, next: Next) -> Response {
    let shared = match req.extensions().get::<Shared>() {
        Some(s) => s.clone(),
        None => {
            tracing::error!("Shared extension not found in request - ensure add_shared_extension middleware is configured");
            let mut res =
                Response::new(Body::from("Internal server error: middleware misconfiguration"));
            *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
            return res;
        }
    };
    let op = req.extensions().get::<OperationMeta>().cloned();

    // First, check the new enhanced failure injection system
    if let Some(failure_injector) = &shared.failure_injector {
        let tags = op.as_ref().map(|o| o.tags.as_slice()).unwrap_or(&[]);
        if let Some((status_code, error_message)) = failure_injector.process_request(tags) {
            let mut res = Response::new(Body::from(error_message));
            *res.status_mut() =
                StatusCode::from_u16(status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
            return res;
        }
    }

    // Fallback to legacy latency profiles system for backward compatibility
    if let Some(op) = &op {
        if let Some((code, msg)) = shared
            .profiles
            .maybe_fault(&op.id, &op.tags.iter().map(|s| s.to_string()).collect::<Vec<_>>())
            .await
        {
            let mut res = Response::new(Body::from(msg));
            *res.status_mut() =
                StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
            return res;
        }
    }

    // Apply traffic shaping (bandwidth throttling and burst loss) to the request
    if shared.traffic_shaping_enabled {
        if let Some(traffic_shaper) = &shared.traffic_shaper {
            // Calculate request size for bandwidth throttling
            let request_size = calculate_request_size(&req);

            let tags = op.as_ref().map(|o| o.tags.as_slice()).unwrap_or(&[]);

            // Apply traffic shaping
            match traffic_shaper.process_transfer(request_size, tags).await {
                Ok(Some(_timeout)) => {
                    // Request was "lost" due to burst loss - return timeout error
                    let mut res =
                        Response::new(Body::from("Request timeout due to traffic shaping"));
                    *res.status_mut() = StatusCode::REQUEST_TIMEOUT;
                    return res;
                }
                Ok(None) => {
                    // Transfer allowed, continue
                }
                Err(e) => {
                    // Traffic shaping error - return internal server error
                    let mut res =
                        Response::new(Body::from(format!("Traffic shaping error: {}", e)));
                    *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
                    return res;
                }
            }
        }
    }

    let (parts, body) = req.into_parts();
    let req = Request::from_parts(parts, body);

    let response = next.run(req).await;

    // Apply traffic shaping to the response
    if shared.traffic_shaping_enabled {
        if let Some(traffic_shaper) = &shared.traffic_shaper {
            // Calculate response size for bandwidth throttling
            let response_size = calculate_response_size(&response);

            let tags = op.as_ref().map(|o| o.tags.as_slice()).unwrap_or(&[]);

            // Apply traffic shaping to response
            match traffic_shaper.process_transfer(response_size, tags).await {
                Ok(Some(_timeout)) => {
                    // Response was "lost" due to burst loss - return timeout error
                    let mut res =
                        Response::new(Body::from("Response timeout due to traffic shaping"));
                    *res.status_mut() = StatusCode::GATEWAY_TIMEOUT;
                    return res;
                }
                Ok(None) => {
                    // Transfer allowed, continue
                }
                Err(e) => {
                    // Traffic shaping error - return internal server error
                    let mut res =
                        Response::new(Body::from(format!("Traffic shaping error: {}", e)));
                    *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
                    return res;
                }
            }
        }
    }

    response
}

/// Calculate the approximate size of an HTTP request for bandwidth throttling
fn calculate_request_size<B>(req: &Request<B>) -> u64 {
    let mut size = 0u64;

    // Add header sizes
    for (name, value) in req.headers() {
        size += name.as_str().len() as u64;
        size += value.as_bytes().len() as u64;
    }

    // Add URI size
    size += req.uri().to_string().len() as u64;

    // Use Content-Length header for body size when available
    if let Some(content_length) = req.headers().get(http::header::CONTENT_LENGTH) {
        if let Ok(len_str) = content_length.to_str() {
            if let Ok(len) = len_str.parse::<u64>() {
                size += len;
                return size;
            }
        }
    }

    // Fallback: estimate body size from method (GET/HEAD/DELETE typically have no body)
    let method = req.method();
    if method == http::Method::POST || method == http::Method::PUT || method == http::Method::PATCH
    {
        size += 256; // Conservative estimate for requests without Content-Length
    }

    size
}

/// Calculate the approximate size of an HTTP response for bandwidth throttling
fn calculate_response_size(res: &Response) -> u64 {
    let mut size = 0u64;

    // Add header sizes
    for (name, value) in res.headers() {
        size += name.as_str().len() as u64;
        size += value.as_bytes().len() as u64;
    }

    // Add status line size
    size += 15; // "HTTP/1.1 200 OK\r\n"

    // Use Content-Length header for body size when available
    if let Some(content_length) = res.headers().get(http::header::CONTENT_LENGTH) {
        if let Ok(len_str) = content_length.to_str() {
            if let Ok(len) = len_str.parse::<u64>() {
                size += len;
                return size;
            }
        }
    }

    // Fallback: estimate from status code (204/304 have no body)
    match res.status().as_u16() {
        204 | 304 => {}   // No body
        _ => size += 256, // Conservative estimate for responses without Content-Length
    }

    size
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::{Request, Response, StatusCode};

    #[test]
    fn test_operation_meta_creation() {
        let meta = OperationMeta {
            id: "getUserById".to_string(),
            tags: vec!["users".to_string(), "public".to_string()],
            path: "/users/{id}".to_string(),
        };

        assert_eq!(meta.id, "getUserById");
        assert_eq!(meta.tags.len(), 2);
        assert_eq!(meta.path, "/users/{id}");
    }

    #[test]
    fn test_shared_creation() {
        let shared = Shared {
            profiles: LatencyProfiles::default(),
            failure_injector: None,
            traffic_shaper: None,
            traffic_shaping_enabled: false,
        };

        assert!(!shared.traffic_shaping_enabled);
        assert!(shared.failure_injector.is_none());
        assert!(shared.traffic_shaper.is_none());
    }

    #[test]
    fn test_shared_with_failure_injector() {
        let failure_injector = FailureInjector::new(None, true);
        let shared = Shared {
            profiles: LatencyProfiles::default(),
            failure_injector: Some(failure_injector),
            traffic_shaper: None,
            traffic_shaping_enabled: false,
        };

        assert!(shared.failure_injector.is_some());
    }

    #[test]
    fn test_calculate_request_size_basic() {
        let req = Request::builder()
            .uri("/test")
            .header("content-type", "application/json")
            .body(())
            .unwrap();

        let size = calculate_request_size(&req);

        // Should be > 0 (includes headers + URI + body estimate)
        assert!(size > 0);
        // Should include at least the URI and header sizes
        assert!(size >= "/test".len() as u64 + "content-type".len() as u64);
    }

    #[test]
    fn test_calculate_request_size_with_multiple_headers() {
        let req = Request::builder()
            .uri("/api/users")
            .header("content-type", "application/json")
            .header("authorization", "Bearer token123")
            .header("user-agent", "test-client")
            .body(())
            .unwrap();

        let size = calculate_request_size(&req);

        // Should account for all headers (URI + header names + header values)
        assert!(size > 50);
    }

    #[test]
    fn test_calculate_response_size_basic() {
        let res = Response::builder()
            .status(StatusCode::OK)
            .header("content-type", "application/json")
            .body(Body::empty())
            .unwrap();

        let size = calculate_response_size(&res);

        // Should be > 0 (includes status line + headers + body estimate)
        assert!(size > 0);
        // Should include at least the status line estimate (50) and header sizes
        assert!(size >= 50);
    }

    #[test]
    fn test_calculate_response_size_with_multiple_headers() {
        let res = Response::builder()
            .status(StatusCode::OK)
            .header("content-type", "application/json")
            .header("cache-control", "no-cache")
            .header("x-request-id", "123-456-789")
            .body(Body::empty())
            .unwrap();

        let size = calculate_response_size(&res);

        // Should account for all headers
        assert!(size > 100);
    }

    #[test]
    fn test_shared_clone() {
        let shared = Shared {
            profiles: LatencyProfiles::default(),
            failure_injector: None,
            traffic_shaper: None,
            traffic_shaping_enabled: true,
        };

        let cloned = shared.clone();

        assert_eq!(shared.traffic_shaping_enabled, cloned.traffic_shaping_enabled);
    }

    #[test]
    fn test_operation_meta_clone() {
        let meta = OperationMeta {
            id: "testOp".to_string(),
            tags: vec!["tag1".to_string()],
            path: "/test".to_string(),
        };

        let cloned = meta.clone();

        assert_eq!(meta.id, cloned.id);
        assert_eq!(meta.tags, cloned.tags);
        assert_eq!(meta.path, cloned.path);
    }
}