use std::{
    pin::Pin,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    task::{Context, Poll},
    time::{Duration, Instant},
};

use axum::{
    body::Body,
    extract::{Request, State},
    http::{header, HeaderValue, StatusCode},
    middleware::Next,
    response::{IntoResponse, Response},
    Json,
};
use bytes::Bytes;
use http_body::Frame;
use rand::Rng;
use serde_json::json;
use subtle::ConstantTimeEq;
use tokio::sync::{mpsc, oneshot};
use tower::{Layer, Service};
use tower_http::trace::{MakeSpan, OnRequest, OnResponse, TraceLayer};
use tracing::{debug, error, field::Empty, info, info_span, warn, Span};

pub use crate::core::token_bucket::TokenBucket;
use crate::{
    observability::{
        inflight_tracker::InFlightRequestTracker,
        metrics::{method_to_static_str, metrics_labels, Metrics},
    },
    routers::error::extract_error_code_from_response,
    server::AppState,
    wasm::{
        module::{MiddlewareAttachPoint, WasmModuleAttachPoint},
        spec::{
            apply_modify_action_to_headers, build_wasm_headers_from_axum_headers,
            smg::gateway::middleware_types::{
                Action, Request as WasmRequest, Response as WasmResponse,
            },
        },
        types::WasmComponentInput,
    },
};

/// A body wrapper that holds a token and returns it when the body is fully consumed or dropped.
/// This ensures that for streaming responses, the token is only returned after the entire
/// stream has been sent to the client.
pub struct TokenGuardBody {
    inner: Body,
    /// The token bucket to return tokens to. Uses Option so we can take() on drop.
    token_bucket: Option<Arc<TokenBucket>>,
    /// Number of tokens to return.
    tokens: f64,
}

impl TokenGuardBody {
    /// Create a new TokenGuardBody that will return tokens when dropped.
    pub fn new(inner: Body, token_bucket: Arc<TokenBucket>, tokens: f64) -> Self {
        Self {
            inner,
            token_bucket: Some(token_bucket),
            tokens,
        }
    }
}

impl Drop for TokenGuardBody {
    fn drop(&mut self) {
        if let Some(bucket) = self.token_bucket.take() {
            debug!(
                "TokenGuardBody: stream ended, returning {} tokens to bucket",
                self.tokens
            );
            // Use lock-free sync return - no runtime needed, guaranteed token return
            bucket.return_tokens_sync(self.tokens);
        }
    }
}

impl http_body::Body for TokenGuardBody {
    type Data = Bytes;
    type Error = axum::Error;

    fn poll_frame(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        // SAFETY: We never move the inner body, and Body is Unpin
        // (it's a type alias for UnsyncBoxBody which is Unpin)
        let this = self.get_mut();
        Pin::new(&mut this.inner).poll_frame(cx)
    }

    fn is_end_stream(&self) -> bool {
        self.inner.is_end_stream()
    }

    fn size_hint(&self) -> http_body::SizeHint {
        self.inner.size_hint()
    }
}

#[derive(Clone)]
pub struct AuthConfig {
    pub api_key: Option<String>,
}

/// Middleware to validate Bearer token against configured API key
/// Only active when router has an API key configured
pub async fn auth_middleware(
    State(auth_config): State<AuthConfig>,
    request: Request<Body>,
    next: Next,
) -> Result<Response, StatusCode> {
    if let Some(expected_key) = &auth_config.api_key {
        // Extract Authorization header
        let auth_header = request
            .headers()
            .get(header::AUTHORIZATION)
            .and_then(|h| h.to_str().ok());

        match auth_header {
            Some(header_value) if header_value.starts_with("Bearer ") => {
                let token = &header_value[7..]; // Skip "Bearer "
                                                // Use constant-time comparison to prevent timing attacks
                let token_bytes = token.as_bytes();
                let expected_bytes = expected_key.as_bytes();

                // Check if lengths match first (this is not constant-time but necessary)
                if token_bytes.len() != expected_bytes.len() {
                    return Err(StatusCode::UNAUTHORIZED);
                }

                // Constant-time comparison of the actual values
                if token_bytes.ct_eq(expected_bytes).unwrap_u8() != 1 {
                    return Err(StatusCode::UNAUTHORIZED);
                }
            }
            _ => return Err(StatusCode::UNAUTHORIZED),
        }
    }

    Ok(next.run(request).await)
}

