use crate::error::{QdrantError, Result};
use crate::protocol::index::{
    CreateFieldIndex, PayloadIndexInfo, PayloadIndexes, PayloadSchemaType,
};
use crate::protocol::{
    CollectionDetail, CollectionExists, CollectionName, CreateCollection, Distance,
};
use crate::protocol::{UpdateStatus, UpdateStatus::*};
use model::collection::CollectionRecord;
use model::error::Error;
use model::meta::{tenant::TenantCollectionMeta, CollectionMeta};
use serde::{Deserialize, Serialize};
use sqlx::MySqlPool;
use std::sync::Arc;
use tokio::try_join;

pub(super) async fn get_collections(
    ctx: Arc<context::Context>,
    meta_db: Arc<MySqlPool>,
    tenant_token: &str,
) -> Result<Vec<CollectionName>> {
    let result = CollectionMeta::all(ctx.clone(), tenant_token, &meta_db)
        .await?
        .into_iter()
        .map(|c| CollectionName::new(c.collection_name))
        .collect();
    ctx.push_phase("get_collections", None);
    Ok(result)
}

pub(crate) async fn get_collection(
    ctx: Arc<context::Context>,
    meta_db: Arc<MySqlPool>,
    tenant_db: Arc<MySqlPool>,
    tenant_token: &str,
    collection_name: &str,
) -> Result<Option<CollectionDetail>> {
    let qdrant_meta: QdrantMeta = serde_json::from_str(
        match &CollectionMeta::get_meta(ctx.clone(), &meta_db, tenant_token, collection_name)
            .await?
        {
            Some(meta) => meta,
            None => return Err(QdrantError::CollectionNotFound),
        },
    )?;

    let (points_count, tenant_meta) = try_join!(
        CollectionRecord::count(ctx.clone(), &tenant_db, collection_name),
        TenantCollectionMeta::get_by_collection_name(
            ctx.clone(),
            &tenant_db,
            collection_name,
            false
        )
    )?;

    let tenant_meta = match tenant_meta {
        Some(payload) => payload,
        None => return Err(QdrantError::CollectionNotFoundWithTraceId(ctx.id().into())),
    };

    Ok(Some(CollectionDetail::new(
        points_count,
        qdrant_meta.qdrant.hnsw_config,
        qdrant_meta.qdrant.vectors,
        tenant_meta.payload_schema,
    )))
}

pub(super) async fn does_collection_exist(
    ctx: Arc<context::Context>,
    tenant_db: Arc<MySqlPool>,
    collection_name: &str,
) -> Result<CollectionExists> {
    Ok(CollectionRecord::table_exists(&tenant_db, collection_name)
        .await
        .map_err(|e| e.ctx(ctx))?
        .into())
}

pub(crate) async fn delete_collection(
    ctx: Arc<context::Context>,
    meta_db: Arc<MySqlPool>,
    tenant_db: Arc<MySqlPool>,
    tenant_token: &str,
    collection_name: &str,
) -> Result<u64> {
    let deleted_collection_id =
        CollectionMeta::delete(meta_db.as_ref(), tenant_token, collection_name)
            .await
            .map_err(|e| e.ctx(ctx.clone()))?;
    try_join!(
        TenantCollectionMeta::delete(tenant_db.as_ref(), collection_name, deleted_collection_id),
        CollectionRecord::drop_table(tenant_db.as_ref(), collection_name)
    )
    .map_err(|e| e.ctx(ctx))?;
    Ok(deleted_collection_id)
}

#[derive(Debug, Serialize, Deserialize)]
pub(crate) struct QdrantMeta {
    pub(crate) qdrant: CreateCollection,
}

pub(crate) async fn create_collection(
    ctx: Arc<context::Context>,
    meta_db: Arc<MySqlPool>,
    tenant_db: Arc<MySqlPool>,
    tenant_token: &str,
    collection_name: &str,
    create_collection: CreateCollection,
) -> Result<u64> {
    if create_collection.vectors.distance != Distance::Cosine {
        return Err(QdrantError::BadRequestWithMsg(format!(
            "[{}] {} distance is not supported",
            ctx.id(),
            create_collection.vectors.distance,
        )));
    }

    let qdrant_meta = QdrantMeta {
        qdrant: create_collection,
    };

    let (version, mut tx) = CollectionMeta::new_version(&meta_db, None)
        .await
        .map_err(|e| e.ctx(ctx.clone()))?;
    let collection_id = CollectionMeta::create(
        &mut tx,
        version,
        tenant_token,
        collection_name,
        &serde_json::to_string(&qdrant_meta)?,
    )
    .await
    .map_err(|e| e.ctx(ctx.clone()))?;
    let vector_size = qdrant_meta.qdrant.vectors.size;
    try_join!(
        TenantCollectionMeta::create(tenant_db.as_ref(), collection_name, collection_id),
        CollectionRecord::create_table(tenant_db.as_ref(), collection_name, vector_size),
    )
    .map_err(|e| e.ctx(ctx.clone()))?;
    tx.commit()
        .await
        .map_err(|e| Error::from(e).ctx(ctx.clone()))?;
    Ok(collection_id)
}

