use fts_index::protocol::Query;
use model::collection::delete::CollectionRecordDeleted;
use model::collection::CollectionRecord;
use model::collection::CollectionRecordQueryBuilder;
use model::error::column_not_found;
use model::error::Error;
use model::meta::CollectionMeta;
use model::meta::DEFAULT_SHARD_NAME;
use sqlx::Execute;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use tokio::spawn;
use tokio::try_join;

use crate::error::*;
use crate::point::try_repair_field_schema;
use crate::protocol::points::*;
use crate::state::State;
use log::*;
use sqlx::MySqlPool;

pub(super) struct MixQuery {
    state: Arc<State>,
    tenant_token: String,
    meta_db: Option<Arc<MySqlPool>>,
    tenant_db: Arc<MySqlPool>,
    collection_name: String,
    collection_id: u64,
    with_payload: bool,
    with_vector: bool,
    with_version: bool,
    offset: Option<u64>,
    limit: Option<u32>,
    point_query_param: Option<PointQueryParameters>,
    score_threshold: Option<f32>,
    vector: Option<DenseVector>,
    ctx: Arc<context::Context>,
}

impl MixQuery {
    pub fn new(
        ctx: Arc<context::Context>,
        state: Arc<State>,
        tenant_token: String,
        meta_db: Option<Arc<MySqlPool>>,
        tenant_db: Arc<MySqlPool>,
        collection_name: String,
        collection_id: u64,
    ) -> Self {
        Self {
            state,
            tenant_token,
            meta_db,
            tenant_db,
            collection_name,
            collection_id,
            with_payload: false,
            with_vector: false,
            with_version: false,
            offset: None,
            limit: None,
            point_query_param: None,
            score_threshold: None,
            vector: None,
            ctx,
        }
    }

    pub fn with_payload(&mut self, payload: bool) -> &mut Self {
        self.with_payload = payload;
        self
    }

    pub fn with_vector(&mut self, vector: bool) -> &mut Self {
        self.with_vector = vector;
        self
    }

    pub fn with_version(&mut self, version: bool) -> &mut Self {
        self.with_version = version;
        self
    }

    pub fn offset(&mut self, offset: Option<u64>) -> &mut Self {
        self.offset = offset;
        self
    }

    pub fn limit(&mut self, limit: Option<u32>) -> &mut Self {
        self.limit = limit;
        self
    }

    pub fn point_query_param(
        &mut self,
        point_query_param: Option<PointQueryParameters>,
    ) -> &mut Self {
        self.point_query_param = point_query_param;
        self
    }

    pub fn vector(&mut self, vector: DenseVector, score_threshold: f32) -> &mut Self {
        self.vector = Some(vector);
        self.score_threshold = Some(score_threshold);
        self
    }

    pub fn need_filter_by_tokens(&self) -> bool {
        let param = match self.point_query_param.as_ref() {
            Some(param) => param,
            None => return false,
        };

        param
            .must
            .iter()
            .any(|predicate| matches!(predicate.1, ValuePredicate::Text(_)))
            || param
                .should
                .iter()
                .any(|predicate| matches!(predicate.1, ValuePredicate::Text(_)))
            || param
                .must_not
                .iter()
                .any(|predicate| matches!(predicate.1, ValuePredicate::Text(_)))
    }

    pub async fn execute<R: From<GenericPointModel>>(&self) -> Result<Vec<R>> {
        let index_version = self.state.get_index_version(DEFAULT_SHARD_NAME).await?;
        debug!(
            "[{}] MixQuery version: {}, query-parameters: {:?}",
            self.ctx.id(),
            index_version,
            self.point_query_param
        );

        if index_version == 0 || !self.need_filter_by_tokens() {
            return self.query_points(self.tenant_db.clone(), 0).await;
        }

        let (fts_result, db_result): (Vec<GenericPointModel>, Vec<GenericPointModel>) = try_join!(
            self.query_points_with_fts(
                self.tenant_token.to_string(),
                self.tenant_db.clone(),
                index_version
            ),
            self.query_points(self.tenant_db.clone(), index_version),
        )?;

        let mut combined_results: Vec<GenericPointModel> = fts_result.into_iter().collect();
        combined_results.extend(db_result);
        let mut unique_ids = HashSet::new();
        combined_results.retain(|point: &GenericPointModel| unique_ids.insert(point.id.clone()));
        combined_results.truncate(self.limit.unwrap_or(100) as usize);

        self.ctx.push_phase("mix_query_execute", None);

        Ok(combined_results.into_iter().map(Into::into).collect())
    }

