iron-shell 14.0.2

use crate::ai_api::{
    config::{APIConfig, RateLimitConfig, ServerConfig},
    models::*,
    AIModelAPI,
};
use anyhow::Result;
use axum::{
    extract::{ConnectInfo, Path, Query, State},
    http::{HeaderMap, StatusCode},
    response::{IntoResponse, Json, Response, Sse},
    routing::{get, post},
    Router,
};
use futures::stream::Stream;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::convert::Infallible;
use std::net::{IpAddr, SocketAddr};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::net::TcpListener;
use tower::ServiceBuilder;
use tower_http::{
    cors::{Any, CorsLayer},
    timeout::TimeoutLayer,
    trace::TraceLayer,
};
use utoipa::ToSchema;
use utoipa_swagger_ui::SwaggerUi;

/// HTTP Server for the AI Model API
pub struct APIServer {
    config: APIConfig,
    api: Arc<AIModelAPI>,
}

/// Shared application state
#[derive(Clone)]
pub struct AppState {
    pub api: Arc<AIModelAPI>,
    pub config: APIConfig,
    pub start_time: std::time::Instant,
}

/// Query parameters for model listing
#[derive(Debug, Deserialize, ToSchema)]
pub struct ModelsQuery {
    /// Filter by provider
    pub provider: Option<String>,
    /// Filter by capability
    pub capability: Option<String>,
    /// Include local models only
    pub local_only: Option<bool>,
}

/// Query parameters for pagination
#[derive(Debug, Deserialize, ToSchema)]
pub struct PaginationQuery {
    /// Page number (0-based)
    pub page: Option<usize>,
    /// Items per page
    pub limit: Option<usize>,
}

impl APIServer {
    pub fn new(config: APIConfig) -> Result<Self> {
        let api = Arc::new(AIModelAPI::new(config.clone())?);
        Ok(Self { config, api })
    }

    /// Start the HTTP server
    pub async fn start(&self) -> Result<()> {
        // Refused before anything is bound: a non-loopback address without keys,
        // or TLS, which this server does not implement.
        crate::ai_api::config::check_startup(&self.config.server, &self.config.security)?;
        let app = self.create_router().await?;

        let addr = format!("{}:{}", self.config.server.host, self.config.server.port);
        let listener = TcpListener::bind(&addr).await?;

        tracing::info!("AI Model API server starting on {}", addr);

        crate::serve::serve_until(
            listener,
            app,
            crate::serve::shutdown_signal(),
            crate::serve::DRAIN,
        )
        .await?;

        Ok(())
    }

    /// Create CORS layer with configured allowed origins
    fn create_cors_layer(&self) -> CorsLayer {
        use axum::http::Method;

        // SECURITY FIX (LOW-003): Configure strict CORS
        let cors = CorsLayer::new();

        // Parse allowed origins from config
        let origins: Vec<_> = self
            .config
            .server
            .cors_origins
            .iter()
            .filter_map(|origin| {
                if origin == "*" {
                    None // Handle wildcard separately
                } else {
                    origin.parse::<axum::http::HeaderValue>().ok()
                }
            })
            .collect();

        let cors =
            if origins.is_empty() && self.config.server.cors_origins.contains(&"*".to_string()) {
                // If wildcard is specified, allow any origin (less secure but flexible)
                cors.allow_origin(Any)
            } else {
                // Use specific origin allowlist (more secure)
                cors.allow_origin(origins)
            };

        cors
            // Allow common HTTP methods for API
            .allow_methods([
                Method::GET,
                Method::POST,
                Method::PUT,
                Method::DELETE,
                Method::OPTIONS,
            ])
            // Allow common headers
            .allow_headers([
                axum::http::header::CONTENT_TYPE,
                axum::http::header::AUTHORIZATION,
                axum::http::header::ACCEPT,
            ])
            // Allow credentials (cookies, auth headers)
            .allow_credentials(false) // Set to true only if needed with specific origins
            // Cache preflight requests for 1 hour
            .max_age(std::time::Duration::from_secs(3600))
    }

    /// Create the main router with all endpoints.
    ///
    /// Two layers wrap every route, including the Swagger UI: the `Host` check
    /// (outermost, so a foreign host gets no response body at all) and the
    /// per-IP rate limit when `security.rate_limiting.enabled` is set.
    pub async fn create_router(&self) -> Result<Router> {
        let state = AppState {
            api: self.api.clone(),
            config: self.config.clone(),
            start_time: std::time::Instant::now(),
        };

        let api_routes = Router::new()
            // OpenAI-compatible endpoints
            .route("/models", get(list_models))
            .route("/chat/completions", post(chat_completions))
            .route("/embeddings", post(embeddings))
            // Extended endpoints for model management
            .route("/models/:model_id", get(get_model))
            .route("/models/:model_id/download", post(download_model))
            .route("/models/:model_id/convert", post(convert_model))
            .route("/models/:model_id", axum::routing::delete(delete_model))
            // Provider management
            .route("/providers", get(list_providers))
            .route("/providers/:provider_id/validate", post(validate_provider))
            // Storage and caching
            .route("/storage/stats", get(storage_stats))
            .route("/storage/cleanup", post(cleanup_storage))
            // Health and status
            .route("/health", get(health_check))
            .route("/status", get(server_status))
            .with_state(state);

        let mut app = Router::new()
            .nest("/v1", api_routes)
            .layer(
                ServiceBuilder::new()
                    .layer(TraceLayer::new_for_http())
                    .layer(TimeoutLayer::new(std::time::Duration::from_secs(
                        self.config.server.request_timeout_seconds,
                    )))
                    // SECURITY FIX (LOW-003): Configure strict CORS with allowlist
                    .layer(if self.config.server.enable_cors {
                        self.create_cors_layer()
                    } else {
                        CorsLayer::new()
                    }),
            )
            // SECURITY FIX (LOW-001): Apply security headers middleware
            .layer(axum::middleware::from_fn(add_security_headers));

        // Add OpenAPI documentation if enabled
        if self.config.server.enable_openapi {
            app = app.merge(
                SwaggerUi::new("/swagger-ui").url(
                    "/api-docs/openapi.json",
                    utoipa::openapi::OpenApiBuilder::new()
                        .info(
                            utoipa::openapi::InfoBuilder::new()
                                .title("AI Model API")
                                .version("1.0.0")
                                .build(),
                        )
                        .build(),
                ),
            );
        }

        if self.config.security.rate_limiting.enabled {
            let limiter = Arc::new(RateLimiter::new(&self.config.security.rate_limiting));
            app = app.layer(axum::middleware::from_fn_with_state(
                limiter,
                rate_limit_middleware,
            ));
        }

        let hosts = Arc::new(HostPolicy::new(&self.config.server));
        app = app.layer(axum::middleware::from_fn_with_state(
            hosts,
            host_guard_middleware,
        ));

        Ok(app)
    }
}

