use std::collections::HashMap;
use std::sync::Arc;

use crate::collection::actions::*;
use crate::error::{QdrantError, Result};
use crate::point::actions::*;
use crate::protocol::index::PayloadIndexes;
use crate::protocol::points::{PointInsertOperations, PointStruct, ScrollRequest, ScrollResult};
use crate::protocol::snapshots::SnapshotDescription;
use crate::protocol::{ApiResponse, CollectionDetail, CreateCollection, PointIdType};
use axum::extract::Multipart;
use log::*;
use model::collection::CollectionRecord;
use model::error::Error;
use model::jobs::{Job, JobAction, JobStatus, JobTarget};
use reqwest::multipart::Part;
use serde_json::Value;
use sqlx::{MySql, Pool};

const SNAPSHOT_FORM_FIELD: &str = "snapshot";

pub(crate) async fn create_qdrant_collection(
    ctx: Arc<context::Context>,
    collection_detail: CollectionDetail,
    host: &str,
    target_collection: &str,
) -> Result<()> {
    debug!(
        "[{}] Creating qdrant collection: {} with: {:?}",
        ctx.id(),
        target_collection,
        collection_detail
    );
    let mut request_map: HashMap<String, Value> = HashMap::new();
    request_map.insert(
        "vectors".to_string(),
        serde_json::to_value(collection_detail.config.params.vectors)?,
    );
    request_map.insert(
        "hnsw_config".to_string(),
        serde_json::to_value(collection_detail.config.hnsw_config)?,
    );
    request_map.insert(
        "wal_config".to_string(),
        serde_json::to_value(collection_detail.config.wal_config)?,
    );
    request_map.insert(
        "optimizers_config".to_string(),
        serde_json::to_value(collection_detail.config.optimizer_config)?,
    );

    let snapshots_url = format!("{}/collections/{}", host, target_collection);
    debug!(
        "[{}] Creating qdrant collection: {} with: {:?}",
        ctx.id(),
        snapshots_url,
        serde_json::to_string(&request_map)?,
    );
    match reqwest::Client::new()
        .put(snapshots_url)
        .json(&request_map)
        .send()
        .await
    {
        Ok(resp) => {
            if !resp.status().is_success() {
                error!(
                    "[{}] Create collection failed: {}",
                    ctx.id(),
                    resp.text().await?
                );
                return Err(QdrantError::InternalServerErrorWithMsg(format!(
                    "[trace-id:{}] Create collection failed",
                    ctx.id(),
                )));
            }
            Ok(())
        }
        Err(err) => {
            error!("[{}] Failed to create collection: {}", ctx.id(), err);
            Err(QdrantError::InternalServerErrorWithMsg(format!(
                "[trace-id:{}] Failed to create collection",
                ctx.id(),
            )))
        }
    }
}

pub(crate) async fn upload_snapshot(
    ctx: Arc<context::Context>,
    meta_db: Arc<Pool<MySql>>,
    tenant_db: Arc<Pool<MySql>>,
    tenants: (&str, &str),
    targets: (&str, &str),
    mut multipart: Multipart,
) -> Result<String> {
    debug!(
        "[{}] Uploading snapshot to collection: {}: {}",
        ctx.id(),
        tenants.0,
        tenants.1
    );
    let snapshots_url = format!("{}/collections/{}/snapshots/upload", targets.0, targets.1,);

    // get snapshot bytes
    let mut bytes = None;
    while let Some(field) = multipart.next_field().await? {
        if field.name().is_none() || field.name().unwrap() != SNAPSHOT_FORM_FIELD {
            continue;
        }

        bytes = Some(field.bytes().await.unwrap());
    }

    let bytes = match bytes {
        None => {
            warn!("[{}] No snapshot provided", ctx.id());
            return Err(crate::error::QdrantError::BadRequestWithMsg(
                "No snapshot provided".to_string(),
            ));
        }
        Some(bytes) => bytes,
    };

    // Upload qdrant snapshot
    debug!("[{}] Uploading snapshot to: {}", ctx.id(), snapshots_url);
    let form = reqwest::multipart::Form::new().part(SNAPSHOT_FORM_FIELD, Part::stream(bytes));
    let response = reqwest::Client::new()
        .post(snapshots_url)
        .multipart(form)
        .send()
        .await?;

    if !response.status().is_success() {
        return Err(crate::error::QdrantError::InternalServerErrorWithMsg(
            format!("Failed to upload snapshot: {}", response.text().await?),
        ));
    }

    // Job queue
    let job_id = Job::new(
        JobTarget::Qdrant,
        JobAction::Delete,
        targets.1.to_string(),
        std::time::Duration::from_secs(3600),
    )
    .push(&meta_db)
    .await?;

    // Get target collection
    let collection_detail = reqwest::Client::new()
        .get(format!("{}/collections/{}", targets.0, targets.1))
        .send()
        .await?
        .text()
        .await?;

    let collection_detail =
        serde_json::from_str::<ApiResponse<CollectionDetail>>(&collection_detail)?
            .result
            .unwrap();
    debug!("[{}] Collection detail: {:?}", ctx.id(), collection_detail);

    // Delete old collection
    debug!(
        "[{}] Deleting collection: {}: {}",
        ctx.id(),
        tenants.0,
        tenants.1
    );
    delete_collection(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        tenants.0,
        tenants.1,
    )
    .await?;

    // Create new collection
    debug!(
        "[{}] Creating collection: {}: {}",
        ctx.id(),
        tenants.0,
        tenants.1
    );
    create_collection(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        tenants.0,
        tenants.1,
        CreateCollection {
            vectors: collection_detail.config.params.vectors,
            hnsw_config: collection_detail.config.hnsw_config,
        },
    )
    .await?;

    // Create collection indexes
    let payload_index_info: String = collection_detail.payload_schema.try_into()?;
    let payload_index_info: PayloadIndexes = payload_index_info.try_into()?;
    for field in payload_index_info.iter() {
        debug!(
            "[{}] Creating field index: {}: {}",
            ctx.id(),
            tenants.1,
            field.0,
        );
        create_field_index(ctx.clone(), tenant_db.clone(), tenants.1, field.into()).await?;
    }

    // Move points from target collection to new collection by /collections/:collection_name/scroll in loop
    let scroll_url = format!("{}/collections/{}/points/scroll", targets.0, targets.1);
    debug!(
        "[{}] Moving points from: {} to: {}:{}",
        ctx.id(),
        scroll_url,
        tenants.0,
        tenants.1,
    );
    upsert_points_qdrant2tidb(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        tenants.0,
        tenants.1,
        &scroll_url,
    )
    .await?;

    // clean up target collection
    clean_up_collection_on_qdrant(ctx.clone(), targets.0, targets.1).await?;

    Job::update_status(&meta_db, &job_id, JobStatus::Done)
        .await
        .map_err(|e| Error::from(e).ctx(ctx))?;

    Ok("Ok".to_string())
}