    pub async fn query_points<R: From<GenericPointModel>>(
        &self,
        tenant_db: Arc<MySqlPool>,
        index_version: u64,
    ) -> Result<Vec<R>> {
        info!(
            "[{}] Querying points from db: collection[{}], index-version[{}], query-param[{:?}]",
            self.ctx.id(),
            self.collection_name,
            index_version,
            self.point_query_param
        );

        let query_params = match self.point_query_param.as_ref() {
            Some(params) => params,
            None => &PointQueryParameters::default(),
        };

        let mut query_builder = CollectionRecordQueryBuilder::select(self.collection_name.clone());
        query_builder
            .with_version(self.with_version)
            .with_payload(self.with_payload)
            .with_vector(self.with_vector);

        if self.vector.is_some() {
            query_builder.expression("Vec_Cosine_distance(`vector`, ?)", "distance");
            // query_builder.having().distance();
            query_builder.predicate();
            query_builder.order_by("distance", true);
        }

        query_params.apply_predicates(query_builder.predicate(), false);
        query_builder.predicate().not_deleted();
        if index_version > 0 {
            query_builder.predicate().version();
        }

        let mut result = Vec::new();
        let mut offset: u64 = self.offset.unwrap_or(0);
        let limit = self.limit.unwrap_or(100);

        let threshold = match self.score_threshold {
            Some(threshold) => (1.0 - threshold).abs(),
            None => 1.0,
        };

        let mut column_retry = true;
        loop {
            let mut query_builder = query_builder.clone();
            query_builder.offset(Some(offset)).limit(Some(limit));
            let sql = query_builder.build();
            let query = match self.vector.as_ref() {
                Some(vectors) => sqlx::query_as(&sql).bind(format!("{:?}", vectors)),
                None => sqlx::query_as(&sql),
            };

            let query = match index_version > 0 {
                true => bind_query_params!(query_params, query, false).bind(index_version),
                false => bind_query_params!(query_params, query, false),
            };

            info!("[{}] query sql: {}", self.ctx.id(), query.sql());

            let result_clips: Vec<GenericPointModel> =
                match query.fetch_all(tenant_db.as_ref()).await {
                    Ok(result) => result.into_iter().collect(),
                    Err(err) => {
                        if column_not_found(&err)
                            && column_retry
                            && err.to_string().contains("__tokens__page_content")
                        {
                            warn!(
                            "[{}] Column not found in [{}:{}], try to repair field schema: {:?}",
                            self.ctx.id(),
                            self.tenant_token,
                            self.collection_name,
                            err,
                        );
                            try_repair_field_schema(
                                self.ctx.clone(),
                                tenant_db.clone(),
                                self.tenant_token.clone(),
                                self.collection_name.clone(),
                            )
                            .await;
                            column_retry = false;
                            continue;
                        }
                        return Err(Error::from(err).ctx(self.ctx.clone()).into());
                    }
                };
            let result_size = result_clips.len();

            let result_clips: Vec<GenericPointModel> = result_clips
                .into_iter()
                .filter(|row: &GenericPointModel| row.deleted_at.is_none())
                .collect();

            self.ctx.push_phase("query_points", Some(sql.clone()));
            if let Some(meta_db) = self.meta_db.clone() {
                if self.ctx.slow(Duration::from_secs(1)) && self.vector.is_some() {
                    spawn(Self::points_optimization(
                        meta_db,
                        tenant_db.clone(),
                        self.tenant_token.clone(),
                        self.collection_name.clone(),
                    ));
                }
            }

            result.extend(result_clips.into_iter().take(limit as usize - result.len()));
            offset += limit as u64;

            if self.need_filter_by_tokens() // JSON_CONTAINS done
            || result.len() >= limit as usize // enough data
            || result_size < limit as usize // no more data in db
            || offset >= (limit as u64 * 10)
            // too many data
            {
                break;
            }
        }

        Ok(result
            .into_iter()
            .filter(|row: &GenericPointModel| {
                row.distance.is_none() || row.distance.unwrap_or(0.0) < threshold
            })
            .map(Into::into)
            .collect())
    }

    async fn points_optimization(
        meta_db: Arc<MySqlPool>,
        tenant_db: Arc<MySqlPool>,
        tenant_token: String,
        collection_name: String,
    ) {
        match CollectionMeta::bump_version(&meta_db, &tenant_token, &collection_name).await {
            Ok(version) => {
                if let Err(err) =
                    CollectionRecordDeleted::create_table(&tenant_db, &collection_name, version)
                        .await
                {
                    error!("Failed to create deleted table: {}", err);
                };
            }
            Err(e) => {
                error!("Failed to bump version: {}", e);
            }
        };

        CollectionRecord::create_vector_idx(&tenant_db, &collection_name).await;
    }

