pub mod delete;
pub(crate) mod select;

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

use crate::error::{is_table_exists, Error, ErrorKind, Result};
use crate::meta::tenant::TENANT_COLLECTION_META_TABLE_NAME;
use delete::{DeleteCollectionRecordBuilder, COLLECTION_DELETED_TABLE};
use lazy_static::lazy_static;
use log::*;
use regex::Regex;
use select::*;
use sqlx::types::chrono::NaiveDateTime;
use sqlx::{FromRow, MySql, MySqlConnection, MySqlPool};

lazy_static! {
    pub static ref COLLECTION_NAME_PATTERN: Regex = Regex::new(r#"^[0-9a-zA-Z$_]+$"#).unwrap();
    static ref PAYLOAD_FIELD_NAME_PATTERN: Regex = Regex::new(r#"^[a-z0-9A-Z$_.]+$"#).unwrap();
}

#[derive(FromRow, Debug, Clone)]
pub struct CollectionRecord {
    pub id: String,
    #[sqlx(default)]
    pub vector: Option<String>,
    #[sqlx(default)]
    pub tokenizer_version: Option<u32>,
    #[sqlx(default)]
    pub distance: Option<f32>,
    #[sqlx(default)]
    pub version: Option<u64>,
    #[sqlx(default)]
    pub deleted_at: Option<NaiveDateTime>,
    #[sqlx(default)]
    pub payload: Option<String>,
}

impl CollectionRecord {
    pub fn new(
        id: String,
        vector: impl Into<Option<String>>,
        payload: impl Into<Option<String>>,
        version: impl Into<Option<u64>>,
    ) -> Self {
        Self {
            id,
            vector: vector.into(),
            tokenizer_version: Some(0),
            version: version.into(),
            payload: payload.into(),
            distance: None,
            deleted_at: None,
        }
    }

    pub async fn updated_between_all(
        db: &MySqlPool,
        collection_name: impl AsRef<str>,
        start_version: u64,
        end_version: u64,
    ) -> Result<Vec<CollectionRecord>> {
        let table_name = collection_name.as_ref();
        let result = sqlx::query_as::<_, CollectionRecord>(
            &format!(
                "SELECT `id`, `version`, CAST(`payload` AS CHAR) as `payload` FROM `{}` WHERE `version` >= ? AND `version` <= ?",
                table_name
            )
        )
            .bind(start_version)
            .bind(end_version)
            .fetch_all(db)
            .await?;

        info!(
            "Got {} updated documents from {} withing version [{}, {}]",
            result.len(),
            table_name,
            start_version,
            end_version
        );
        Ok(result)
    }

    pub async fn deleted_between_all(
        db: &MySqlPool,
        collection_name: impl AsRef<str>,
        start_version: u64,
        end_version: u64,
    ) -> Result<Vec<CollectionRecord>> {
        let table_name = collection_name.as_ref();
        let result = sqlx::query_as::<_, CollectionRecord>(&format!(
            "SELECT `id`, `version` FROM `{}{}` WHERE `version` > ? AND `version` <= ?",
            table_name, COLLECTION_DELETED_TABLE
        ))
        .bind(start_version)
        .bind(end_version)
        .fetch_all(db)
        .await?;

        info!(
            "Reading {} Deleted documents from [{}{}] withing version [{}, {}]",
            result.len(),
            table_name,
            COLLECTION_DELETED_TABLE,
            start_version,
            end_version
        );
        Ok(result)
    }

    fn tokens_fields_tokens(
        input: Option<HashMap<String, Vec<String>>>,
    ) -> (Vec<String>, Vec<Vec<String>>) {
        let input = match input {
            None => return (vec![], vec![]),
            Some(input) => input,
        };

        let mut fields = Vec::with_capacity(input.len());
        let mut tokens = Vec::with_capacity(input.len());
        for (field, ts) in input.iter() {
            fields.push(format!("`{}`", get_tokens_column_name(field).unwrap()));
            tokens.push(ts.clone());
        }

        (fields, tokens)
    }

    pub async fn upsert(
        ctx: Arc<context::Context>,
        db: &MySqlPool,
        collection_name: &str,
        records: Vec<CollectionRecord>,
        tokens: Vec<Option<HashMap<String, Vec<String>>>>,
    ) -> Result<u64> {
        assert!(records.len() == tokens.len());

        let mut tx = db
            .begin()
            .await
            .map_err(|e| Error::from(e).ctx(ctx.clone()))?;

        // for record in records.iter() {
        for i in 0..records.len() {
            let record = &records[i];
            let tokens = Self::tokens_fields_tokens(tokens[i].clone());
            assert!(tokens.0.len() == tokens.1.len());
            let mut sql = format!(
                "REPLACE INTO `{}`(`id`, `vector`, `payload`, `version`",
                collection_name
            );
            if !tokens.0.is_empty() {
                sql.push_str(", ");
                sql.push_str(&tokens.0.join(", "));
            }
            sql.push_str(") VALUES (?, ?, ?, ?");
            if !tokens.0.is_empty() {
                sql.push_str(&", ?".repeat(tokens.0.len()));
            }
            sql.push(')');
            log::info!("[{}] upsert sql: {}", ctx.id(), sql);

            let mut query = sqlx::query::<MySql>(&sql);
            query = query
                .bind(&record.id)
                .bind(&record.vector)
                .bind(&record.payload);
            if !tokens.0.is_empty() {
                query = query.bind(record.version);
            } else {
                query = query.bind(0);
            }
            for ts in tokens.1.iter() {
                query = query.bind(serde_json::to_string(ts).unwrap());
            }
            query
                .execute(&mut *tx)
                .await
                .map_err(|e| Error::from(e).ctx(ctx.clone()))?;
        }

        tx.commit()
            .await
            .map_err(|e| Error::from(e).ctx(ctx.clone()))?;

        Ok(records.len() as u64)
    }

    pub async fn count(
        ctx: Arc<context::Context>,
        db: &MySqlPool,
        collection_name: &str,
    ) -> Result<u64> {
        match sqlx::query_scalar(&CollectionRecordQueryBuilder::count(collection_name).build())
            .fetch_one(db)
            .await
        {
            Ok(count) => Ok(count),
            Err(e) => {
                error!(
                    "[{}] Failed to count collection '{}'. {}",
                    ctx.id(),
                    collection_name,
                    e
                );
                Err(Error::from(e).ctx(ctx))
            }
        }
    }

    pub fn is_valid_name(collection_name: &str) -> bool {
        COLLECTION_NAME_PATTERN.is_match(collection_name)
            && collection_name != TENANT_COLLECTION_META_TABLE_NAME
    }

    pub async fn drop_table(pool: &MySqlPool, collection_name: &str) -> Result<()> {
        sqlx::query(&format!("DROP TABLE IF EXISTS `{}`", collection_name))
            .execute(pool)
            .await?;
        Ok(())
    }

    pub async fn table_exists(pool: &MySqlPool, collection_name: &str) -> Result<bool> {
        match sqlx::query_as::<_, (String, String)>(&format!(
            "SHOW CREATE TABLE `{}`",
            collection_name
        ))
        .fetch_one(pool)
        .await
        {
            Ok(_) => Ok(true),
            Err(e) => {
                error!(
                    "Failed to check table exists for collection '{}'. {}",
                    collection_name, e
                );
                Ok(false)
            }
        }
    }

    pub async fn create_table(
        pool: &MySqlPool,
        collection_name: &str,
        vector_size: u64,
    ) -> Result<()> {
        let mut retry = true;
        let query = format!(
            r#"
                CREATE TABLE `{}` (
                    `id` VARCHAR(36) PRIMARY KEY,
                    `vector` VECTOR<FLOAT>({}) NOT NULL,
                    `version` BIGINT UNSIGNED NOT NULL,
                    `tokenizer_version` INT UNSIGNED NOT NULL DEFAULT 0,
                    `payload` JSON DEFAULT NULL,
                    `created_at` DATETIME DEFAULT CURRENT_TIMESTAMP,
                    `updated_at` DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    `deleted_at` DATETIME DEFAULT NULL,
                    KEY idx_deleted_at(`deleted_at`),
                    KEY idx_updated_at(`updated_at`),
                    KEY idx_data_version(`version`)
                );
                "#,
            collection_name, vector_size,
        );
        loop {
            return match sqlx::query(&query).execute(pool).await {
                Err(e) => {
                    if retry && is_table_exists(&e) {
                        Self::drop_table(pool, collection_name).await?;
                        retry = false;
                        continue;
                    }
                    warn!(
                        "Failed to create table for collection '{}'. {}",
                        collection_name, e
                    );
                    Err(e.into())
                }
                Ok(_) => Ok(()),
            };
        }
    }

    pub async fn add_tokens_column(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        collection_name: &str,
        text_field: &str,
    ) -> Result<()> {
        let query = format!(
            r#"
            ALTER TABLE `{}` ADD COLUMN `{}` TEXT;
            "#,
            collection_name,
            get_tokens_column_name(text_field)?,
        );
        if let Err(e) = sqlx::query(&query).execute(pool).await {
            error!(
                "[{}] Failed to add column 'tokens' for collection '{}'. {}",
                ctx.id(),
                collection_name,
                e
            );
            return Err(Error::from(e).ctx(ctx));
        }

        Ok(())
    }

    pub async fn drop_tokens_column(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        collection_name: &str,
        text_field: &str,
    ) -> Result<()> {
        let query = format!(
            r#"
            ALTER TABLE `{}` DROP COLUMN `{}`;
            "#,
            collection_name,
            get_tokens_column_name(text_field)?,
        );
        if let Err(e) = sqlx::query(&query).execute(pool).await {
            error!(
                "[{}] Failed to drop column 'tokens' for collection '{}'. {}",
                ctx.id(),
                collection_name,
                e
            );
            return Err(Error::from(e).ctx(ctx));
        }

        Ok(())
    }

    pub async fn drop_columns(pool: &mut MySqlConnection, collection_name: &str) {
        let query = format!(
            r#"
            ALTER TABLE `{}` DROP COLUMN `meta`, DROP COLUMN `content`;
            "#,
            collection_name,
        );
        if let Err(e) = sqlx::query(&query).execute(pool).await {
            error!(
                "Failed to drop columns[meta, content] for collection '{}'. {}",
                collection_name, e
            );
        }
    }

    pub async fn try_create_table(pool: &MySqlPool, collection_name: &str) -> Result<()> {
        let query = format!(
            r#"
                CREATE TABLE `{}` (
                    `id` VARCHAR(36) PRIMARY KEY,
                    `version` BIGINT UNSIGNED NOT NULL,
                    `created_at` DATETIME DEFAULT CURRENT_TIMESTAMP,
                    `updated_at` DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    KEY idx_data_version(`version`)
                );
                "#,
            collection_name,
        );
        match sqlx::query(&query).execute(pool).await {
            Err(e) => Err(e.into()),
            Ok(_) => Ok(()),
        }
    }

    pub async fn drop_vector_idx(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        collection_name: &str,
    ) {
        let query = format!("DROP INDEX IF EXISTS idx_vector ON `{}`", collection_name);
        match sqlx::query(&query).execute(pool).await {
            Ok(_) => {
                info!(
                    "[{}] Dropped vector index for collection '{}'",
                    ctx.id(),
                    collection_name
                );
            }
            Err(e) => {
                warn!(
                    "[{}] Failed to drop vector index for collection '{}'. {}",
                    ctx.id(),
                    collection_name,
                    e
                );
            }
        }
    }

    pub async fn create_vector_idx(pool: &MySqlPool, collection_name: &str) {
        let query: String = format!("ALTER TABLE `{}` SET TIFLASH REPLICA 2", collection_name);
        match sqlx::query(&query).execute(pool).await {
            Ok(_) => {
                info!("Setted TIFLASH REPLICA 2: '{}'", collection_name);
                let query = format!(
                    "CREATE VECTOR INDEX IF NOT EXISTS idx_vector ON `{}` ((VEC_COSINE_DISTANCE(`vector`)))",
                    collection_name
                );
                match sqlx::query(&query).execute(pool).await {
                    Ok(_) => {
                        info!("Created vector index for collection '{}'", collection_name);
                    }
                    Err(e) => {
                        warn!(
                            "Failed to create vector index for collection '{}'. {}",
                            collection_name, e
                        );
                    }
                }
            }
            Err(e) => {
                warn!(
                    "Failed to SET TIFLASH REPLICA 2 for collection '{}'. {}",
                    collection_name, e
                );
            }
        };
    }

    pub async fn move_tokens(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        collection_name: &str,
    ) -> Result<()> {
        let query = format!(
            r#"
            UPDATE `{}` SET `{}` = `tokens`
            "#,
            collection_name,
            get_tokens_column_name("page_content")?,
        );
        match sqlx::query(&query).execute(pool).await {
            Ok(_) => Ok(()),
            Err(e) => Err(Error::from(e).ctx(ctx)),
        }
    }

    pub async fn scroll(
        ctx: Arc<context::Context>,
        db: &MySqlPool,
        collection_name: &str,
        offset: u64,
        limit: u64,
    ) -> Result<Vec<CollectionRecord>> {
        let result = sqlx::query_as::<_, CollectionRecord>(
            &format!(
                "SELECT `id`, `vector`, CAST(`payload` AS CHAR) as `payload` FROM `{}` ORDER BY `version` LIMIT ?, ?",
                collection_name
            )
        )
            .bind(offset)
            .bind(limit)
            .fetch_all(db)
            .await?;

        info!(
            "[{}] Scrolled {} documents from collection '{}'",
            ctx.id(),
            result.len(),
            collection_name
        );
        Ok(result)
    }
}

fn is_valid_field_name(key: &str) -> bool {
    PAYLOAD_FIELD_NAME_PATTERN.is_match(key)
}

pub fn get_valid_field_name(key: &str) -> Result<String> {
    if !is_valid_field_name(key) {
        return Err(Error::new(
            ErrorKind::BadRequest,
            format!("Invalid field name: '{}'", key),
        ));
    }
    Ok(key.to_string())
}

pub fn get_payload_key(key: &str) -> Result<String> {
    if !is_valid_field_name(key) {
        return Err(Error::new(
            ErrorKind::BadRequest,
            format!("Invalid field name: '{}'", key),
        ));
    }
    Ok(format!("`payload`->>'$.{}'", key))
}

pub fn get_tokens_column_name(text_field: &str) -> Result<String> {
    if !is_valid_field_name(text_field) {
        return Err(Error::new(
            ErrorKind::BadRequest,
            format!("Invalid field name: '{}'", text_field),
        ));
    }
    Ok(format!("__tokens__{}", text_field))
}

pub fn restore_tokens_column_name(text_field: &str) -> String {
    text_field.replace("__tokens__", "")
}

#[derive(Clone, Debug, Default)]
pub struct CollectionRecordPredicateBuilder {
    sql: String,
    not_deleted: bool,
}

impl CollectionRecordPredicateBuilder {
    pub fn apply_filters(&mut self, filters_sql: Option<String>) {
        if let Some(filters_sql) = filters_sql {
            if !self.sql.is_empty() {
                self.sql.push_str(" AND ");
            }
            self.sql.push_str(&filters_sql);
        }
    }

    fn prepend(&mut self) {
        if !self.sql.is_empty() {
            self.sql.push_str(" AND ");
        }
    }
    pub fn is_empty(&self) -> bool {
        self.sql.is_empty()
    }
    pub fn id(&mut self) -> &mut Self {
        self.prepend();
        self.sql.push_str(" (`id` = ?)");
        self
    }

    pub fn id_in(&mut self, len: usize) -> &mut Self {
        if len == 0 {
            return self;
        }
        self.prepend();
        self.sql.push_str("(`id` IN (");
        for _ in 0..len {
            self.sql.push_str("?,");
        }
        self.sql.truncate(self.sql.len() - 1);
        self.sql.push_str("))");
        self
    }

    pub fn payload(&mut self, key: &str) -> &mut Self {
        self.prepend();
        self.sql.push_str("(`payload`->>'$.");
        self.sql.push_str(key);
        self.sql.push_str("' = ?)");
        self
    }

    pub fn not_deleted(&mut self) -> &mut Self {
        if self.not_deleted {
            return self;
        }
        self.prepend();
        self.sql.push_str("(`deleted_at` IS NULL)");
        self.not_deleted = true;
        self
    }

    pub fn deleted(&mut self) -> &mut Self {
        self.prepend();
        self.sql.push_str("(`deleted_at` IS NOT NULL)");
        self
    }

    pub fn distance(&mut self) -> &mut Self {
        self.prepend();
        self.sql.push_str("(`distance` < ?)");
        self
    }

    pub fn version(&mut self) -> &mut Self {
        self.prepend();
        self.sql.push_str("(`version` > ?)");
        self
    }

    pub fn version_lt(&mut self) -> &mut Self {
        self.prepend();
        self.sql.push_str("(`version` <= ?)");
        self
    }

    pub fn tokens(&mut self, text_field: &str) -> &mut Self {
        self.prepend();
        self.sql
            .push_str(&format!("(JSON_CONTAINS(`{}`, ?))", text_field));
        self
    }
}

#[derive(Clone, Debug, Default)]
pub struct CountCollectionRecordBuilder {
    sql: String,
    predicate: Option<CollectionRecordPredicateBuilder>,
}

impl CountCollectionRecordBuilder {
    fn new(table_name: &str) -> Self {
        Self {
            sql: format!("SELECT CAST(COUNT(id) AS UNSIGNED) FROM `{}`", table_name),
            predicate: None,
        }
    }

    pub fn predicate(&mut self) -> &mut CollectionRecordPredicateBuilder {
        self.predicate.get_or_insert_with(Default::default)
    }

    pub fn build(self) -> String {
        let mut predicate = self.predicate.unwrap_or_default();
        predicate.not_deleted();
        format!("{} WHERE {}", self.sql, predicate.sql)
    }
}

#[derive(Clone, Debug)]
pub struct UpdateCollectionRecordBuilder {
    sql: String,
    predicate: Option<CollectionRecordPredicateBuilder>,
}

impl UpdateCollectionRecordBuilder {
    fn new(table_name: &str, version: u64) -> Self {
        let sql = if version == 0 {
            format!("UPDATE `{}` SET `version` = `version`", table_name)
        } else {
            format!("UPDATE `{}` SET `version` = {}", table_name, version)
        };
        Self {
            sql,
            predicate: None,
        }
    }

    pub fn set_vector(&mut self) -> &mut Self {
        self.sql.push_str(", `vector` = ?");
        self
    }

    pub fn delete(&mut self) -> &mut Self {
        self.sql.push_str(", `deleted_at` = CURRENT_TIMESTAMP");
        self
    }

    pub fn set_payload(&mut self) -> &mut Self {
        self.sql
            .push_str(", `payload` = JSON_MERGE_PATCH(COALESCE(`payload`, '{}'), ?)");
        self
    }

    pub fn set_tokens(&mut self, text_field: &str) -> &mut Self {
        self.sql.push_str(&format!(
            ", `{}` = ?",
            get_tokens_column_name(text_field).unwrap()
        ));
        self
    }

    pub fn set_payload_with_path(&mut self, path: &str) -> &mut Self {
        self.sql.push_str(&format!(", `payload` = JSON_MERGE_PATCH(COALESCE(`payload`, '{{}}'), JSON_INSERT('{{}}', '$.{}', CAST(? AS JSON)))", get_valid_field_name(path).unwrap()));
        self
    }

    pub fn clear_payload(&mut self) -> &mut Self {
        self.sql.push_str(", `payload` = NULL");
        self
    }

    pub fn overwrite_payload(&mut self) -> &mut Self {
        self.sql.push_str(", `payload` = ?");
        self
    }

    pub fn delete_payload(&mut self, num: usize) -> &mut Self {
        self.sql
            .push_str(", `payload` = IF(`payload` IS NULL, NULL, JSON_REMOVE(`payload`");
        self.sql.push_str(&", ?".repeat(num));
        self.sql.push_str("))");
        self
    }

    pub fn predicate(&mut self) -> &mut CollectionRecordPredicateBuilder {
        self.predicate.get_or_insert_with(Default::default)
    }

    pub fn build(self) -> String {
        if let Some(mut predicate) = self.predicate {
            predicate.not_deleted();
            format!("{} WHERE {}", self.sql, predicate.sql)
        } else {
            self.sql
        }
    }
}

pub struct CollectionRecordQueryBuilder {}

impl CollectionRecordQueryBuilder {
    pub fn delete(collection_name: impl AsRef<str>, version: u64) -> DeleteCollectionRecordBuilder {
        DeleteCollectionRecordBuilder::new(collection_name.as_ref().to_string(), version)
    }

    pub fn select(collection_name: impl AsRef<str>) -> SelectCollectionRecordBuilder {
        SelectCollectionRecordBuilder::new(collection_name.as_ref().to_string())
    }

    pub fn count(collection_name: impl AsRef<str>) -> CountCollectionRecordBuilder {
        CountCollectionRecordBuilder::new(collection_name.as_ref())
    }

    pub fn update(collection_name: impl AsRef<str>, version: u64) -> UpdateCollectionRecordBuilder {
        UpdateCollectionRecordBuilder::new(collection_name.as_ref(), version)
    }
}