use std::collections::HashMap;

use anyhow::Context;
use http_body_util::BodyExt;
use hyper::HeaderMap;
use jiff::Zoned;
use neon_rs::core::ctx::wrap;
use salvo::{
    Depot, FlowCtrl, Handler, Request, Response, Writer, async_trait,
    http::{ReqBody, ResBody},
};
use serde_json::Value;

use crate::status::api_err::ApiErr;

const MAX_LOG_SIZE: usize = 32 << 10; // 32KB

#[derive(Default)]
pub struct Log;

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

#[async_trait]
impl Handler for Log {
    async fn handle(&self, req: &mut Request, depot: &mut Depot, resp: &mut Response, ctrl: &mut FlowCtrl) {
        let enter_time = Zoned::now();
        let req_method = req.method().to_string();
        let req_uri = req.uri().path_and_query().unwrap().to_string();
        let req_header = header_to_string(req.headers());

        // 获取请求体
        let req_body = match drain_req_body(req).await {
            Err(e) => {
                e.write(req, depot, resp).await;
                ctrl.skip_rest();
                return;
            }
            Ok(v) => v,
        };
        ctrl.call_next(req, depot, resp).await;

        // 获取响应体
        let resp_body = drain_resp_body(resp).await;
        let status = resp.status_code.unwrap_or_default().as_u16();

        // 请求时长
        let duration = Zoned::now().duration_since(&enter_time).as_millis_f64();

        // 日志记录
        tracing::info!(
            uri = req_uri,
            method = req_method,
            headers = req_header,
            request = req_body,
            response = resp_body,
            status = status,
            duration = duration,
        );
    }
}

fn header_to_string(h: &HeaderMap) -> String {
    let mut map: HashMap<String, Vec<String>> = HashMap::new();
    for k in h.keys() {
        let mut vals: Vec<String> = Vec::new();
        for v in h.get_all(k) {
            if let Ok(s) = v.to_str() {
                vals.push(s.to_string())
            }
        }
        map.insert(k.to_string(), vals);
    }
    serde_json::to_string(&map).unwrap_or_else(|_| String::from("<none>"))
}

async fn drain_req_body(req: &mut Request) -> Result<Option<String>, ApiErr> {
    // 取出 Body
    let body = req.take_body();
    // this wont work if the body is a long running stream
    let bytes = body
        .collect()
        .await
        .with_context(|| wrap("request body collect failed"))?
        .to_bytes();

    // 处理 Body
    let body_str = std::str::from_utf8(&bytes[..bytes.len().min(MAX_LOG_SIZE)])
        .map(|s| {
            // 解析JSON字符串
            if let Ok(v) = serde_json::from_str::<Value>(s) {
                // 将JSON数据格式化为紧凑的字符串
                if let Ok(s) = serde_json::to_string(&v) {
                    return s;
                }
            }
            s.to_string()
        })
        .ok();

    // 重置 Body
    req.replace_body(ReqBody::Once(bytes));

    Ok(body_str)
}

async fn drain_resp_body(resp: &mut Response) -> Option<String> {
    // 取出 Body
    let body = resp.take_body();
    // this wont work if the body is a long running stream
    let bytes = match body.collect().await {
        Ok(v) => v.to_bytes(),
        Err(e) => {
            tracing::error!(err = ?e, "[response] body.collect");
            return None;
        }
    };

    // 处理 Body
    let body_str = std::str::from_utf8(&bytes[..bytes.len().min(MAX_LOG_SIZE)])
        .map(|s| s.to_string())
        .ok();

    // 重置 Body
    resp.replace_body(ResBody::Once(bytes));
    body_str
}
