use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use axum::body::Body;
use axum::http::{HeaderName, HeaderValue, Request, Response, header};
use bytes::Bytes;
use http_body_util::BodyExt;
use serde_json::Value;
use tower::{Layer, Service};
use uuid::Uuid;

pub const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id");
pub const LOG_ID_HEADER: HeaderName = HeaderName::from_static("x-log-id");

const MAX_INBOUND_ID_LEN: usize = 128;

tokio::task_local! {
    static REQUEST_ID: String;
}

#[derive(Debug, Clone)]
pub struct RequestId(pub String);

#[derive(Debug, Clone)]
pub struct LogId(pub String);

/// The request id of the request currently being handled, when one is set.
pub fn request_id() -> Option<String> {
    REQUEST_ID.try_with(|id| id.clone()).ok()
}

/// Runs `future` with `request_id` available to nested code, so queue
/// consumers keep the id of the request that produced the message.
pub async fn with_request_id<F: Future>(request_id: Option<&str>, future: F) -> F::Output {
    match request_id {
        Some(request_id) => REQUEST_ID.scope(request_id.to_string(), future).await,
        None => future.await,
    }
}

pub fn new_request_id() -> String {
    new_id()
}

#[derive(Debug, Clone, Copy, Default)]
pub struct AppTraceLayer;

impl AppTraceLayer {
    pub fn new() -> Self {
        Self
    }
}

impl<S> Layer<S> for AppTraceLayer {
    type Service = AppTraceService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        AppTraceService { inner }
    }
}

#[derive(Debug, Clone)]
pub struct AppTraceService<S> {
    inner: S,
}

impl<S> Service<Request<Body>> for AppTraceService<S>
where
    S: Service<Request<Body>, Response = Response<Body>> + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn 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<Body>) -> Self::Future {
        let request_id = inbound_id(req.headers().get(&REQUEST_ID_HEADER)).unwrap_or_else(new_id);
        let log_id = new_id();

        req.extensions_mut().insert(RequestId(request_id.clone()));
        req.extensions_mut().insert(LogId(log_id.clone()));

        let fut = REQUEST_ID.scope(request_id.clone(), self.inner.call(req));
        Box::pin(async move {
            let mut res = fut.await?;
            if let Ok(value) = HeaderValue::try_from(request_id.as_str()) {
                res.headers_mut().insert(REQUEST_ID_HEADER, value);
            }
            if let Ok(value) = HeaderValue::try_from(log_id.as_str()) {
                res.headers_mut().insert(LOG_ID_HEADER, value);
            }
            inject_ids(&mut res, &request_id, &log_id).await;
            Ok(res)
        })
    }
}

fn new_id() -> String {
    Uuid::now_v7().to_string()
}

fn inbound_id(value: Option<&HeaderValue>) -> Option<String> {
    let value = value?.to_str().ok()?.trim();
    if value.is_empty() || value.len() > MAX_INBOUND_ID_LEN {
        return None;
    }
    Some(value.to_string())
}

async fn inject_ids(res: &mut Response<Body>, request_id: &str, log_id: &str) {
    let is_json = res
        .headers()
        .get(header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .is_some_and(|value| value.starts_with("application/json"));
    if !is_json {
        return;
    }

    let (mut parts, body) = std::mem::take(res).into_parts();
    let bytes = match body.collect().await {
        Ok(collected) => collected.to_bytes(),
        Err(_) => {
            *res = Response::from_parts(parts, Body::empty());
            return;
        }
    };

    let output = match serde_json::from_slice::<Value>(&bytes) {
        Ok(Value::Object(mut map))
            if map.contains_key("request_id") && map.contains_key("log_id") =>
        {
            map.insert(
                "request_id".to_string(),
                Value::String(request_id.to_string()),
            );
            map.insert("log_id".to_string(), Value::String(log_id.to_string()));
            match serde_json::to_vec(&map) {
                Ok(rewritten) => Bytes::from(rewritten),
                Err(_) => bytes,
            }
        }
        _ => bytes,
    };

    if let Ok(length) = HeaderValue::try_from(output.len().to_string()) {
        parts.headers.insert(header::CONTENT_LENGTH, length);
    }
    *res = Response::from_parts(parts, Body::from(output));
}