/// Alphanumeric characters for request ID generation (as bytes for O(1) indexing)
const REQUEST_ID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";

/// Generate OpenAI-compatible request ID based on endpoint.
fn generate_request_id(path: &str) -> String {
    let prefix = if path.contains("/chat/completions") {
        "chatcmpl-"
    } else if path.contains("/completions") {
        "cmpl-"
    } else if path.contains("/generate") {
        "gnt-"
    } else if path.contains("/responses") {
        "resp-"
    } else {
        "req-"
    };

    // Generate a random string similar to OpenAI's format
    // Use byte array indexing (O(1)) instead of chars().nth() (O(n))
    let mut rng = rand::rng();
    let random_part: String = (0..24)
        .map(|_| {
            let idx = rng.random_range(0..REQUEST_ID_CHARS.len());
            REQUEST_ID_CHARS[idx] as char
        })
        .collect();

    format!("{}{}", prefix, random_part)
}

/// Extension type for storing request ID
#[derive(Clone, Debug)]
pub struct RequestId(pub String);

/// Tower Layer for request ID middleware
#[derive(Clone)]
pub struct RequestIdLayer {
    headers: Arc<Vec<String>>,
}

impl RequestIdLayer {
    pub fn new(headers: Vec<String>) -> Self {
        Self {
            headers: Arc::new(headers),
        }
    }
}

impl<S> Layer<S> for RequestIdLayer {
    type Service = RequestIdMiddleware<S>;

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

/// Tower Service for request ID middleware
#[derive(Clone)]
pub struct RequestIdMiddleware<S> {
    inner: S,
    headers: Arc<Vec<String>>,
}

impl<S> Service<Request> for RequestIdMiddleware<S>
where
    S: Service<Request, Response = Response> + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future =
        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, mut req: Request) -> Self::Future {
        let headers = self.headers.clone();

        // Extract request ID from headers or generate new one
        let mut request_id = None;

        for header_name in headers.iter() {
            if let Some(header_value) = req.headers().get(header_name) {
                if let Ok(value) = header_value.to_str() {
                    request_id = Some(value.to_string());
                    break;
                }
            }
        }

        let request_id = request_id.unwrap_or_else(|| generate_request_id(req.uri().path()));

        // Insert request ID into request extensions for other middleware/handlers to use
        req.extensions_mut().insert(RequestId(request_id.clone()));

        // Call the inner service
        let future = self.inner.call(req);

        Box::pin(async move {
            let mut response = future.await?;

            // Add request ID to response headers
            response.headers_mut().insert(
                "x-request-id",
                HeaderValue::from_str(&request_id)
                    .unwrap_or_else(|_| HeaderValue::from_static("invalid-request-id")),
            );

            Ok(response)
        })
    }
}

/// Custom span maker that includes request ID
#[derive(Clone, Debug)]
pub struct RequestSpan;

impl<B> MakeSpan<B> for RequestSpan {
    fn make_span(&mut self, request: &Request<B>) -> Span {
        // Don't try to extract request ID here - it won't be available yet
        // The RequestIdLayer runs after TraceLayer creates the span
        info_span!(
            target: "smg::otel-trace",
            "http_request",
            method = %request.method(),
            uri = %request.uri(),
            version = ?request.version(),
            request_id = Empty,  // Will be set later
            status_code = Empty,
            latency = Empty,
            error = Empty,
            module = "smg"
        )
    }
}

/// Custom on_request handler
#[derive(Clone, Debug)]
pub struct RequestLogger;

impl<B> OnRequest<B> for RequestLogger {
    fn on_request(&mut self, request: &Request<B>, span: &Span) {
        let _enter = span.enter();

        // Try to get the request ID from extensions
        // This will work if RequestIdLayer has already run
        if let Some(request_id) = request.extensions().get::<RequestId>() {
            span.record("request_id", request_id.0.as_str());
        }

        let method = method_to_static_str(request.method().as_str());
        let path = normalize_path_for_metrics(request.uri().path());
        Metrics::record_http_request(method, &path);

        // Log the request start
        info!(
            target: "smg::request",
            "started processing request"
        );
    }
}

/// Custom on_response handler
#[derive(Clone, Debug)]
pub struct ResponseLogger {
    _start_time: Instant,
}

