use std::time::Instant;

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

use crate::transport::HEADER_X_REQUEST_ID;

const HEADER_X_TRACE_ID: axum::http::HeaderName = axum::http::HeaderName::from_static("x-trace-id");

#[derive(Clone, Debug)]
pub(super) struct RequestContext {
    pub(super) correlation: CorrelationContext,
    started_at: Instant,
}

impl RequestContext {
    pub(super) fn request_id(&self) -> &str {
        self.correlation.request_id().as_str()
    }

    pub(super) fn started_at(&self) -> Instant {
        self.started_at
    }

    pub(super) fn exemplar(&self) -> Option<ExemplarTraceId> {
        self.correlation
            .trace_context()
            .is_sampled()
            .then(|| ExemplarTraceId::parse(self.correlation.trace_id().as_str()))
            .flatten()
    }
}

pub(super) async fn apply_request_context(mut request: Request, next: Next) -> Response {
    let (request_id, remote_parent) = resolve_correlation(request.headers());
    let span = info_span!(
        target: "crank::trace",
        "mcp.request",
        request_id = %request_id,
        trace_id = tracing::field::Empty,
    );
    if let Some(remote_parent) = remote_parent.as_ref() {
        set_canonical_parent(&span, remote_parent);
    }
    let trace_context = crank_trace::trace_context_for_span(&span).unwrap_or_else(|| {
        remote_parent
            .as_ref()
            .map_or_else(TraceContext::generate, TraceContext::continue_local)
    });
    span.record("trace_id", trace_context.trace_id().as_str());
    let context = RequestContext {
        correlation: CorrelationContext::new(request_id, trace_context),
        started_at: Instant::now(),
    };
    request.extensions_mut().insert(context.clone());

    with_request_correlation(
        context.correlation.request_id().to_string(),
        context.correlation.trace_id().to_string(),
        async move {
            let mut response = next.run(request).instrument(span).await;
            if let Ok(value) = HeaderValue::from_str(context.correlation.request_id().as_str()) {
                response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
            }
            if let Ok(value) = HeaderValue::from_str(context.correlation.trace_id().as_str()) {
                response.headers_mut().insert(HEADER_X_TRACE_ID, value);
            }
            if let Some(exemplar) = context.exemplar() {
                response.extensions_mut().insert(exemplar);
            }
            response
        },
    )
    .await
}

fn resolve_correlation(headers: &axum::http::HeaderMap) -> (RequestId, Option<TraceContext>) {
    let _tracestate_accepted = one_auxiliary_header_within_budget(
        headers,
        "tracestate",
        TraceContext::tracestate_within_budget,
    );
    let _baggage_accepted =
        one_auxiliary_header_within_budget(headers, "baggage", TraceContext::baggage_within_budget);
    let mut request_ids = headers.get_all(HEADER_X_REQUEST_ID).iter();
    let request_id = request_ids.next().and_then(|value| value.to_str().ok());
    let request_id = if request_ids.next().is_some() {
        RequestId::generate()
    } else {
        RequestId::resolve(request_id)
    };
    let mut traceparents = headers.get_all("traceparent").iter();
    let traceparent = traceparents.next().and_then(|value| value.to_str().ok());
    let remote_parent = if traceparents.next().is_some() {
        None
    } else {
        traceparent.and_then(|value| TraceContext::parse(value).ok())
    };
    (request_id, remote_parent)
}

fn one_auxiliary_header_within_budget(
    headers: &axum::http::HeaderMap,
    name: &'static str,
    validate: fn(&str) -> bool,
) -> bool {
    let mut values = headers.get_all(name).iter();
    let value = values.next().and_then(|value| value.to_str().ok());
    values.next().is_none() && value.is_some_and(validate)
}

fn set_canonical_parent(span: &tracing::Span, context: &TraceContext) {
    let mut headers = axum::http::HeaderMap::new();
    if let Ok(value) = HeaderValue::from_str(context.traceparent()) {
        headers.insert("traceparent", value);
        set_remote_trace_parent(span, &headers);
    }
}