use crate::state::ApiState;
use crate::{collection, metrics, point, reports, snapshots, state, Args};
use anyhow::{anyhow, Result};
use axum::body::{Body, Bytes};
use axum::extract::{DefaultBodyLimit, FromRequestParts, Request, State};
use axum::http::request::Parts;
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use axum::routing::*;
use axum::{async_trait, middleware, Router};
use axum_server::tls_rustls::RustlsConfig;
use filetime::FileTime;
use http_body_util::BodyExt;
use log::Level::Debug;
use log::*;
use std::cmp::minmax;
use std::net::SocketAddr;
use std::time::Duration;
use tokio::time::sleep;
use tokio::{fs, spawn, try_join};
use tower::ServiceBuilder;
use tower_http::compression::CompressionLayer;
use tower_http::trace::TraceLayer;

pub(crate) struct ExtractApiKey(pub(crate) String);

#[async_trait]
impl<S> FromRequestParts<S> for ExtractApiKey
where
    S: Send + Sync,
{
    type Rejection = (StatusCode, &'static str);
    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        if let Some(api_key) = parts.headers.get("api-key") {
            let api_key = api_key
                .to_str()
                .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid `Api-Key` header"))?;
            Ok(ExtractApiKey(api_key.to_string()))
        } else {
            Err((StatusCode::UNAUTHORIZED, "`Api-Key` header is missing"))
        }
    }
}

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, None).await?;
    let req = Request::from_parts(parts, Body::from(bytes));

    let res = next.run(req).await;
    let code = res.status();
    let (mut parts, body) = res.into_parts();
    let bytes = buffer_and_print("response", body, Some(code)).await?;
    parts.headers.insert("Connection", "close".parse().unwrap());
    let res = Response::from_parts(parts, Body::from(bytes));

    Ok(res)
}

async fn buffer_and_print<B>(
    direction: &str,
    body: B,
    status: Option<StatusCode>,
) -> 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) {
        if status.is_some() && !status.unwrap().is_success() {
            metrics::error_metrics("axum", status.unwrap().as_str());
            warn!("{direction} body = {body:?}");
        } else {
            debug!("{direction} body = {body:?}");
        }
    }

    Ok(bytes)
}

pub(crate) async fn check_health(State(state): State<ApiState>) -> crate::error::Result<String> {
    state.check_health().await?;
    Ok("Ok".to_string())
}

pub(crate) async fn not_implemented() -> crate::error::Result<()> {
    Err(crate::error::QdrantError::NotImplemented)
}

async fn reload_certificate(config: RustlsConfig, tls_cert: String, tls_key: String) {
    let mut last_modified = None;
    loop {
        let [earliest, latest] = match (fs::metadata(&tls_cert).await, fs::metadata(&tls_key).await)
        {
            (Ok(cert), Ok(key)) => minmax(
                FileTime::from_last_modification_time(&cert),
                FileTime::from_last_modification_time(&key),
            ),
            (Err(e), _) => {
                warn!("failed to get cert file metadata: {e}");
                continue;
            }
            (_, Err(e)) => {
                warn!("failed to get key file metadata: {e}");
                continue;
            }
        };
        if last_modified.is_none() {
            last_modified = Some(latest);
        }
        if last_modified < Some(earliest) {
            match config.reload_from_pem_file(&tls_cert, &tls_key).await {
                Ok(_) => {
                    info!(
                        "Reload the updated certificate from {}/{} with timestamp {}",
                        &tls_cert,
                        &tls_key,
                        latest.unix_seconds(),
                    );
                    last_modified = Some(latest);
                }
                Err(e) => warn!("failed to reload TLS certificate: {e}"),
            }
        }
        sleep(Duration::from_secs(3600)).await;
    }
}