pub(crate) async fn clean_up_collection_on_qdrant(
    ctx: Arc<context::Context>,
    host: &str,
    collection_name: &str,
) -> Result<()> {
    debug!(
        "[{}] Cleanup collection: {}: {}",
        ctx.id(),
        host,
        collection_name
    );
    reqwest::Client::new()
        .delete(format!("{}/collections/{}", host, collection_name))
        .send()
        .await?
        .text()
        .await?;

    Ok(())
}

pub(crate) async fn upsert_points_qdrant2tidb(
    ctx: Arc<context::Context>,
    meta_db: Arc<Pool<MySql>>,
    tenant_db: Arc<Pool<MySql>>,
    tenant_token: &str,
    collection_name: &str,
    scroll_url: &str,
) -> Result<()> {
    let mut scroll_id: Option<PointIdType> = None;
    let batch_size = 10;
    loop {
        let request = ScrollRequest {
            limit: Some(batch_size),
            with_payload: true,
            with_vector: true,
            offset: scroll_id.clone(),
            filter: None,
        };
        let response = reqwest::Client::new()
            .post(scroll_url)
            .body(request.to_string())
            .send()
            .await?
            .text()
            .await?;
        let response = serde_json::from_str::<ApiResponse<ScrollResult>>(&response)?.result;
        let response = match response {
            None => break,
            Some(response) => response,
        };
        debug!("Scroll response: {:?}", response);

        let points: Vec<PointStruct> = response.points.iter().map(|point| point.into()).collect();
        upsert_points_list(
            ctx.clone(),
            meta_db.clone(),
            tenant_db.clone(),
            tenant_token,
            collection_name,
            points,
        )
        .await?;

        scroll_id = match response.next_page_offset {
            None => break,
            Some(scroll_id) => Some(scroll_id),
        };
    }

    Ok(())
}

pub(crate) async fn upsert_points_tidb2qdrant(
    ctx: Arc<context::Context>,
    tenant_db: Arc<Pool<MySql>>,
    collection_name: &str,
    upsert_url: &str,
) -> Result<()> {
    debug!(
        "[{}] Upsert points from TiDB to qdrant: {}",
        ctx.id(),
        upsert_url
    );
    let mut offset = 0;
    let limit = 1000;
    loop {
        let points =
            CollectionRecord::scroll(ctx.clone(), &tenant_db, collection_name, offset, limit)
                .await?;
        let points_size = points.len();
        if points_size == 0 {
            break;
        }

        let upsert_points: PointInsertOperations = points.into();
        debug!(
            "[{}] Upserting points: {}",
            ctx.id(),
            serde_json::to_string(&upsert_points)?
        );
        let response = reqwest::Client::new()
            .put(upsert_url)
            .json(&upsert_points)
            .send()
            .await?;
        if response.status().is_success() {
            offset += limit;
        } else {
            return Err(QdrantError::InternalServerErrorWithMsg(format!(
                "Failed to upsert points: {}",
                response.text().await?
            )));
        }

        if points_size < limit as usize {
            break;
        }
    }

    Ok(())
}

pub(crate) async fn create_qdrant_snapshot(
    ctx: Arc<context::Context>,
    snapshot_url: &str,
) -> Result<SnapshotDescription> {
    debug!(
        "[{}] Creating snapshot on qdrant: {}",
        ctx.id(),
        snapshot_url
    );
    let response = reqwest::Client::new().post(snapshot_url).send().await?;
    if !response.status().is_success() {
        return Err(QdrantError::InternalServerErrorWithMsg(format!(
            "Failed to create snapshot: {}",
            response.text().await?
        )));
    }

    let response = response.text().await?;
    let response: ApiResponse<SnapshotDescription> = serde_json::from_str(&response)?;
    let response = response.result.unwrap();
    Ok(response)
}

pub(crate) fn tidb_collection_snapshot_on_s3(
    tenant_token: &str,
    collection_name: &str,
    snapshot_name: &str,
) -> String {
    format!(
        "ai-gateway/{}/{}/{}",
        tenant_token, collection_name, snapshot_name
    )
}

pub(crate) fn qdrant_collection_snapshot_on_s3(
    tenant_token: &str,
    collection_name: &str,
    snapshot_name: &str,
) -> String {
    format!(
        "snapshots/{}__{}/{}",
        tenant_token, collection_name, snapshot_name
    )
}