pub(crate) mod actions;

use std::sync::Arc;

use crate::collection::actions::*;
use crate::collection_mutexes::{Action, CollectionMutexe, CollectionMutexeDbImpl};
use crate::error::{validate_collection_name, QdrantError, Result};
use crate::metrics;
use crate::protocol::snapshots::SnapshotDescription;
use crate::protocol::ApiResponse;
use crate::server::ExtractApiKey;
use crate::state::ApiState;
use axum::extract::{Multipart, Path, State};
use axum::response::IntoResponse;
use axum::Json;
use cloud::s3::S3Service;
use cloud::Storage;
use log::*;
use model::jobs::{Job, JobAction, JobTarget};
use model::snapshots::SnapshotDescription as SnapshotDescriptionModel;

pub(crate) async fn upload_snapshot(
    State(state): State<ApiState>,
    ExtractApiKey(api_key): ExtractApiKey,
    Path(collection_name): Path<String>,
    multipart: Multipart,
) -> Result<Json<ApiResponse<String>>> {
    let target_host = state.get_qudrant_api()?;

    let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());

    info!(
        "[{}] POST /collections/:collection_name/snapshots/upload with: collection_name:[ {} ]",
        ctx.id(),
        &collection_name
    );
    let collection_name = validate_collection_name(collection_name)?;
    let tmp_collection_name = format!("{}__tmp", collection_name);
    metrics::ROUTER_COLLECTION_WITH_PATH
        .with_label_values(&[&collection_name, "/snapshots/upload", "post"])
        .inc();

    // check api-key
    let (meta_db, tenant_db, tenant_token) = state.get_meta_and_tenant_db(&ctx, &api_key).await?;
    let meta_db = meta_db.ok_or(QdrantError::InternalServerErrorWithMsg(
        "meta-db not configured".to_string(),
    ))?;

    let target_collection = format!("{}__{}", tenant_token, collection_name);

    // check collection exists
    if let Ok(Some(collection_detail)) = get_collection(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        tenant_token,
        &collection_name,
    )
    .await
    {
        if let Err(err) = actions::create_qdrant_collection(
            ctx.clone(),
            collection_detail,
            target_host,
            &target_collection,
        )
        .await
        {
            warn!(
                "[{}] Failed to create collection: {}",
                ctx.id(),
                err
            );
        }
    }

    CollectionMutexeDbImpl::new(
        meta_db.clone(),
        tenant_token.into(),
        collection_name.clone(),
        Action::CreateSnapshot,
        std::time::Duration::from_secs(600),
    )
    .lock()
    .await
    .map_err(|e| e.ctx(ctx.clone()))?;

    // upload snapshot to collection__tmp
    let result = match actions::upload_snapshot(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        (tenant_token, &tmp_collection_name),
        (target_host, &target_collection),
        multipart,
    )
    .await
    {
        Ok(result) => {
            // replace collection
            replace_collection(
                ctx.clone(),
                meta_db.clone(),
                tenant_db.clone(),
                tenant_token,
                &tmp_collection_name,
                &collection_name,
            )
            .await?;
            result
        }
        Err(err) => {
            return Err(err);
        }
    };

    metrics::ROUTER_COLLECTION_WITH_PATH_HISTOGRAM
        .with_label_values(&[&collection_name, "/snapshots/upload", "post"])
        .observe(metrics::since(&ctx.started()));

    Ok(Json(ApiResponse::ok(ctx.started(), result)))
}