impl Default for ResponseLogger {
    fn default() -> Self {
        Self {
            _start_time: Instant::now(),
        }
    }
}

impl<B> OnResponse<B> for ResponseLogger {
    fn on_response(self, response: &Response<B>, latency: Duration, span: &Span) {
        let status = response.status();
        let status_code = status.as_u16();

        let error_code = extract_error_code_from_response(response);

        // Layer 1: HTTP metrics
        Metrics::record_http_response(status_code, error_code);

        // Record these in the span for structured logging/observability tools
        span.record("status_code", status_code);
        // Use microseconds as integer to avoid format! string allocation
        span.record("latency", latency.as_micros() as u64);

        // Log the response completion
        let _enter = span.enter();
        if status.is_server_error() {
            error!(
                target: "smg::response",
                "request failed with server error"
            );
        } else if status.is_client_error() {
            warn!(
                target: "smg::response",
                "request failed with client error"
            );
        } else {
            info!(
                target: "smg::response",
                "finished processing request"
            );
        }
    }
}

/// Create a configured TraceLayer for HTTP logging
/// Note: Actual request/response logging with request IDs is done in RequestIdService
pub fn create_logging_layer() -> TraceLayer<
    tower_http::classify::SharedClassifier<tower_http::classify::ServerErrorsAsFailures>,
    RequestSpan,
    RequestLogger,
    ResponseLogger,
> {
    TraceLayer::new_for_http()
        .make_span_with(RequestSpan)
        .on_request(RequestLogger)
        .on_response(ResponseLogger::default())
}

/// Request queue entry
pub struct QueuedRequest {
    /// Time when the request was queued
    queued_at: Instant,
    /// Channel to send the permit back when acquired
    permit_tx: oneshot::Sender<Result<(), StatusCode>>,
}

/// Queue metrics for monitoring
#[derive(Debug, Default)]
pub struct QueueMetrics {
    pub total_queued: AtomicU64,
    pub current_queued: AtomicU64,
    pub total_timeout: AtomicU64,
    pub total_rejected: AtomicU64,
}

/// Queue processor that handles queued requests
pub struct QueueProcessor {
    token_bucket: Arc<TokenBucket>,
    queue_rx: mpsc::Receiver<QueuedRequest>,
    queue_timeout: Duration,
}

impl QueueProcessor {
    pub fn new(
        token_bucket: Arc<TokenBucket>,
        queue_rx: mpsc::Receiver<QueuedRequest>,
        queue_timeout: Duration,
    ) -> Self {
        Self {
            token_bucket,
            queue_rx,
            queue_timeout,
        }
    }

    pub async fn run(mut self) {
        debug!("Starting concurrency queue processor");

        // Process requests in a single task to reduce overhead
        while let Some(queued) = self.queue_rx.recv().await {
            // Check timeout immediately
            let elapsed = queued.queued_at.elapsed();
            if elapsed >= self.queue_timeout {
                warn!("Request already timed out in queue");
                let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
                continue;
            }

            let remaining_timeout = self.queue_timeout - elapsed;

            // Try to acquire token for this request
            if self.token_bucket.try_acquire(1.0).await.is_ok() {
                // Got token immediately
                debug!("Queue: acquired token immediately for queued request");
                let _ = queued.permit_tx.send(Ok(()));
            } else {
                // Need to wait for token
                let token_bucket = self.token_bucket.clone();

                // Spawn task only when we actually need to wait
                tokio::spawn(async move {
                    if token_bucket
                        .acquire_timeout(1.0, remaining_timeout)
                        .await
                        .is_ok()
                    {
                        debug!("Queue: acquired token after waiting");
                        let _ = queued.permit_tx.send(Ok(()));
                    } else {
                        warn!("Queue: request timed out waiting for token");
                        let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
                    }
                });
            }
        }

        warn!("Concurrency queue processor shutting down");
    }
}

/// State for the concurrency limiter
pub struct ConcurrencyLimiter {
    pub queue_tx: Option<mpsc::Sender<QueuedRequest>>,
}