// API Handlers

/// List available models
async fn list_models(
    State(state): State<AppState>,
    Query(query): Query<ModelsQuery>,
) -> Result<Json<Vec<ModelInfo>>, (StatusCode, Json<APIError>)> {
    match state.api.list_models().await {
        Ok(mut models) => {
            // Apply filters
            if let Some(provider) = &query.provider {
                models.retain(|m| &m.provider == provider);
            }

            if let Some(capability) = &query.capability {
                models.retain(|m| match capability.as_str() {
                    "chat" => m.capabilities.chat,
                    "embeddings" => m.capabilities.embeddings,
                    "image_generation" => m.capabilities.image_generation,
                    "image_understanding" => m.capabilities.image_understanding,
                    "function_calling" => m.capabilities.function_calling,
                    _ => true,
                });
            }

            if query.local_only.unwrap_or(false) {
                models.retain(|m| m.local_path.is_some());
            }

            Ok(Json(models))
        }
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: e.to_string(),
                    error_type: "internal_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )),
    }
}

/// Get specific model information
async fn get_model(
    State(state): State<AppState>,
    Path(model_id): Path<String>,
) -> Result<Json<ModelInfo>, (StatusCode, Json<APIError>)> {
    match state.api.list_models().await {
        Ok(models) => {
            if let Some(model) = models.iter().find(|m| m.id == model_id) {
                Ok(Json(model.clone()))
            } else {
                Err((
                    StatusCode::NOT_FOUND,
                    Json(APIError {
                        error: ErrorDetail {
                            message: format!("Model {} not found", model_id),
                            error_type: "not_found".to_string(),
                            param: Some("model_id".to_string()),
                            code: None,
                        },
                    }),
                ))
            }
        }
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: e.to_string(),
                    error_type: "internal_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )),
    }
}

/// Chat completions endpoint - supports both streaming and non-streaming
async fn chat_completions(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(request): Json<ChatCompletionRequest>,
) -> Response {
    // Validate API key if required
    if state.config.security.require_api_key {
        if let Err((status, json)) = validate_api_key(&headers, &state.config) {
            return (status, json).into_response();
        }
    }

    // Check if streaming is requested
    if request.stream.unwrap_or(false) {
        // Streaming response
        match stream_chat_completion(state.api.clone(), request).await {
            Ok(stream) => Sse::new(stream)
                .keep_alive(
                    axum::response::sse::KeepAlive::new()
                        .interval(Duration::from_secs(15))
                        .text("keep-alive"),
                )
                .into_response(),
            Err(e) => {
                let error = APIError {
                    error: ErrorDetail {
                        message: e.to_string(),
                        error_type: "streaming_error".to_string(),
                        param: None,
                        code: None,
                    },
                };
                (StatusCode::INTERNAL_SERVER_ERROR, Json(error)).into_response()
            }
        }
    } else {
        // Non-streaming response
        match state.api.chat_completion(request).await {
            Ok(response) => Json(response).into_response(),
            Err(e) => {
                let error = APIError {
                    error: ErrorDetail {
                        message: e.to_string(),
                        error_type: "invalid_request_error".to_string(),
                        param: None,
                        code: None,
                    },
                };
                (StatusCode::BAD_REQUEST, Json(error)).into_response()
            }
        }
    }
}