pub(crate) async fn create_snapshot(
    State(state): State<ApiState>,
    ExtractApiKey(api_key): ExtractApiKey,
    Path(collection_name): Path<String>,
) -> Result<Json<ApiResponse<SnapshotDescription>>> {
    let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());

    info!(
        "[{}] POST /collections/:collection_name/snapshots with: collection_name:[ {} ]",
        ctx.id(),
        &collection_name
    );

    // validate collection name
    let collection_name = validate_collection_name(collection_name)?;

    // check api-key
    let (meta_db, tenant_db, tenant_token) = state.get_meta_and_tenant_db(&ctx, &api_key).await?;
    let meta_db = meta_db.ok_or(QdrantError::InternalServerErrorWithMsg(
        "meta-db not configured".to_string(),
    ))?;

    // check collection exists, if not exists, return 404
    let collection_detail = match get_collection(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        tenant_token,
        &collection_name,
    )
    .await?
    {
        Some(collection_detail) => collection_detail,
        None => return Err(QdrantError::CollectionNotFoundWithTraceId(ctx.id().into())),
    };

    // create collection on qdrant
    let target_host = state.get_qudrant_api()?;
    let target_collection = format!("{}__{}", tenant_token, collection_name);
    match actions::clean_up_collection_on_qdrant(ctx.clone(), target_host, &target_collection).await
    {
        Ok(_) => {}
        Err(err) => {
            warn!(
                "[{}] Failed to clean up collection: {}",
                ctx.id(),
                err
            );
        }
    };
    actions::create_qdrant_collection(
        ctx.clone(),
        collection_detail,
        target_host,
        &target_collection,
    )
    .await?;

    // Job to delete collection on qdrant
    let job_id = Job::new(
        JobTarget::Qdrant,
        JobAction::Delete,
        target_collection.clone(),
        std::time::Duration::from_secs(3600),
    )
    .push(&meta_db)
    .await?;

    // upsert points from TiDB to qdrant
    let upsert_url = format!("{}/collections/{}/points", target_host, target_collection);
    actions::upsert_points_tidb2qdrant(
        ctx.clone(),
        tenant_db.clone(),
        &collection_name,
        upsert_url.as_str(),
    )
    .await?;

    // create snapshot on qdrant
    let snapshot_url = format!(
        "{}/collections/{}/snapshots",
        target_host, target_collection
    );
    let description = actions::create_qdrant_snapshot(ctx.clone(), &snapshot_url).await?;

    // move snapshot from qdrant to TiDB
    let s3_service: Arc<dyn Storage> = Arc::new(S3Service::default());
    let from_path = actions::qdrant_collection_snapshot_on_s3(
        tenant_token,
        &collection_name,
        &description.name,
    );
    let to_path =
        actions::tidb_collection_snapshot_on_s3(tenant_token, &collection_name, &description.name);
    s3_service.move_file(&from_path, &to_path).await?;

    // save snapshot meta to TiDB
    let snapshot_meta = SnapshotDescriptionModel::new(
        tenant_token.into(),
        collection_name.clone(),
        description.name,
        description.size as i64,
        description.checksum.unwrap_or_default(),
        to_path,
    );
    snapshot_meta.clone().save(&meta_db).await?;

    // cleanup collection on qdrant
    actions::clean_up_collection_on_qdrant(ctx.clone(), target_host, &target_collection).await?;

    // job status update
    Job::update_status(&meta_db, &job_id, model::jobs::JobStatus::Done).await?;

    metrics::ROUTER_COLLECTION_WITH_PATH_HISTOGRAM
        .with_label_values(&[&collection_name, "/snapshots", "post"])
        .observe(metrics::since(&ctx.started()));

    Ok(Json(ApiResponse::ok(
        ctx.started(),
        SnapshotDescription::from(snapshot_meta),
    )))
}