impl ConcurrencyLimiter {
    /// Create new concurrency limiter with optional queue
    pub fn new(
        token_bucket: Option<Arc<TokenBucket>>,
        queue_size: usize,
        queue_timeout: Duration,
    ) -> (Self, Option<QueueProcessor>) {
        match (token_bucket, queue_size) {
            (None, _) => (Self { queue_tx: None }, None),
            (Some(bucket), size) if size > 0 => {
                let (queue_tx, queue_rx) = mpsc::channel(size);
                let processor = QueueProcessor::new(bucket, queue_rx, queue_timeout);
                (
                    Self {
                        queue_tx: Some(queue_tx),
                    },
                    Some(processor),
                )
            }
            (Some(_), _) => (Self { queue_tx: None }, None),
        }
    }
}

/// Middleware function for concurrency limiting with optional queuing
pub async fn concurrency_limit_middleware(
    State(app_state): State<Arc<AppState>>,
    request: Request<Body>,
    next: Next,
) -> Response {
    // Check mesh global rate limit first if mesh is enabled
    // If mesh is not enabled, this check is skipped and local rate limiting is used
    if let Some(sync_manager) = &app_state.mesh_sync_manager {
        let (is_exceeded, current_count, limit) = sync_manager.check_global_rate_limit();
        if is_exceeded {
            debug!(
                "Global rate limit exceeded: {}/{} req/s",
                current_count, limit
            );
            return (
                StatusCode::TOO_MANY_REQUESTS,
                Json(json!({
                    "error": "Rate limit exceeded",
                    "current_count": current_count,
                    "limit": limit
                })),
            )
                .into_response();
        }
    }

    let token_bucket = match &app_state.context.rate_limiter {
        Some(bucket) => bucket.clone(),
        None => {
            // Rate limiting disabled, pass through immediately
            return next.run(request).await;
        }
    };

    // Static counter for embeddings queue size
    static EMBEDDINGS_QUEUE_SIZE: AtomicU64 = AtomicU64::new(0);

    // Identify if this is an embeddings request based on path
    let is_embeddings = request.uri().path().contains("/v1/embeddings");

    // Try to acquire token immediately
    if token_bucket.try_acquire(1.0).await.is_ok() {
        debug!("Acquired token immediately");
        Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED);
        let response = next.run(request).await;

        // Wrap the response body with TokenGuardBody to return token when stream ends
        // This ensures that for streaming responses, the token is only returned
        // after the entire stream has been sent to the client.
        let (parts, body) = response.into_parts();
        let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0);
        Response::from_parts(parts, Body::new(guarded_body))
    } else {
        // No tokens available, try to queue if enabled
        if let Some(queue_tx) = &app_state.concurrency_queue_tx {
            debug!("No tokens available, attempting to queue request");

            // Create a channel for the token response
            let (permit_tx, permit_rx) = oneshot::channel();

            let queued = QueuedRequest {
                queued_at: Instant::now(),
                permit_tx,
            };

            // Try to send to queue
            match queue_tx.try_send(queued) {
                Ok(_) => {
                    // On successful enqueue, update embeddings queue counter if applicable
                    if is_embeddings {
                        EMBEDDINGS_QUEUE_SIZE.fetch_add(1, Ordering::Relaxed);
                    }

                    // Wait for token from queue processor
                    match permit_rx.await {
                        Ok(Ok(())) => {
                            debug!("Acquired token from queue");
                            Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED);
                            // Dequeue for embeddings
                            if is_embeddings {
                                EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed);
                            }

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

                            // Wrap the response body with TokenGuardBody to return token when stream ends
                            let (parts, body) = response.into_parts();
                            let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0);
                            Response::from_parts(parts, Body::new(guarded_body))
                        }
                        Ok(Err(status)) => {
                            warn!("Queue returned error status: {}", status);
                            Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
                            // Dequeue for embeddings on error
                            if is_embeddings {
                                EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed);
                            }
                            status.into_response()
                        }
                        Err(_) => {
                            error!("Queue response channel closed");
                            Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
                            // Dequeue for embeddings on channel error
                            if is_embeddings {
                                EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed);
                            }
                            StatusCode::INTERNAL_SERVER_ERROR.into_response()
                        }
                    }
                }
                Err(_) => {
                    warn!("Request queue is full, returning 429");
                    Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
                    StatusCode::TOO_MANY_REQUESTS.into_response()
                }
            }
        } else {
            warn!("No tokens available and queuing is disabled, returning 429");
            Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
            StatusCode::TOO_MANY_REQUESTS.into_response()
        }
    }
}

