use axum::{
    extract::Request,
    http::{HeaderValue, header::HeaderName},
    middleware::Next,
    response::Response,
};
use shared_contracts::api::{
    API_VERSION, API_VERSION_HEADER, RESPONSE_TIME_HEADER, ROUTE_VERSION_HEADER,
};

use crate::request_context::{RequestContext, X_REQUEST_ID_HEADER, resolve_request_id};

pub async fn propagate_request_id_header(request: Request, next: Next) -> Response {
    let request_id = resolve_request_id(&request);
    let request_context = request.extensions().get::<RequestContext>().cloned();
    let mut response = next.run(request).await;

    if let Some(request_id) = request_id
        && let Ok(header_value) = HeaderValue::from_str(&request_id)
    {
        response
            .headers_mut()
            .insert(HeaderName::from_static(X_REQUEST_ID_HEADER), header_value);
    }

    if let Ok(header_value) = HeaderValue::from_str(API_VERSION) {
        response.headers_mut().insert(
            HeaderName::from_static(API_VERSION_HEADER),
            header_value.clone(),
        );
        response
            .headers_mut()
            .insert(HeaderName::from_static(ROUTE_VERSION_HEADER), header_value);
    }

    if let Some(request_context) = request_context
        && let Ok(header_value) = HeaderValue::from_str(&request_context.elapsed().to_string())
    {
        response
            .headers_mut()
            .insert(HeaderName::from_static(RESPONSE_TIME_HEADER), header_value);
    }

    response
}