use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
use crank_observability::{RequestId, set_remote_trace_parent, with_request_correlation};
use tracing::{Instrument, info_span};

use crate::transport::HEADER_X_REQUEST_ID;

#[derive(Clone, Debug)]
pub(super) struct RequestContext {
    pub(super) request_id: String,
}

pub(super) async fn apply_request_context(mut request: Request, next: Next) -> Response {
    let request_id = RequestId::resolve_from_headers(request.headers()).into_string();
    let context = RequestContext {
        request_id: request_id.clone(),
    };
    let span = info_span!(
        target: "crank::trace",
        "mcp.request",
        request_id = %request_id,
    );
    set_remote_trace_parent(&span, request.headers());
    request.extensions_mut().insert(context);

    with_request_correlation(request_id.clone(), async move {
        let mut response = next.run(request).instrument(span).await;
        if let Ok(value) = HeaderValue::from_str(&request_id) {
            response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
        }
        response
    })
    .await
}