/// Embeddings endpoint
async fn embeddings(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(request): Json<EmbeddingRequest>,
) -> Result<Json<EmbeddingResponse>, (StatusCode, Json<APIError>)> {
    // Validate API key if required
    if state.config.security.require_api_key {
        validate_api_key(&headers, &state.config)?
    }

    match state.api.embeddings(request).await {
        Ok(response) => Ok(Json(response)),
        Err(e) => Err((
            StatusCode::BAD_REQUEST,
            Json(APIError {
                error: ErrorDetail {
                    message: e.to_string(),
                    error_type: "invalid_request_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )),
    }
}

/// Download model request body
#[derive(Debug, Deserialize)]
struct DownloadModelRequest {
    source: Option<String>,
    format: Option<String>,
    quantization: Option<String>,
}

/// Download model endpoint
async fn download_model(
    State(state): State<AppState>,
    Path(model_id): Path<String>,
    headers: HeaderMap,
    peer: Option<ConnectInfo<SocketAddr>>,
    body: Option<Json<DownloadModelRequest>>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<APIError>)> {
    use crate::ai_api::downloader::{DownloadRequest, ModelDownloader};
    use crate::ai_api::models::{ModelFormat, ModelSource};
    use crate::ai_api::storage::{ModelStorage, StorageConfig};

    authorize_mutation(&headers, &state.config, peer_addr(peer))?;

    let req_body = body.map(|b| b.0).unwrap_or_else(|| DownloadModelRequest {
        source: None,
        format: None,
        quantization: None,
    });

    // Initialize storage
    let storage_config = StorageConfig::default();
    let storage = ModelStorage::new(&storage_config).map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to initialize storage: {}", e),
                    error_type: "storage_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
    })?;

    // Initialize downloader
    let mut downloader = ModelDownloader::new(storage).map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to initialize downloader: {}", e),
                    error_type: "downloader_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
    })?;

    // Determine source (default to huggingface)
    let source = req_body.source.unwrap_or_else(|| "huggingface".to_string());

    // For a "url" source the path segment is the URL itself. Its stored id is
    // derived from the URL; for any other source the segment is the id, and it
    // must pass the id rules before it can name a file.
    let (storage_id, url) = if source == "url" {
        let id = crate::ai_api::storage::url_model_id(&model_id).map_err(|e| {
            (
                StatusCode::BAD_REQUEST,
                Json(APIError {
                    error: ErrorDetail {
                        message: format!("Invalid model URL: {}", e),
                        error_type: "invalid_request_error".to_string(),
                        param: Some("model_id".to_string()),
                        code: None,
                    },
                }),
            )
        })?;
        (id, Some(model_id.clone()))
    } else {
        validate_id_or_bad_request(&model_id)?;
        (model_id.clone(), None)
    };
    let format_preference =
        req_body
            .format
            .as_ref()
            .and_then(|f| match f.to_lowercase().as_str() {
                "gguf" => Some(ModelFormat::GGUF),
                "safetensors" => Some(ModelFormat::SafeTensors),
                "pytorch" => Some(ModelFormat::PyTorch),
                "onnx" => Some(ModelFormat::ONNX),
                _ => None,
            });

    let download_request = DownloadRequest {
        model_id: storage_id,
        source: ModelSource {
            origin: source.clone(),
            url,
            repository: if source == "huggingface" {
                Some(model_id.clone())
            } else {
                None
            },
            commit: None,
            license: None,
        },
        format_preference,
        quantization: req_body.quantization,
        validate_checksum: true,
    };

    // Download the model
    match downloader.download_model(download_request).await {
        Ok(metadata) => Ok(Json(serde_json::json!({
            "success": true,
            "model_id": metadata.id,
            "size_bytes": metadata.size_bytes,
            "format": format!("{:?}", metadata.format),
            "file_path": metadata.file_path,
            "sha256": metadata.sha256,
            "downloaded_at": metadata.downloaded_at.to_rfc3339()
        }))),
        Err(e) => Err((
            StatusCode::BAD_REQUEST,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to download model: {}", e),
                    error_type: "download_error".to_string(),
                    param: Some("model_id".to_string()),
                    code: None,
                },
            }),
        )),
    }
}

/// Convert model format request body
#[derive(Debug, Deserialize)]
struct ConvertModelRequest {
    target_format: String,
    output_id: Option<String>,
    #[allow(dead_code)]
    quantization: Option<String>,
}

/// Convert model format endpoint
async fn convert_model(
    State(state): State<AppState>,
    Path(model_id): Path<String>,
    headers: HeaderMap,
    peer: Option<ConnectInfo<SocketAddr>>,
    Json(request): Json<ConvertModelRequest>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<APIError>)> {
    use crate::ai_api::converters::ModelConverter;
    use crate::ai_api::models::ModelFormat;
    use crate::ai_api::storage::{ModelStorage, StorageConfig};

    authorize_mutation(&headers, &state.config, peer_addr(peer))?;

    // Initialize storage
    let storage_config = StorageConfig::default();
    let storage = ModelStorage::new(&storage_config).map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to initialize storage: {}", e),
                    error_type: "storage_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
    })?;

    // Check if model exists
    let source_metadata = storage.get_model_metadata(&model_id).ok_or_else(|| {
        (
            StatusCode::NOT_FOUND,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Model {} not found", model_id),
                    error_type: "not_found".to_string(),
                    param: Some("model_id".to_string()),
                    code: None,
                },
            }),
        )
    })?;

    // Parse target format
    let target_format = match request.target_format.to_lowercase().as_str() {
        "gguf" => ModelFormat::GGUF,
        "safetensors" => ModelFormat::SafeTensors,
        "pytorch" | "pt" => ModelFormat::PyTorch,
        "onnx" => ModelFormat::ONNX,
        _ => {
            return Err((
                StatusCode::BAD_REQUEST,
                Json(APIError {
                    error: ErrorDetail {
                        message: format!("Unsupported target format: {}", request.target_format),
                        error_type: "invalid_format".to_string(),
                        param: Some("target_format".to_string()),
                        code: None,
                    },
                }),
            ))
        }
    };

    // Initialize converter
    let converter = ModelConverter::new();
    let output_id = request
        .output_id
        .unwrap_or_else(|| format!("{}_{}", model_id, request.target_format));

    // Check conversion feasibility
    if !converter.can_convert(&source_metadata.format, &target_format) {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(APIError {
                error: ErrorDetail {
                    message: format!(
                        "Cannot convert from {:?} to {:?}",
                        source_metadata.format, target_format
                    ),
                    error_type: "conversion_not_supported".to_string(),
                    param: None,
                    code: None,
                },
            }),
        ));
    }

    // Return info about conversion (actual conversion would be async/long-running)
    Ok(Json(serde_json::json!({
        "status": "queued",
        "source_model": model_id,
        "source_format": format!("{:?}", source_metadata.format),
        "target_format": format!("{:?}", target_format),
        "output_id": output_id,
        "message": "Model conversion has been queued. This may take several minutes depending on model size."
    })))
}

/// Delete model endpoint
async fn delete_model(
    State(state): State<AppState>,
    Path(model_id): Path<String>,
    headers: HeaderMap,
    peer: Option<ConnectInfo<SocketAddr>>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<APIError>)> {
    use crate::ai_api::storage::{ModelStorage, StorageConfig};

    authorize_mutation(&headers, &state.config, peer_addr(peer))?;
    validate_id_or_bad_request(&model_id)?;

    // Initialize storage
    let storage_config = StorageConfig::default();
    let mut storage = ModelStorage::new(&storage_config).map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to initialize storage: {}", e),
                    error_type: "storage_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
    })?;

    // Delete the model
    match storage.remove_model(&model_id) {
        Ok(()) => Ok(Json(serde_json::json!({
            "success": true,
            "model_id": model_id,
            "message": "Model successfully deleted"
        }))),
        Err(e) => Err((
            StatusCode::NOT_FOUND,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to delete model: {}", e),
                    error_type: "not_found".to_string(),
                    param: Some("model_id".to_string()),
                    code: None,
                },
            }),
        )),
    }
}