pub(crate) async fn list_snapshots(
    State(state): State<ApiState>,
    ExtractApiKey(api_key): ExtractApiKey,
    Path(collection_name): Path<String>,
) -> Result<Json<ApiResponse<Vec<SnapshotDescription>>>> {
    let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());

    info!(
        "[{}] GET /collections/:collection_name/snapshots with: collection_name:[ {} ]",
        ctx.id(),
        &collection_name
    );

    // validate collection name
    let collection_name = validate_collection_name(collection_name)?;

    // check api-key
    let (meta_db, _, tenant_token) = state.get_meta_and_tenant_db(&ctx, &api_key).await?;
    let meta_db = meta_db.ok_or(QdrantError::InternalServerErrorWithMsg(
        "meta-db not configured".to_string(),
    ))?;

    let descriptions =
        SnapshotDescriptionModel::query_list(&meta_db, tenant_token, &collection_name).await?;
    let descriptions: Vec<SnapshotDescription> = descriptions
        .iter()
        .map(|d| SnapshotDescription::from(d.clone()))
        .collect();

    metrics::ROUTER_COLLECTION_WITH_PATH_HISTOGRAM
        .with_label_values(&[&collection_name, "/snapshots", "get"])
        .observe(metrics::since(&ctx.started()));

    Ok(Json(ApiResponse::ok(ctx.started(), descriptions)))
}

pub(crate) async fn delete_snapshot(
    State(state): State<ApiState>,
    ExtractApiKey(api_key): ExtractApiKey,
    Path((collection_name, snapshot_name)): Path<(String, String)>,
) -> Result<Json<ApiResponse<()>>> {
    let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());

    info!(
        "[{}] DELETE /collections/:collection_name/snapshots/:snapshot_name with: collection_name:[ {} ], snapshot_name:[ {} ]",
        ctx.id(),
        &collection_name,
        &snapshot_name
    );

    // validate collection name
    let collection_name = validate_collection_name(collection_name)?;

    // check api-key
    let (meta_db, _, tenant_token) = state.get_meta_and_tenant_db(&ctx, &api_key).await?;
    let meta_db = meta_db.ok_or(QdrantError::InternalServerErrorWithMsg(
        "meta-db not configured".to_string(),
    ))?;

    SnapshotDescriptionModel::delete(&meta_db, tenant_token, &collection_name, &snapshot_name)
        .await?;
    let snapshot_path =
        actions::tidb_collection_snapshot_on_s3(tenant_token, &collection_name, &snapshot_name);
    let s3_service: Arc<dyn Storage> = Arc::new(S3Service::default());
    s3_service.delete_file(&snapshot_path).await?;

    metrics::ROUTER_COLLECTION_WITH_PATH_HISTOGRAM
        .with_label_values(&[&collection_name, "/snapshots/:snapshot_name", "delete"])
        .observe(metrics::since(&ctx.started()));

    Ok(Json(ApiResponse::ok(ctx.started(), ())))
}

pub(crate) async fn download_snapshot(
    State(state): State<ApiState>,
    ExtractApiKey(api_key): ExtractApiKey,
    Path((collection_name, snapshot_name)): Path<(String, String)>,
) -> Result<impl IntoResponse> {
    let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());

    info!(
        "[{}] GET /collections/:collection_name/snapshots/:snapshot_name with: collection_name:[ {} ], snapshot_name:[ {} ]",
        ctx.id(),
        &collection_name,
        &snapshot_name
    );

    // validate collection name
    let collection_name = validate_collection_name(collection_name)?;

    // check api-key
    let (_, _, tenant_token) = state.get_meta_and_tenant_db(&ctx, &api_key).await?;

    let snapshot_path =
        actions::tidb_collection_snapshot_on_s3(tenant_token, &collection_name, &snapshot_name);
    let s3_service: Arc<dyn Storage> = Arc::new(S3Service::default());

    // download snapshot from s3
    let tmp_path = format!("/tmp/{}", snapshot_name);
    s3_service.download_file(&snapshot_path, &tmp_path).await?;

    metrics::ROUTER_COLLECTION_WITH_PATH_HISTOGRAM
        .with_label_values(&[&collection_name, "/snapshots/:snapshot_name", "get"])
        .observe(metrics::since(&ctx.started()));

    let content_disposition = format!("attachment; filename={}", snapshot_name);
    let body = axum::http::Response::builder()
        .header("Content-Disposition", content_disposition)
        .header("Content-Type", "application/octet-stream")
        .body(std::fs::read(tmp_path).unwrap())
        .unwrap();

    Ok(body.into_body())
}