pub(crate) async fn serve(args: Args) -> Result<()> {
    let max_body_size_mb = std::env::var("MAX_BODY_SIZE_MB")
        .unwrap_or("500".to_string())
        .parse::<usize>()
        .unwrap_or(20);
    assert!(max_body_size_mb > 0);
    info!("Max body size: {} MB", max_body_size_mb);
    let middleware_stack = ServiceBuilder::new()
        .layer(TraceLayer::new_for_http())
        .layer(CompressionLayer::new())
        .layer(middleware::from_fn(debug_middleware))
        .layer(DefaultBodyLimit::max(max_body_size_mb * 1024 * 1024))
        .into_inner();

    let tls_config = RustlsConfig::from_pem_file(&args.tls_cert, &args.tls_key)
        .await
        .map_err(|e| {
            anyhow!(
                "Failed to load TLS certificate and key from {}/{}: {}",
                args.tls_cert,
                args.tls_key,
                e
            )
        })?;

    spawn(reload_certificate(
        tls_config.clone(),
        args.tls_cert.clone(),
        args.tls_key.clone(),
    ));

    let state = state::new(args).await?;

    // build our application with a route
    let app = Router::new()
        .route("/healthz", get(check_health))
        .route("/collections", get(collection::get_collections))
        .route(
            "/collections/:collection_name",
            get(collection::get_collection)
                .delete(collection::delete_collection)
                .put(collection::create_collection)
                .patch(collection::update_collection),
        )
        .route(
            "/collections/:collection_name/facet",
            post(collection::facet),
        )
        .route(
            "/collections/:collection_name/index",
            put(collection::create_field_index),
        )
        .route(
            "/collections/:collection_name/index/:field_name",
            delete(collection::delete_field_index),
        )
        .route(
            "/collections/:collection_name/exists",
            get(collection::does_collection_exist),
        )
        .route(
            "/collections/:collection_name/points",
            put(point::upsert_points).post(point::get_points),
        )
        .route(
            "/collections/:collection_name/points/:id",
            get(point::get_point),
        )
        .route(
            "/collections/:collection_name/points/vectors",
            put(point::update_vectors),
        )
        .route(
            "/collections/:collection_name/points/vectors/delete",
            post(point::delete_vectors),
        )
        .route(
            "/collections/:collection_name/points/delete",
            post(point::delete_points),
        )
        .route(
            "/collections/:collection_name/points/search",
            post(point::search_points),
        )
        .route(
            "/collections/:collection_name/points/search/:actions",
            post(not_implemented),
        )
        .route(
            "/collections/:collection_name/points/scroll",
            post(point::scroll_points),
        )
        .route(
            "/collections/:collection_name/points/count",
            post(point::count_points),
        )
        .route(
            "/collections/:collection_name/points/batch",
            post(not_implemented),
        )
        .route(
            "/collections/:collection_name/points/payload",
            post(point::set_payload).put(point::overwrite_payload),
        )
        .route(
            "/collections/:collection_name/points/payload/delete",
            post(point::delete_payload),
        )
        .route(
            "/collections/:collection_name/points/payload/clear",
            post(point::clear_payload),
        )
        .route("/pipeline/reports", get(reports::list_pipeline_reports))
        .route(
            "/collections/:collection_name/snapshots/upload",
            post(snapshots::upload_snapshot),
        )
        .route(
            "/collections/:collection_name/snapshots",
            post(snapshots::create_snapshot).get(snapshots::list_snapshots),
        )
        .route(
            "/collections/:collection_name/snapshots/:snapshot_name",
            delete(snapshots::delete_snapshot).get(snapshots::download_snapshot),
        )
        .route("/cluster", get(crate::storage::cluster))
        .with_state(state.clone())
        .layer(middleware_stack.clone());

    // run it
    let addr = SocketAddr::from(([0, 0, 0, 0], 6333));

    let metrics_app = Router::new()
        .route("/metrics", get(metrics::collection_metrics))
        .with_state(state);
    let metrics_addr = SocketAddr::from(([0, 0, 0, 0], 10080));

    info!(
        "Listening on {}, prometheus metrics on {}",
        addr, metrics_addr
    );
    try_join!(
        axum_server::bind_rustls(addr, tls_config).serve(app.into_make_service()),
        // axum_server::bind(addr).serve(app.into_make_service()),
        axum_server::bind(metrics_addr).serve(metrics_app.into_make_service()),
    )?;

    Ok(())
}