/// List providers endpoint
async fn list_providers(
    State(_state): State<AppState>,
) -> Result<Json<Vec<ProviderInfo>>, (StatusCode, Json<APIError>)> {
    let providers = vec![
        ProviderInfo {
            id: "openai".to_string(),
            name: "OpenAI".to_string(),
            description: "OpenAI GPT models".to_string(),
            enabled: true,
            status: "active".to_string(),
        },
        ProviderInfo {
            id: "anthropic".to_string(),
            name: "Anthropic".to_string(),
            description: "Anthropic Claude models".to_string(),
            enabled: true,
            status: "active".to_string(),
        },
        ProviderInfo {
            id: "local".to_string(),
            name: "Local Models".to_string(),
            description: "Locally hosted models".to_string(),
            enabled: true,
            status: "active".to_string(),
        },
    ];

    Ok(Json(providers))
}

/// Validate provider endpoint
async fn validate_provider(
    State(state): State<AppState>,
    Path(provider_id): Path<String>,
    headers: HeaderMap,
    peer: Option<ConnectInfo<SocketAddr>>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<APIError>)> {
    // Reports whether provider keys are set and probes a local port, so it is
    // not open to an anonymous caller.
    authorize_mutation(&headers, &state.config, peer_addr(peer))?;

    // Check if provider is configured and test connectivity
    let valid_providers = ["openai", "anthropic", "local", "ollama", "huggingface"];

    if !valid_providers.contains(&provider_id.as_str()) {
        return Ok(Json(serde_json::json!({
            "valid": false,
            "provider_id": provider_id,
            "error": "Unknown provider"
        })));
    }

    // Check for API keys based on provider
    let (has_credentials, test_result) = match provider_id.as_str() {
        "openai" => {
            let has_key = std::env::var("OPENAI_API_KEY").is_ok();
            (
                has_key,
                if has_key {
                    "API key configured"
                } else {
                    "Missing OPENAI_API_KEY"
                },
            )
        }
        "anthropic" => {
            let has_key = std::env::var("ANTHROPIC_API_KEY").is_ok();
            (
                has_key,
                if has_key {
                    "API key configured"
                } else {
                    "Missing ANTHROPIC_API_KEY"
                },
            )
        }
        "local" => (true, "Local models always available"),
        "ollama" => {
            // Check if Ollama is running
            let ollama_running = std::net::TcpStream::connect("127.0.0.1:11434").is_ok();
            (
                ollama_running,
                if ollama_running {
                    "Ollama is running"
                } else {
                    "Ollama not accessible on port 11434"
                },
            )
        }
        "huggingface" => {
            let has_token =
                std::env::var("HF_TOKEN").is_ok() || std::env::var("HUGGINGFACE_TOKEN").is_ok();
            (
                true,
                if has_token {
                    "HuggingFace token configured"
                } else {
                    "No token (public models only)"
                },
            )
        }
        _ => (false, "Unknown provider"),
    };

    Ok(Json(serde_json::json!({
        "valid": has_credentials,
        "provider_id": provider_id,
        "status": test_result,
        "checked_at": chrono::Utc::now().to_rfc3339()
    })))
}

/// Storage statistics endpoint
async fn storage_stats(
    State(state): State<AppState>,
    headers: HeaderMap,
    peer: Option<ConnectInfo<SocketAddr>>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<APIError>)> {
    use crate::ai_api::storage::{ModelStorage, StorageConfig};

    // The response carries the data and config directories and the model list.
    authorize_mutation(&headers, &state.config, peer_addr(peer))?;

    // Initialize storage
    let storage_config = StorageConfig::default();
    let storage = ModelStorage::new(&storage_config).map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to initialize storage: {}", e),
                    error_type: "storage_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
    })?;

    let stats = storage.get_storage_stats();

    Ok(Json(serde_json::json!({
        "total_models": stats.model_count,
        "total_size": stats.total_size_human(),
        "total_size_bytes": stats.total_size,
        "cache_size": stats.cache_size_human(),
        "cache_size_bytes": stats.cache_size,
        "data_directory": stats.data_dir.to_string_lossy(),
        "config_directory": stats.config_dir.to_string_lossy(),
        "format_breakdown": stats.format_breakdown.iter().map(|(k, v)| {
            (format!("{:?}", k), *v)
        }).collect::<std::collections::HashMap<_, _>>()
    })))
}