// ============================================================================
// HTTP Metrics Layer (Layer 1: SMG metrics)
// ============================================================================

/// Global counter for active HTTP connections (handlers currently executing)
static ACTIVE_HTTP_CONNECTIONS: AtomicU64 = AtomicU64::new(0);

/// Tower Layer for HTTP metrics collection (SMG Layer 1 metrics)
#[derive(Clone)]
pub struct HttpMetricsLayer {
    tracker: Arc<InFlightRequestTracker>,
}

impl HttpMetricsLayer {
    pub fn new(tracker: Arc<InFlightRequestTracker>) -> Self {
        Self { tracker }
    }
}

impl<S> Layer<S> for HttpMetricsLayer {
    type Service = HttpMetricsMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        HttpMetricsMiddleware {
            inner,
            in_flight_request_tracker: self.tracker.clone(),
        }
    }
}

/// Tower Service for HTTP metrics collection
#[derive(Clone)]
pub struct HttpMetricsMiddleware<S> {
    inner: S,
    in_flight_request_tracker: Arc<InFlightRequestTracker>,
}

impl<S> Service<Request> for HttpMetricsMiddleware<S>
where
    S: Service<Request, Response = Response> + Send + Clone + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future =
        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, req: Request) -> Self::Future {
        // Convert method to static string to avoid allocation
        let method = method_to_static_str(req.method().as_str());
        let path = normalize_path_for_metrics(req.uri().path());
        let start = Instant::now();

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

        Box::pin(async move {
            // Increment inside async block - ensures no leak if future is dropped before polling
            let active = ACTIVE_HTTP_CONNECTIONS.fetch_add(1, Ordering::Relaxed) + 1;
            Metrics::set_http_connections_active(active as usize);

            let guard = in_flight_request_tracker.track();

            // Capture result before decrementing to ensure decrement happens on error too
            let result = inner.call(req).await;

            drop(guard);

            // Always decrement, regardless of success or failure
            let active = ACTIVE_HTTP_CONNECTIONS.fetch_sub(1, Ordering::Relaxed) - 1;
            Metrics::set_http_connections_active(active as usize);

            let response = result?;

            let duration = start.elapsed();
            Metrics::record_http_duration(method, &path, duration);

            Ok(response)
        })
    }
}

/// Normalize path for metrics to avoid high cardinality.
/// Replaces dynamic segments (IDs, UUIDs) with `{id}` placeholder.
/// Only allocates when normalization is needed; uses single-pass with byte offsets.
fn normalize_path_for_metrics(path: &str) -> String {
    let bytes = path.as_bytes();
    let mut segment_start = 0;
    let mut segment_idx = 0;
    let mut result: Option<String> = None;

    for (pos, &b) in bytes.iter().enumerate() {
        if b == b'/' || pos == bytes.len() - 1 {
            // Determine segment end (include last char if not a slash)
            let segment_end = if b == b'/' { pos } else { pos + 1 };
            let segment = &path[segment_start..segment_end];

            // Check segments after index 2 for dynamic IDs
            if segment_idx > 2 && !segment.is_empty() && is_dynamic_id(segment) {
                // Initialize result with everything before this segment
                let result = result.get_or_insert_with(|| {
                    let mut s = String::with_capacity(path.len());
                    s.push_str(&path[..segment_start]);
                    s
                });
                result.push_str("{id}");
            } else if let Some(ref mut r) = result {
                // Already normalizing, append this segment as-is
                r.push_str(segment);
            }

            // Add slash after segment (except at end)
            if b == b'/' {
                if let Some(ref mut r) = result {
                    r.push('/');
                }
                segment_start = pos + 1;
                segment_idx += 1;
            }
        }
    }

    result.unwrap_or_else(|| path.to_owned())
}

/// Check if segment looks like a dynamic ID (prefixed ID, UUID, or numeric).
#[inline]
fn is_dynamic_id(s: &str) -> bool {
    // Prefixed IDs: resp_xxx, chatcmpl_xxx (len > 10 with underscore)
    if s.len() > 10 && s.contains('_') {
        return true;
    }
    // UUIDs: 32+ hex chars with dashes
    if s.len() >= 32 && s.bytes().all(|b| b.is_ascii_hexdigit() || b == b'-') {
        return true;
    }
    // Numeric IDs
    !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit())
}