    async fn query_points_with_fts<R: From<GenericPointModel>>(
        &self,
        tenant_token: String,
        tenant_db: Arc<MySqlPool>,
        index_version: u64,
    ) -> Result<Vec<R>> {
        info!(
            "[{}] Querying points from fts: collection[{}], index-version[{}], query-param[{:?}]",
            self.ctx.id(),
            self.collection_name,
            index_version,
            self.point_query_param
        );
        let query_params = match self.point_query_param.as_ref() {
            Some(params) => params,
            None => {
                warn!(
                    "[{}] FTS query without query parameters, fallback to db query",
                    self.ctx.id()
                );
                return Ok(Vec::default());
            }
        };

        let (docs, fts_version, fts_search) = search_points_from_fts(
            self.ctx.clone(),
            tenant_token.to_string(),
            self.collection_id,
            query_params.into(),
            self.state.fts_client()?,
        )
        .await?;

        // fallback to db query
        if !fts_search {
            return self.query_points(tenant_db.clone(), fts_version).await;
        }

        // fts filter no result
        if self.need_filter_by_tokens() && docs.is_empty() {
            return Ok(Vec::default());
        }

        let mut query_builder = CollectionRecordQueryBuilder::select(self.collection_name.clone());
        query_params.apply_predicates(query_builder.predicate(), true);
        query_builder
            .with_version(self.with_version)
            .with_payload(self.with_payload)
            .with_vector(self.with_vector);
        query_builder.predicate().not_deleted();

        if self.vector.clone().is_some() {
            query_builder.expression("Vec_Cosine_distance(`vector`, ?)", "distance");
            query_builder.having().distance();
            query_builder.order_by("distance", true);
        }

        query_builder.predicate().id_in(docs.len());
        query_builder.predicate().version_lt().not_deleted();

        let sql = query_builder.build();
        let query = bind_query_params!(query_params, sqlx::query_as(&sql), true);

        // if vector query
        let query = match self.vector.as_ref() {
            Some(vectors) => query.bind(format!("{:?}", vectors)),
            None => query,
        };

        // if filter query
        let query = match self.need_filter_by_tokens() {
            true => docs.into_iter().fold(query, |query, doc| query.bind(doc.0)),
            false => query,
        };

        let query = query.bind(index_version);

        // if vector score_threshold
        let query = match self.score_threshold {
            Some(threshold) => query.bind((1.0 - threshold).abs()),
            None => query,
        };

        let sql = query.sql();
        info!("[{}] query with fts sql: {}", self.ctx.id(), sql,);

        let mut result: Vec<GenericPointModel> = match query.fetch_all(tenant_db.as_ref()).await {
            Ok(result) => result,
            Err(err) => return Err(Error::from(err).ctx(self.ctx.clone()).into()),
        };

        if index_version < fts_version {
            warn!(
                "[{}] FTS version {} is ahead of Qdrant-API {}",
                self.ctx.id(),
                fts_version,
                index_version
            );
        }

        // when fts version rollback but qdrant-api not sync yet, must complement the result from db
        if index_version > fts_version {
            warn!(
                "[{}] FTS version {} is behind of Qdrant-API {}, complementing the result from DB",
                self.ctx.id(),
                fts_version,
                index_version
            );
            let complements: Vec<GenericPointModel> =
                self.query_points(tenant_db, fts_version).await?;
            result.extend(complements);
        }

        self.ctx.push_phase("fts_join_query", Some(sql.into()));
        self.ctx.slow(Duration::from_secs(1));

        Ok(result.into_iter().map(Into::into).collect())
    }
}

const DEFAULT_FTS_RECALL_SIZE: u32 = 300;
pub(super) async fn search_points_from_fts(
    ctx: Arc<context::Context>,
    tenant_token: String,
    collection_id: u64,
    query: Option<Query>,
    fts_client: fts_index::client::Client,
) -> Result<(HashMap<String, f32>, u64, bool)> {
    let query = match query {
        Some(query) => query,
        None => return Ok((HashMap::new(), 0, false)),
    };

    match fts_client {
        fts_index::client::Client::Dummy {
            indexes: _,
            version: _,
        } => Ok((HashMap::new(), 0, false)),
        _ => {
            Ok(
                match fts_client
                    .search(
                        model::meta::DEFAULT_SHARD_NAME,
                        collection_id,
                        tenant_token,
                        query,
                        Some(DEFAULT_FTS_RECALL_SIZE),
                        None,
                    )
                    .await
                {
                    Ok(response) => {
                        ctx.push_phase("search_points_from_fts", None);
                        (
                            response
                                .hits
                                .unwrap_or_default()
                                .into_iter()
                                .map(|hit| (hit.doc_id, hit.score))
                                .collect(),
                            response.index_version,
                            true,
                        )
                    }
                    Err(err) => {
                        let err = err.ctx(ctx.clone());
                        warn!(
                            "[{}] Failed to search from search service fallback to filtering: {}",
                            ctx.id(),
                            err
                        );
                        // fallback to filtering
                        (HashMap::new(), 0, false)
                    }
                },
            )
        }
    }
}