use std::sync::LazyLock;

use http::{HeaderName, HeaderValue, header::AUTHORIZATION};
use salvo::{Depot, FlowCtrl, Handler, Request, Response, async_trait};
use tracing::Instrument;
use uuid::Uuid;

use crate::util::iden::Identity;

pub const TRACE_ID: HeaderName = HeaderName::from_static("x-trace-id");

static HOSTNAME: LazyLock<String> = LazyLock::new(|| hostname::get().ok().and_then(|s| s.into_string().ok()).unwrap_or_default());

#[derive(Default)]
pub struct Trace;

impl Trace {
    #[inline]
    pub fn new() -> Self {
        Trace {}
    }
}

#[async_trait]
impl Handler for Trace {
    async fn handle(&self, req: &mut Request, depot: &mut Depot, resp: &mut Response, ctrl: &mut FlowCtrl) {
        // 生成 TraceId
        let trace_id = match req.header::<String>(TRACE_ID) {
            Some(v) => {
                if !v.is_empty() {
                    v
                } else {
                    gen_trace_id(req)
                }
            }
            None => gen_trace_id(req),
        };

        // 解析 Identity
        let token = req.header::<String>(AUTHORIZATION);
        let iden = match token {
            None => Identity::default(),
            Some(v) => Identity::from_auth_token(v),
        };
        let iden_str = iden.to_string();

        // 设置 Identity
        depot.insert("trace_id", trace_id.clone());
        depot.inject(iden);

        // 设置 trace span
        let span = tracing::info_span!("trace", hostname = HOSTNAME.as_str(), trace_id, identity = iden_str);
        ctrl.call_next(req, depot, resp).instrument(span).await;
        // 设置 Header
        resp.headers_mut()
            .insert(TRACE_ID, HeaderValue::from_str(&trace_id).unwrap_or(HeaderValue::from_static("")));
    }
}

fn gen_trace_id(req: &mut Request) -> String {
    let id = Uuid::new_v4().simple().to_string();
    req.headers_mut()
        .insert(TRACE_ID, HeaderValue::from_str(&id).unwrap_or(HeaderValue::from_static("")));
    id
}