pub async fn wasm_middleware(
    State(app_state): State<Arc<AppState>>,
    request: Request<Body>,
    next: Next,
) -> Result<Response, StatusCode> {
    // Check if WASM is enabled
    if !app_state.context.router_config.enable_wasm {
        return Ok(next.run(request).await);
    }

    // Get WASM manager
    let wasm_manager = match &app_state.context.wasm_manager {
        Some(manager) => manager,
        None => {
            return Ok(next.run(request).await);
        }
    };

    // Get request ID from extensions or generate one
    let request_id = request
        .extensions()
        .get::<RequestId>()
        .map(|r| r.0.clone())
        .unwrap_or_else(|| generate_request_id(request.uri().path()));

    // ===== OnRequest Phase =====
    let on_request_attach_point =
        WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest);

    let modules_on_request =
        match wasm_manager.get_modules_by_attach_point(on_request_attach_point.clone()) {
            Ok(modules) => modules,
            Err(e) => {
                error!("Failed to get WASM modules for OnRequest: {}", e);
                return Ok(next.run(request).await);
            }
        };

    let response = if modules_on_request.is_empty() {
        next.run(request).await
    } else {
        // Extract request body once before processing modules
        let method = request.method().clone();
        let uri = request.uri().clone();
        let mut headers = request.headers().clone();
        let max_body_size = wasm_manager.get_max_body_size();
        let body_bytes = match axum::body::to_bytes(request.into_body(), max_body_size).await {
            Ok(bytes) => bytes.to_vec(),
            Err(e) => {
                error!("Failed to read request body: {}", e);
                // Create a minimal request with empty body for error recovery
                let error_request = Request::builder()
                    .uri(uri)
                    .body(Body::empty())
                    .unwrap_or_else(|_| Request::new(Body::empty()));
                return Ok(next.run(error_request).await);
            }
        };

        // Process each OnRequest module
        let mut modified_body = body_bytes;

        // Pre-compute strings once before the loop to avoid repeated allocations
        let method_str = method.to_string();
        let path_str = uri.path().to_string();
        let query_str = uri.query().unwrap_or("").to_string();

        for module in modules_on_request {
            // Build WebAssembly request from collected data
            let wasm_headers = build_wasm_headers_from_axum_headers(&headers);
            let wasm_request = WasmRequest {
                method: method_str.clone(),
                path: path_str.clone(),
                query: query_str.clone(),
                headers: wasm_headers,
                body: modified_body.clone(),
                request_id: request_id.clone(),
                now_epoch_ms: std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_else(|_| Duration::from_millis(0))
                    .as_millis() as u64,
            };

            // Execute WASM component
            let action = match wasm_manager
                .execute_module_for_attach_point(
                    &module,
                    on_request_attach_point.clone(),
                    WasmComponentInput::MiddlewareRequest(wasm_request),
                )
                .await
            {
                Some(action) => action,
                None => continue, // Continue to next module on error
            };

            // Process action
            match action {
                Action::Continue => {}
                Action::Reject(status) => {
                    return Err(StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST));
                }
                Action::Modify(modify) => {
                    // Apply modifications to headers and body
                    apply_modify_action_to_headers(&mut headers, &modify);
                    if let Some(body_bytes) = modify.body_replace {
                        modified_body = body_bytes;
                    }
                }
            }
        }

        // Reconstruct request with modifications
        let mut final_request = Request::builder()
            .method(method)
            .uri(uri)
            .body(Body::from(modified_body))
            .unwrap_or_else(|_| Request::new(Body::empty()));
        *final_request.headers_mut() = headers;

        // Continue with request processing
        next.run(final_request).await
    };

    // ===== OnResponse Phase =====
    let on_response_attach_point =
        WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse);

    let modules_on_response =
        match wasm_manager.get_modules_by_attach_point(on_response_attach_point.clone()) {
            Ok(modules) => modules,
            Err(e) => {
                error!("Failed to get WASM modules for OnResponse: {}", e);
                return Ok(response);
            }
        };
    if modules_on_response.is_empty() {
        return Ok(response);
    }
    // Extract response data once before processing modules
    let mut status = response.status();
    let mut headers = response.headers().clone();
    let max_body_size = wasm_manager.get_max_body_size();
    let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await {
        Ok(bytes) => bytes.to_vec(),
        Err(e) => {
            error!("Failed to read response body: {}", e);
            // Create a minimal response with empty body for error recovery
            let error_response = Response::builder()
                .status(status)
                .body(Body::empty())
                .unwrap_or_else(|_| Response::new(Body::empty()));
            return Ok(error_response);
        }
    };

    // Process each OnResponse module
    for module in modules_on_response {
        // Build WebAssembly response from collected data
        let wasm_headers = build_wasm_headers_from_axum_headers(&headers);
        let wasm_response = WasmResponse {
            status: status.as_u16(),
            headers: wasm_headers,
            body: body_bytes.clone(),
        };

        // Execute WASM component
        let action = match wasm_manager
            .execute_module_for_attach_point(
                &module,
                on_response_attach_point.clone(),
                WasmComponentInput::MiddlewareResponse(wasm_response),
            )
            .await
        {
            Some(action) => action,
            None => continue, // Continue to next module on error
        };

        // Process action - apply modifications incrementally
        match action {
            Action::Continue => {
                // Continue to next module
            }
            Action::Reject(status_code) => {
                // Override response status
                status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_REQUEST);
                // Return immediately with current state
                let final_response = Response::builder()
                    .status(status)
                    .body(Body::from(body_bytes))
                    .unwrap_or_else(|_| Response::new(Body::empty()));
                let mut final_response = final_response;
                *final_response.headers_mut() = headers;
                return Ok(final_response);
            }
            Action::Modify(modify) => {
                // Apply status modification
                if let Some(new_status) = modify.status {
                    status = StatusCode::from_u16(new_status).unwrap_or(status);
                }
                // Apply headers modifications
                apply_modify_action_to_headers(&mut headers, &modify);
                // Apply body_replace
                if let Some(new_body) = modify.body_replace {
                    body_bytes = new_body;
                }
            }
        }
    }

    // Reconstruct final response with all modifications
    let final_response = Response::builder()
        .status(status)
        .body(Body::from(body_bytes))
        .unwrap_or_else(|_| Response::new(Body::empty()));
    let mut final_response = final_response;
    *final_response.headers_mut() = headers;
    Ok(final_response)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_normalize_path_no_ids() {
        // Common API paths should pass through unchanged
        assert_eq!(
            normalize_path_for_metrics("/v1/chat/completions"),
            "/v1/chat/completions"
        );
        assert_eq!(
            normalize_path_for_metrics("/v1/completions"),
            "/v1/completions"
        );
        assert_eq!(normalize_path_for_metrics("/v1/models"), "/v1/models");
        assert_eq!(normalize_path_for_metrics("/health"), "/health");
    }

    #[test]
    fn test_normalize_path_with_prefixed_id() {
        // Prefixed IDs (resp_xxx, chatcmpl_xxx) should be normalized
        assert_eq!(
            normalize_path_for_metrics("/v1/responses/resp_abc123def456"),
            "/v1/responses/{id}"
        );
        assert_eq!(
            normalize_path_for_metrics("/v1/chat/completions/chatcmpl_abc123xyz"),
            "/v1/chat/completions/{id}"
        );
    }

    #[test]
    fn test_normalize_path_with_uuid() {
        assert_eq!(
            normalize_path_for_metrics("/v1/responses/550e8400-e29b-41d4-a716-446655440000"),
            "/v1/responses/{id}"
        );
    }

    #[test]
    fn test_normalize_path_with_numeric_id() {
        assert_eq!(
            normalize_path_for_metrics("/v1/workers/12345"),
            "/v1/workers/{id}"
        );
    }

    #[test]
    fn test_is_dynamic_id() {
        // Prefixed IDs
        assert!(is_dynamic_id("resp_abc123def"));
        assert!(is_dynamic_id("chatcmpl_xyz789"));
        assert!(!is_dynamic_id("short_id")); // Too short

        // UUIDs
        assert!(is_dynamic_id("550e8400-e29b-41d4-a716-446655440000"));
        assert!(is_dynamic_id("550e8400e29b41d4a716446655440000")); // No dashes

        // Numeric
        assert!(is_dynamic_id("12345"));
        assert!(!is_dynamic_id("")); // Empty

        // Regular words
        assert!(!is_dynamic_id("completions"));
        assert!(!is_dynamic_id("chat"));
    }
}
