use std::collections::HashMap;

use anyhow::Context;
use axum::{
    body::Body,
    extract::Request,
    middleware::Next,
    response::{IntoResponse, Response},
};
use http_body_util::BodyExt;
use hyper::HeaderMap;
use jiff::Zoned;
use neon_rs::core::ctx::wrap;
use serde_json::Value;

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

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

pub async fn log_mw(request: Request, next: Next) -> Response {
    let enter_time = Zoned::now();
    let req_method = request.method().to_string();
    let req_uri = request.uri().to_string();
    let req_header = header_to_string(request.headers());

    // 获取body
    let (response, (req_body, resp_body)) = match drain_body(request, next).await {
        Err(e) => return e.into_response(),
        Ok(v) => v,
    };

    // 请求时长
    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,
        duration = duration,
    );

    response
}

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_body(request: Request, next: Next) -> Result<(Response, (Option<String>, Option<String>)), ApiErr> {
    // Request
    // 取出 Body
    let (req_parts, req_body) = request.into_parts();
    // this wont work if the body is a long running stream
    let req_bytes = req_body
        .collect()
        .await
        .with_context(|| wrap("request body collect failed"))?
        .to_bytes();
    // 处理 Body
    let req_body_str = std::str::from_utf8(&req_bytes[..req_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
    let request = Request::from_parts(req_parts, Body::from(req_bytes));

    // Response
    let response = next.run(request).await;
    // 取出 Body
    let (resp_parts, resp_body) = response.into_parts();
    // this wont work if the body is a long running stream
    let resp_bytes = resp_body
        .collect()
        .await
        .with_context(|| wrap("response body collect failed"))?
        .to_bytes();
    // 处理 Body
    let resp_body_str = std::str::from_utf8(&resp_bytes[..resp_bytes.len().min(MAX_LOG_SIZE)])
        .map(|s| s.to_string())
        .ok();
    // 重置 Body
    let response = Response::from_parts(resp_parts, Body::from(resp_bytes));

    Ok((response, (req_body_str, resp_body_str)))
}