pub(crate) async fn replace_collection(
    ctx: Arc<context::Context>,
    meta_db: Arc<MySqlPool>,
    tenant_db: Arc<MySqlPool>,
    tenant_token: &str,
    from_collection_name: &str,
    to_collection_name: &str,
) -> Result<()> {
    delete_collection(
        ctx.clone(),
        meta_db.clone(),
        tenant_db.clone(),
        tenant_token,
        to_collection_name,
    )
    .await?;
    try_join!(
        CollectionMeta::rename(
            &meta_db,
            tenant_token,
            from_collection_name,
            to_collection_name
        ),
        TenantCollectionMeta::rename(&tenant_db, from_collection_name, to_collection_name),
    )
    .map_err(|e| e.ctx(ctx))?;
    Ok(())
}

pub(crate) async fn create_field_index(
    ctx: Arc<context::Context>,
    tenant_db: Arc<MySqlPool>,
    collection_name: &str,
    create_field_index: CreateFieldIndex,
) -> Result<UpdateStatus> {
    let (tenant_meta, points_count) = try_join!(
        TenantCollectionMeta::get_by_collection_name(
            ctx.clone(),
            &tenant_db,
            collection_name,
            false
        ),
        CollectionRecord::count(ctx.clone(), &tenant_db, collection_name)
    )?;

    let tenant_meta =
        tenant_meta.ok_or(QdrantError::CollectionNotFoundWithTraceId(ctx.id().into()))?;

    // check if we can create text field
    let field_name = create_field_index.field_name.clone();
    let payload_index_info = match &create_field_index.field_schema {
        Some(schema) => {
            if schema.is_text() && points_count > 0 {
                return Err(QdrantError::BadRequestWithMsg(format!(
                    "[{}] Can't create text field for non-empty collection",
                    ctx.id()
                )));
            }
            PayloadIndexInfo::from(schema)
        }
        None => {
            return Err(QdrantError::BadRequestWithMsg(format!(
                "[{}] Field schema is empty",
                ctx.id()
            )))
        }
    };

    let mut payload_indexes: PayloadIndexes = match tenant_meta.payload_schema {
        Some(payload) => payload.try_into()?,
        None => PayloadIndexes::default(),
    };

    if payload_indexes.exists(&field_name) {
        return Ok(Completed);
    }

    if payload_indexes.iter().any(|(_, v)| {
        v.data_type == PayloadSchemaType::Text
            && payload_index_info.data_type == PayloadSchemaType::Text
    }) {
        return Err(QdrantError::BadRequestWithMsg(format!(
            "[{}] Can't create multiple text fields",
            ctx.id()
        )));
    }

    // create index
    index_ddl(
        ctx.clone(),
        tenant_db.clone(),
        collection_name,
        &create_field_index,
    )
    .await?;
    payload_indexes.insert(field_name, payload_index_info);

    // update payload schema
    TenantCollectionMeta::update_payload_index(
        ctx,
        &tenant_db,
        collection_name,
        payload_indexes.try_into()?,
    )
    .await?;

    Ok(Acknowledged)
}

async fn index_ddl(
    ctx: Arc<context::Context>,
    tenant_db: Arc<MySqlPool>,
    collection_name: &str,
    create_field_index: &CreateFieldIndex,
) -> Result<()> {
    if !create_field_index.field_schema.as_ref().unwrap().is_text() {
        TenantCollectionMeta::add_payload_idx(
            ctx.clone(),
            &tenant_db,
            collection_name,
            create_field_index.clone().try_into()?,
        )
        .await?;
    } else {
        CollectionRecord::add_tokens_column(
            ctx.clone(),
            &tenant_db,
            collection_name,
            &create_field_index.field_name,
        )
        .await?;
    }
    Ok(())
}

