use axum::body::Body;
use axum::http::{HeaderMap, HeaderValue, Request, StatusCode};
use axum::response::Response;
use serde_json::json;

pub(crate) const PROXY_REQUEST_ID_HEADER: &str = "x-athena-proxy-request-id";
pub(crate) const REQUEST_ID_HEADER: &str = "x-request-id";
pub(crate) const RESPONSE_TRACE_ID_HEADER: &str = "x-athena-trace-id";

fn header_value(headers: &HeaderMap, name: &str) -> Option<String> {
    headers
        .get(name)
        .and_then(|value| value.to_str().ok())
        .map(str::to_string)
}

pub(crate) fn request_trace_id(headers: &HeaderMap) -> Option<String> {
    header_value(headers, PROXY_REQUEST_ID_HEADER)
        .or_else(|| header_value(headers, REQUEST_ID_HEADER))
}

pub(crate) fn ensure_request_trace_id<B>(req: &mut Request<B>) -> String {
    if let Some(trace_id) = request_trace_id(req.headers()) {
        return trace_id;
    }

    let trace_id = uuid::Uuid::new_v4().to_string();
    if let Ok(header_value) = HeaderValue::from_str(trace_id.as_str()) {
        req.headers_mut().insert(REQUEST_ID_HEADER, header_value);
    }
    trace_id
}

pub(crate) fn ensure_response_trace_headers<B>(response: &mut Response<B>, trace_id: &str) {
    let Ok(header_value) = HeaderValue::from_str(trace_id) else {
        return;
    };

    if !response.headers().contains_key(REQUEST_ID_HEADER) {
        response
            .headers_mut()
            .insert(REQUEST_ID_HEADER, header_value.clone());
    }

    if (response.status().is_client_error() || response.status().is_server_error())
        && !response.headers().contains_key(RESPONSE_TRACE_ID_HEADER)
    {
        response
            .headers_mut()
            .insert(RESPONSE_TRACE_ID_HEADER, header_value);
    }
}

pub(crate) fn response_trace_id<B>(response: &Response<B>) -> Option<String> {
    header_value(response.headers(), RESPONSE_TRACE_ID_HEADER)
        .or_else(|| header_value(response.headers(), REQUEST_ID_HEADER))
        .or_else(|| header_value(response.headers(), PROXY_REQUEST_ID_HEADER))
}

pub(crate) fn json_error_response(
    status: StatusCode,
    message: impl Into<String>,
    trace_id: &str,
) -> Response {
    let payload = json!({
        "message": message.into(),
        "version": env!("CARGO_PKG_VERSION"),
        "traceId": trace_id,
    })
    .to_string();

    let mut builder = Response::builder()
        .status(status)
        .header("content-type", "application/json");

    if let Ok(header_value) = HeaderValue::from_str(trace_id) {
        builder = builder
            .header(REQUEST_ID_HEADER, header_value.clone())
            .header(RESPONSE_TRACE_ID_HEADER, header_value);
    }

    builder
        .body(Body::from(payload))
        .unwrap_or_else(|_| Response::new(Body::from("{\"message\":\"Internal server error\"}")))
}