/// Cleanup storage endpoint
async fn cleanup_storage(
    State(state): State<AppState>,
    headers: HeaderMap,
    peer: Option<ConnectInfo<SocketAddr>>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<APIError>)> {
    use crate::ai_api::storage::{ModelStorage, StorageConfig};

    authorize_mutation(&headers, &state.config, peer_addr(peer))?;

    // Initialize storage
    let storage_config = StorageConfig::default();
    let storage = ModelStorage::new(&storage_config).map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(APIError {
                error: ErrorDetail {
                    message: format!("Failed to initialize storage: {}", e),
                    error_type: "storage_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
    })?;

    // Clean up cache files older than 30 days
    let cleanup_days = 30;
    match storage.cleanup_cache(cleanup_days) {
        Ok(freed_bytes) => {
            let freed_human = if freed_bytes < 1024 {
                format!("{} B", freed_bytes)
            } else if freed_bytes < 1024 * 1024 {
                format!("{:.2} KB", freed_bytes as f64 / 1024.0)
            } else if freed_bytes < 1024 * 1024 * 1024 {
                format!("{:.2} MB", freed_bytes as f64 / (1024.0 * 1024.0))
            } else {
                format!("{:.2} GB", freed_bytes as f64 / (1024.0 * 1024.0 * 1024.0))
            };

            Ok(Json(serde_json::json!({
                "success": true,
                "freed_bytes": freed_bytes,
                "freed_human": freed_human,
                "cleanup_age_days": cleanup_days,
                "timestamp": chrono::Utc::now().to_rfc3339()
            })))
        }
        Err(e) => Ok(Json(serde_json::json!({
            "success": false,
            "error": format!("Cleanup failed: {}", e),
            "freed_bytes": 0
        }))),
    }
}

/// Health check endpoint
async fn health_check() -> Json<serde_json::Value> {
    Json(serde_json::json!({
        "status": "healthy",
        "timestamp": chrono::Utc::now().to_rfc3339()
    }))
}

/// Server status endpoint
async fn server_status(State(state): State<AppState>) -> Json<serde_json::Value> {
    let uptime = state.start_time.elapsed();
    let uptime_str = format!(
        "{}d {}h {}m {}s",
        uptime.as_secs() / 86400,
        (uptime.as_secs() % 86400) / 3600,
        (uptime.as_secs() % 3600) / 60,
        uptime.as_secs() % 60
    );

    Json(serde_json::json!({
        "status": "running",
        "version": env!("CARGO_PKG_VERSION"),
        "config": {
            "cors_enabled": state.config.server.enable_cors,
            "openapi_enabled": state.config.server.enable_openapi,
        },
        "uptime": uptime_str,
        "uptime_seconds": uptime.as_secs(),
        "timestamp": chrono::Utc::now().to_rfc3339()
    }))
}

// SECURITY FIX (LOW-001): Middleware to add security headers
async fn add_security_headers(
    req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> axum::response::Response {
    let mut response = next.run(req).await;
    let headers = response.headers_mut();

    // Prevent MIME-type sniffing
    headers.insert(
        axum::http::header::HeaderName::from_static("x-content-type-options"),
        axum::http::HeaderValue::from_static("nosniff"),
    );

    // Prevent clickjacking
    headers.insert(
        axum::http::header::HeaderName::from_static("x-frame-options"),
        axum::http::HeaderValue::from_static("DENY"),
    );

    // XSS protection (legacy browsers)
    headers.insert(
        axum::http::header::HeaderName::from_static("x-xss-protection"),
        axum::http::HeaderValue::from_static("1; mode=block"),
    );

    // HSTS - enforce HTTPS (when using TLS)
    headers.insert(
        axum::http::header::HeaderName::from_static("strict-transport-security"),
        axum::http::HeaderValue::from_static("max-age=31536000; includeSubDomains"),
    );

    // Referrer policy - limit referrer information
    headers.insert(
        axum::http::header::HeaderName::from_static("referrer-policy"),
        axum::http::HeaderValue::from_static("strict-origin-when-cross-origin"),
    );

    // Permissions policy - restrict browser features
    headers.insert(
        axum::http::header::HeaderName::from_static("permissions-policy"),
        axum::http::HeaderValue::from_static("geolocation=(), microphone=(), camera=()"),
    );

    response
}

// ==================== Streaming Implementation ====================

/// Streaming chat completion chunk (OpenAI-compatible format)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionChunk {
    pub id: String,
    pub object: String,
    pub created: i64,
    pub model: String,
    pub choices: Vec<StreamChoice>,
    pub system_fingerprint: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamChoice {
    pub index: u32,
    pub delta: StreamDelta,
    pub finish_reason: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamDelta {
    pub role: Option<String>,
    pub content: Option<String>,
}

/// Create a streaming chat completion response
async fn stream_chat_completion(
    api: Arc<AIModelAPI>,
    request: ChatCompletionRequest,
) -> Result<impl Stream<Item = Result<axum::response::sse::Event, Infallible>>> {
    // Get the non-streaming response first
    let response = api.chat_completion(request.clone()).await?;

    // Extract the full content
    let full_content = response
        .choices
        .first()
        .and_then(|c| c.message.content.as_ref())
        .cloned()
        .unwrap_or_default();

    let model = response.model.clone();
    let id = response.id.clone();
    let created = response.created;

    // Simulate streaming by chunking the response
    // In production, this would connect to actual streaming endpoints
    let chunk_size = 4; // Characters per chunk
    let chunks: Vec<String> = full_content
        .chars()
        .collect::<Vec<_>>()
        .chunks(chunk_size)
        .map(|c| c.iter().collect())
        .collect();

    let stream = async_stream::stream! {
        // First chunk with role
        let first_chunk = ChatCompletionChunk {
            id: id.clone(),
            object: "chat.completion.chunk".to_string(),
            created,
            model: model.clone(),
            choices: vec![StreamChoice {
                index: 0,
                delta: StreamDelta {
                    role: Some("assistant".to_string()),
                    content: None,
                },
                finish_reason: None,
            }],
            system_fingerprint: None,
        };

        let data = serde_json::to_string(&first_chunk).unwrap_or_default();
        yield Ok(axum::response::sse::Event::default().data(data));

        // Content chunks
        for (i, chunk_content) in chunks.iter().enumerate() {
            // Small delay to simulate real streaming
            tokio::time::sleep(Duration::from_millis(20)).await;

            let is_last = i == chunks.len() - 1;

            let chunk = ChatCompletionChunk {
                id: id.clone(),
                object: "chat.completion.chunk".to_string(),
                created,
                model: model.clone(),
                choices: vec![StreamChoice {
                    index: 0,
                    delta: StreamDelta {
                        role: None,
                        content: Some(chunk_content.clone()),
                    },
                    finish_reason: if is_last { Some("stop".to_string()) } else { None },
                }],
                system_fingerprint: None,
            };

            let data = serde_json::to_string(&chunk).unwrap_or_default();
            yield Ok(axum::response::sse::Event::default().data(data));
        }

        // Final [DONE] marker (OpenAI convention)
        yield Ok(axum::response::sse::Event::default().data("[DONE]"));
    };

    Ok(stream)
}

// Helper functions

/// Validate API key from request headers
fn validate_api_key(
    headers: &HeaderMap,
    config: &APIConfig,
) -> Result<(), (StatusCode, Json<APIError>)> {
    let auth_header = headers
        .get("Authorization")
        .and_then(|h| h.to_str().ok())
        .and_then(|h| h.strip_prefix("Bearer "));

    if let Some(provided_key) = auth_header {
        if key_accepted(provided_key, &config.security.api_keys) {
            Ok(())
        } else {
            Err((
                StatusCode::UNAUTHORIZED,
                Json(APIError {
                    error: ErrorDetail {
                        message: "Invalid API key".to_string(),
                        error_type: "authentication_error".to_string(),
                        param: None,
                        code: None,
                    },
                }),
            ))
        }
    } else {
        Err((
            StatusCode::UNAUTHORIZED,
            Json(APIError {
                error: ErrorDetail {
                    message: "Missing API key".to_string(),
                    error_type: "authentication_error".to_string(),
                    param: None,
                    code: None,
                },
            }),
        ))
    }
}

/// Whether `presented` equals one of `keys`. Each key is compared in full, with
/// no early exit, so the time taken does not show which key, or how much of a
/// key, matched. The byte comparison is the same as `mcp.rs`
/// (`mcp_tokens_match`) and `agent_api.rs` (`tokens_match`).
fn key_accepted(presented: &str, keys: &[String]) -> bool {
    let mut accepted = false;
    for key in keys {
        accepted |= token_eq(presented.as_bytes(), key.as_bytes());
    }
    accepted
}

/// Constant-time byte comparison. A length mismatch is decided first, so the
/// length of a key can be learned, not its contents.
fn token_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff = 0u8;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}

/// The peer address a handler was served from, if the server attached it.
/// Handlers called directly, without a socket, get `None`.
fn peer_addr(peer: Option<ConnectInfo<SocketAddr>>) -> Option<SocketAddr> {
    peer.map(|ConnectInfo(addr)| addr)
}

/// Whether the peer is on this machine. IPv4-mapped IPv6 loopback counts.
fn peer_is_loopback(peer: Option<SocketAddr>) -> bool {
    match peer.map(|a| a.ip()) {
        Some(IpAddr::V4(v4)) => v4.is_loopback(),
        Some(IpAddr::V6(v6)) => {
            v6.is_loopback() || v6.to_ipv4_mapped().is_some_and(|v4| v4.is_loopback())
        }
        None => false,
    }
}

/// Map a model-id check failure to a `400` with the reason.
fn validate_id_or_bad_request(model_id: &str) -> Result<(), (StatusCode, Json<APIError>)> {
    crate::ai_api::storage::validate_model_id(model_id).map_err(|e| {
        (
            StatusCode::BAD_REQUEST,
            Json(APIError {
                error: ErrorDetail {
                    message: e.to_string(),
                    error_type: "invalid_request_error".to_string(),
                    param: Some("model_id".to_string()),
                    code: None,
                },
            }),
        )
    })
}

/// Authorize a request that reads or changes local state: model download,
/// convert, delete, storage cleanup and stats, and provider validation.
/// Unlike the inference routes this check does not depend on
/// `require_api_key`, because a default server leaves that flag off.
///
/// With API keys configured, the request needs a matching bearer token.
///
/// With no keys, the server is bound to loopback (`check_startup` refuses a
/// non-loopback bind without keys), and the request must carry an
/// `X-IronShell-Client` header and come from a loopback peer. The header is
/// self-asserted, so it is not a credential. It is a browser-CSRF measure: a
/// page on another origin cannot send a custom header without a CORS
/// preflight, and this server never approves one for that header. The
/// loopback-peer check is what closes the rest: a request that reached a
/// loopback socket from elsewhere (a forwarded port, a reverse proxy on this
/// machine) is refused.
fn authorize_mutation(
    headers: &HeaderMap,
    config: &APIConfig,
    peer: Option<SocketAddr>,
) -> Result<(), (StatusCode, Json<APIError>)> {
    if config.security.require_api_key || !config.security.api_keys.is_empty() {
        return validate_api_key(headers, config);
    }

    let has_client_header = headers
        .get("X-IronShell-Client")
        .and_then(|h| h.to_str().ok())
        .is_some_and(|v| !v.is_empty());

    if !peer_is_loopback(peer) {
        return Err((
            StatusCode::FORBIDDEN,
            Json(APIError {
                error: ErrorDetail {
                    message: "Without an API key, this route accepts only loopback peers"
                        .to_string(),
                    error_type: "forbidden".to_string(),
                    param: None,
                    code: None,
                },
            }),
        ));
    }

    if has_client_header {
        Ok(())
    } else {
        Err((
            StatusCode::FORBIDDEN,
            Json(APIError {
                error: ErrorDetail {
                    message: "Missing X-IronShell-Client header".to_string(),
                    error_type: "forbidden".to_string(),
                    param: None,
                    code: None,
                },
            }),
        ))
    }
}

/// The `Host` values this server answers to. A browser that has been
/// rebound to this machine by DNS sends the attacker's name here, so any
/// name outside this set is refused.
pub struct HostPolicy {
    extra: Vec<String>,
}

impl HostPolicy {
    fn new(server: &ServerConfig) -> Self {
        let mut extra: Vec<String> = server
            .allowed_hosts
            .iter()
            .map(|h| normalize_host(h))
            .collect();
        // A specific bind address is named by its clients too. A wildcard is
        // not a name any client sends, so it adds nothing.
        let bind = normalize_host(&server.host);
        if !bind.is_empty() && bind != "0.0.0.0" && bind != "::" {
            extra.push(bind);
        }
        Self { extra }
    }

    fn allows(&self, host: &str) -> bool {
        host == "localhost"
            || host.parse::<IpAddr>().is_ok_and(|ip| ip.is_loopback())
            || self.extra.iter().any(|h| h == host)
    }
}

fn normalize_host(value: &str) -> String {
    value
        .trim()
        .trim_start_matches('[')
        .trim_end_matches(']')
        .to_ascii_lowercase()
}

/// The host name from a `Host` header value, `host[:port]`, with any port
/// removed and brackets stripped from an IPv6 literal. `None` when the value
/// is not that shape, which includes a bare IPv6 address.
fn host_name(value: &str) -> Option<String> {
    let v = value.trim();
    let (host, port) = if let Some(rest) = v.strip_prefix('[') {
        let end = rest.find(']')?;
        let after = &rest[end + 1..];
        let port = match after {
            "" => None,
            a => Some(a.strip_prefix(':')?),
        };
        (&rest[..end], port)
    } else {
        match v.rsplit_once(':') {
            Some((h, p)) if !h.contains(':') => (h, Some(p)),
            Some(_) => return None,
            None => (v, None),
        }
    };
    if let Some(p) = port {
        if p.is_empty() || !p.bytes().all(|b| b.is_ascii_digit()) {
            return None;
        }
    }
    if host.is_empty() {
        return None;
    }
    Some(normalize_host(host))
}

/// Refuses a request whose `Host` header is not an allowed name. Missing or
/// malformed `Host` is refused too. `/health` is not exempt: it is reachable
/// on the loopback names, which is all a local monitor needs.
async fn host_guard_middleware(
    State(policy): State<Arc<HostPolicy>>,
    req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Response {
    let host = req
        .headers()
        .get(axum::http::header::HOST)
        .and_then(|v| v.to_str().ok())
        .and_then(host_name);
    match host {
        Some(h) if policy.allows(&h) => next.run(req).await,
        _ => (
            StatusCode::FORBIDDEN,
            Json(APIError {
                error: ErrorDetail {
                    message: "Host header is not a name this server answers to".to_string(),
                    error_type: "forbidden".to_string(),
                    param: None,
                    code: None,
                },
            }),
        )
            .into_response(),
    }
}

const RATE_WINDOW_MINUTE: Duration = Duration::from_secs(60);
const RATE_WINDOW_HOUR: Duration = Duration::from_secs(3600);
/// Past this many tracked clients, windows that have expired are dropped.
const RATE_TRACKED_MAX: usize = 10_000;

#[derive(Clone, Copy)]
struct RateWindow {
    minute_start: Instant,
    minute_count: u32,
    hour_start: Instant,
    hour_count: u32,
}

/// A fixed-window request counter per client IP. It is in-process and
/// holds no more than one entry per recent client.
pub struct RateLimiter {
    per_minute: u32,
    per_hour: u32,
    windows: Mutex<HashMap<Option<IpAddr>, RateWindow>>,
}

impl RateLimiter {
    pub fn new(config: &RateLimitConfig) -> Self {
        Self {
            per_minute: config.requests_per_minute,
            per_hour: config.requests_per_hour,
            windows: Mutex::new(HashMap::new()),
        }
    }

    /// Count a request from `key` (`None` when no peer address is known). A
    /// request over either limit is refused and not counted. A limit of `0`
    /// is no limit.
    pub fn admit(&self, key: Option<IpAddr>, now: Instant) -> bool {
        let mut windows = match self.windows.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        };
        if windows.len() > RATE_TRACKED_MAX {
            windows.retain(|_, w| now.duration_since(w.hour_start) < RATE_WINDOW_HOUR);
        }
        let w = windows.entry(key).or_insert(RateWindow {
            minute_start: now,
            minute_count: 0,
            hour_start: now,
            hour_count: 0,
        });
        if now.duration_since(w.minute_start) >= RATE_WINDOW_MINUTE {
            w.minute_start = now;
            w.minute_count = 0;
        }
        if now.duration_since(w.hour_start) >= RATE_WINDOW_HOUR {
            w.hour_start = now;
            w.hour_count = 0;
        }
        let over_minute = self.per_minute > 0 && w.minute_count >= self.per_minute;
        let over_hour = self.per_hour > 0 && w.hour_count >= self.per_hour;
        if over_minute || over_hour {
            return false;
        }
        w.minute_count += 1;
        w.hour_count += 1;
        true
    }
}

async fn rate_limit_middleware(
    State(limiter): State<Arc<RateLimiter>>,
    req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Response {
    let key = req
        .extensions()
        .get::<ConnectInfo<SocketAddr>>()
        .map(|c| c.0.ip());
    if limiter.admit(key, Instant::now()) {
        return next.run(req).await;
    }
    let mut response = (
        StatusCode::TOO_MANY_REQUESTS,
        Json(APIError {
            error: ErrorDetail {
                message: "Rate limit exceeded; retry later".to_string(),
                error_type: "rate_limit_exceeded".to_string(),
                param: None,
                code: None,
            },
        }),
    )
        .into_response();
    response.headers_mut().insert(
        axum::http::header::RETRY_AFTER,
        axum::http::HeaderValue::from_static("60"),
    );
    response
}

// Additional types for API responses

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ProviderInfo {
    pub id: String,
    pub name: String,
    pub description: String,
    pub enabled: bool,
    pub status: String,
}

// OpenAPI documentation
// Temporarily disabled OpenAPI documentation for compilation
// #[derive(OpenApi)]
pub struct ApiDoc;

/// Start the API server with the given configuration
pub async fn start_server(config: APIConfig) -> Result<()> {
    let server = APIServer::new(config)?;
    server.start().await
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_security_headers_list() {
        // Verify that we're setting all required security headers
        // This is a documentation test to ensure we don't forget any headers

        let expected_headers = [
            ("x-content-type-options", "nosniff"),
            ("x-frame-options", "DENY"),
            ("x-xss-protection", "1; mode=block"),
            (
                "strict-transport-security",
                "max-age=31536000; includeSubDomains",
            ),
            ("referrer-policy", "strict-origin-when-cross-origin"),
            (
                "permissions-policy",
                "geolocation=(), microphone=(), camera=()",
            ),
        ];

        // Verify we have 6 security headers configured
        assert_eq!(
            expected_headers.len(),
            6,
            "Should have 6 security headers configured"
        );

        // Test passes - actual middleware testing would require integration tests
        // with a running server, which is beyond the scope of unit tests
    }

    use super::{authorize_mutation, host_name, key_accepted, HostPolicy, RateLimiter};
    use crate::ai_api::config::{APIConfig, RateLimitConfig};
    use axum::http::{HeaderMap, HeaderValue, StatusCode};
    use std::net::{IpAddr, SocketAddr};
    use std::time::{Duration, Instant};

    fn headers_with(name: &'static str, value: &'static str) -> HeaderMap {
        let mut headers = HeaderMap::new();
        headers.insert(name, HeaderValue::from_static(value));
        headers
    }

    fn loopback_peer() -> Option<SocketAddr> {
        Some("127.0.0.1:40000".parse().unwrap())
    }

    fn remote_peer() -> Option<SocketAddr> {
        Some("192.0.2.7:40000".parse().unwrap())
    }

    #[test]
    fn authorize_mutation_no_key_no_header_is_forbidden() {
        let config = APIConfig::default();

        let err = authorize_mutation(&HeaderMap::new(), &config, loopback_peer())
            .expect_err("mutation without key or client header must be refused");
        assert_eq!(err.0, StatusCode::FORBIDDEN);

        // An empty header value does not count as the client marker.
        let empty = headers_with("X-IronShell-Client", "");
        assert!(authorize_mutation(&empty, &config, loopback_peer()).is_err());
    }

    #[test]
    fn authorize_mutation_no_key_with_header_is_allowed() {
        let config = APIConfig::default();
        let headers = headers_with("X-IronShell-Client", "1");

        assert!(authorize_mutation(&headers, &config, loopback_peer()).is_ok());
    }

    #[test]
    fn authorize_mutation_no_key_needs_a_loopback_peer() {
        // The header alone is self-asserted; a remote peer, or no peer at all
        // (a handler called without a socket), is refused even with it.
        let config = APIConfig::default();
        let headers = headers_with("X-IronShell-Client", "1");

        assert!(authorize_mutation(&headers, &config, remote_peer()).is_err());
        assert!(authorize_mutation(&headers, &config, None).is_err());
        let mapped: SocketAddr = "[::ffff:127.0.0.1]:40000".parse().unwrap();
        assert!(authorize_mutation(&headers, &config, Some(mapped)).is_ok());
    }

    #[test]
    fn authorize_mutation_key_required_checks_bearer() {
        let mut config = APIConfig::default();
        config.security.require_api_key = true;
        config.security.api_keys = vec!["secret".to_string()];

        // Missing bearer, and a wrong bearer, are both refused even with the header.
        assert!(authorize_mutation(&HeaderMap::new(), &config, loopback_peer()).is_err());

        let mut wrong = headers_with("X-IronShell-Client", "1");
        wrong.insert("Authorization", HeaderValue::from_static("Bearer nope"));
        assert!(authorize_mutation(&wrong, &config, loopback_peer()).is_err());

        let mut right = HeaderMap::new();
        right.insert("Authorization", HeaderValue::from_static("Bearer secret"));
        assert!(authorize_mutation(&right, &config, remote_peer()).is_ok());
    }

    #[test]
    fn key_match_checks_every_configured_key() {
        let keys = vec!["alpha-key".to_string(), "beta-key".to_string()];
        assert!(key_accepted("alpha-key", &keys));
        assert!(key_accepted("beta-key", &keys));
        assert!(!key_accepted("alpha-kez", &keys));
        assert!(!key_accepted("alpha", &keys));
        assert!(!key_accepted("", &keys));
        assert!(!key_accepted("alpha-key", &[]));
    }

    #[test]
    fn host_header_names_are_parsed() {
        assert_eq!(host_name("localhost:8080").as_deref(), Some("localhost"));
        assert_eq!(host_name("127.0.0.1").as_deref(), Some("127.0.0.1"));
        assert_eq!(host_name("[::1]:8080").as_deref(), Some("::1"));
        assert_eq!(host_name("[::1]").as_deref(), Some("::1"));
        assert_eq!(
            host_name("Evil.Example:443").as_deref(),
            Some("evil.example")
        );
        // A bare IPv6 address is not a Host value, and neither is a bad port.
        assert_eq!(host_name("::1"), None);
        assert_eq!(host_name("localhost:"), None);
        assert_eq!(host_name("localhost:80x"), None);
        assert_eq!(host_name(":80"), None);
    }

    #[test]
    fn host_policy_allows_loopback_and_the_bind_host_only() {
        let mut server = crate::ai_api::config::ServerConfig::default();
        server.host = "192.168.1.10".to_string();
        let policy = HostPolicy::new(&server);
        assert!(policy.allows("localhost"));
        assert!(policy.allows("127.0.0.1"));
        assert!(policy.allows("::1"));
        assert!(policy.allows("192.168.1.10"));
        assert!(!policy.allows("evil.example"));
        assert!(!policy.allows("10.0.0.5"));

        // A wildcard bind names no client, so it adds no host.
        server.host = "0.0.0.0".to_string();
        let policy = HostPolicy::new(&server);
        assert!(!policy.allows("0.0.0.0"));
        assert!(!policy.allows("evil.example"));
    }

    #[test]
    fn rate_limit_refuses_past_the_minute_window_and_resets() {
        let limits = RateLimitConfig {
            requests_per_minute: 2,
            requests_per_hour: 100,
            ..RateLimitConfig::default()
        };
        let limiter = RateLimiter::new(&limits);
        let ip: IpAddr = "203.0.113.9".parse().unwrap();
        let other: IpAddr = "203.0.113.10".parse().unwrap();
        let t0 = Instant::now();

        assert!(limiter.admit(Some(ip), t0));
        assert!(limiter.admit(Some(ip), t0));
        assert!(!limiter.admit(Some(ip), t0), "third request in the window");
        // Another client has its own count.
        assert!(limiter.admit(Some(other), t0));
        // The window ends after a minute.
        assert!(limiter.admit(Some(ip), t0 + Duration::from_secs(61)));
    }
}