pub(crate) async fn delete_field_index(
    ctx: Arc<context::Context>,
    tenant_db: Arc<MySqlPool>,
    collection_name: &str,
    field_name: &str,
) -> Result<UpdateStatus> {
    let (tenant_meta, points_count) = try_join!(
        TenantCollectionMeta::get_by_collection_name(
            ctx.clone(),
            &tenant_db,
            collection_name,
            false
        ),
        CollectionRecord::count(ctx.clone(), &tenant_db, collection_name)
    )?;

    let tenant_meta =
        tenant_meta.ok_or(QdrantError::CollectionNotFoundWithTraceId(ctx.id().into()))?;

    let mut indexes: PayloadIndexes = match &tenant_meta.payload_schema {
        Some(indexes) => indexes.clone().try_into()?,
        None => PayloadIndexes::default(),
    };

    // check filed_name already deleted
    if !indexes.exists(field_name) {
        return Ok(Completed);
    }

    // check if we can delete text field
    let index = indexes.remove(field_name).unwrap();
    if index.data_type == PayloadSchemaType::Text && points_count > 0 {
        return Err(QdrantError::BadRequestWithMsg(format!(
            "[{}] Can't delete text field for non-empty collection",
            ctx.id()
        )));
    }

    // drop index
    if index.data_type == PayloadSchemaType::Text {
        CollectionRecord::drop_tokens_column(ctx.clone(), &tenant_db, collection_name, field_name)
            .await?;
    } else {
        TenantCollectionMeta::drop_payload_idx(
            ctx.clone(),
            &tenant_db,
            collection_name,
            field_name,
        )
        .await?;
    }

    // update payload schema
    TenantCollectionMeta::update_payload_index(
        ctx,
        &tenant_db,
        collection_name,
        indexes.try_into()?,
    )
    .await?;

    Ok(Acknowledged)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::protocol::{index::*, VectorsConfig};
    use crate::state::ApiState;
    use crate::{metrics, Args};

    async fn get_state() -> ApiState {
        crate::state::new(Args::default()).await.unwrap()
    }

    fn get_api_key() -> String {
        std::env::var("QDRANT_GATEWAY_APIKEY").unwrap()
    }

    #[tokio::test]
    async fn test_collection_actions() {
        let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());
        let state = get_state().await;
        let api_key = get_api_key();
        let (meta_db, tenant_db, tenant_token) =
            state.get_meta_and_tenant_db(&ctx, &api_key).await.unwrap();
        let meta_db = meta_db.unwrap();
        let collection_name = "test_collection_100";

        // Create collection
        let collection_id = create_collection(
            ctx.clone(),
            meta_db.clone(),
            tenant_db.clone(),
            &tenant_token,
            collection_name,
            CreateCollection {
                hnsw_config: None,
                vectors: VectorsConfig::default(),
            },
        )
        .await
        .unwrap();
        assert!(collection_id > 0);

        // Collection exists
        let exists = does_collection_exist(ctx.clone(), tenant_db.clone(), collection_name)
            .await
            .unwrap();
        assert!(bool::from(exists));

        // Get collection
        let detail = get_collection(
            ctx.clone(),
            meta_db.clone(),
            tenant_db.clone(),
            tenant_token,
            collection_name,
        )
        .await
        .unwrap();
        assert!(detail.unwrap().points_count == 0);

        // Get collections
        let collections = get_collections(ctx.clone(), meta_db.clone(), tenant_token)
            .await
            .unwrap();
        assert!(collections.len() > 0);

        // Delete collection
        let deleted_collection_id =
            delete_collection(ctx, meta_db, tenant_db, tenant_token, collection_name)
                .await
                .unwrap();
        assert!(deleted_collection_id == collection_id);
    }

    #[tokio::test]
    async fn test_field_index() {
        let ctx = context::new();
        let state = get_state().await;
        let api_key = get_api_key();
        let (meta_db, tenant_db, tenant_token) =
            state.get_meta_and_tenant_db(&ctx, &api_key).await.unwrap();
        let meta_db = meta_db.unwrap();
        let collection_name = "test_collection_101";

        // Delete collection
        let _ = delete_collection(
            ctx.clone(),
            meta_db.clone(),
            tenant_db.clone(),
            tenant_token,
            collection_name,
        )
        .await;

        // Create collection
        let collection_id = create_collection(
            ctx.clone(),
            meta_db.clone(),
            tenant_db.clone(),
            &tenant_token,
            collection_name,
            CreateCollection {
                hnsw_config: None,
                vectors: VectorsConfig::default(),
            },
        )
        .await
        .unwrap();
        assert!(collection_id > 0);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "page_content".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Text)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Delete field index
        let status = delete_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            "page_content",
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "group_id".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Keyword)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Delete field index
        let status =
            delete_field_index(ctx.clone(), tenant_db.clone(), collection_name, "group_id")
                .await
                .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "page_content_2".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldParams(
                    crate::protocol::index::PayloadSchemaParams::Text(TextIndexParams {
                        r#type: TextIndexType::Text,
                        tokenizer: TokenizerType::Multilingual,
                        min_token_len: Some(2),
                        max_token_len: Some(20),
                        lowercase: Some(true),
                        on_disk: None,
                    }),
                )),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "metadata.doc_id".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Uuid)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "metadata.updated".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Datetime)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "metadata.deleted".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Bool)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "metadata.score".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Float)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Create field index
        let status = create_field_index(
            ctx.clone(),
            tenant_db.clone(),
            collection_name,
            CreateFieldIndex {
                field_name: "metadata.score_int".to_string(),
                field_schema: Some(PayloadFieldSchema::FieldType(PayloadSchemaType::Bool)),
            },
        )
        .await
        .unwrap();
        assert_eq!(status, Acknowledged);

        // Delete collection
        let deleted_collection_id =
            delete_collection(ctx, meta_db, tenant_db, tenant_token, collection_name)
                .await
                .unwrap();
        assert!(deleted_collection_id == collection_id);
    }
}