use crate::{index, metrics, state, Args};
use anyhow::Result;
use axum::body::{Body, Bytes};
use axum::extract::Request;
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{middleware, Router};
use axum_server::tls_rustls::RustlsConfig;
use http_body_util::BodyExt;
use log::Level::Debug;
use log::{debug, info, log_enabled};
use std::net::SocketAddr;
use tower::ServiceBuilder;
use tower_http::compression::CompressionLayer;
use tower_http::trace::TraceLayer;

async fn debug_middleware(
    req: Request,
    next: Next,
) -> std::result::Result<impl IntoResponse, (StatusCode, String)> {
    if !log_enabled!(Debug) {
        return Ok(next.run(req).await);
    }
    let (parts, body) = req.into_parts();
    let bytes = buffer_and_print("request", body).await?;
    let req = Request::from_parts(parts, Body::from(bytes));

    let res = next.run(req).await;

    let (parts, body) = res.into_parts();
    let bytes = buffer_and_print("response", body).await?;
    let res = Response::from_parts(parts, Body::from(bytes));

    Ok(res)
}

async fn buffer_and_print<B>(
    direction: &str,
    body: B,
) -> std::result::Result<Bytes, (StatusCode, String)>
where
    B: axum::body::HttpBody<Data = Bytes>,
    B::Error: std::fmt::Display,
{
    let bytes = match body.collect().await {
        Ok(collected) => collected.to_bytes(),
        Err(err) => {
            return Err((
                StatusCode::BAD_REQUEST,
                format!("failed to read {direction} body: {err}"),
            ));
        }
    };

    if let Ok(body) = std::str::from_utf8(&bytes) {
        debug!("{direction} body = {body:?}");
    }

    Ok(bytes)
}

pub(crate) async fn check_health() -> crate::error::Result<String> {
    Ok("Ok".to_string())
}

pub(crate) async fn serve(args: Args) -> Result<()> {
    let middleware_stack = ServiceBuilder::new()
        .layer(TraceLayer::new_for_http())
        .layer(CompressionLayer::new())
        .layer(middleware::from_fn(debug_middleware))
        .into_inner();

    let tls_config = RustlsConfig::from_pem_file(&args.tls_cert, &args.tls_key).await?;

    // build our application with a route
    // Customized verb can be enabled with matchit 0.8 after axum upgrades its dependency.
    let app = Router::new()
        .route("/healthz", get(check_health))
        .route("/metrics", get(|| async move { metrics::index_metrics() }))
        .route("/indexes", get(index::get_indexes))
        .route("/indexes/:index_name/search", post(index::search_index))
        .route("/indexes/:index_name/version", get(index::index_version))
        .route("/indexes", post(index::update_index))
        .with_state(state::new(args).await?)
        .layer(middleware_stack);
    // run it
    let addr = SocketAddr::from(([0, 0, 0, 0], 8443));
    info!("serving-service is listening on {}", addr);
    axum_server::bind_rustls(addr, tls_config)
        .serve(app.into_make_service())
        .await?;
    Ok(())
}