pub mod tenant;

use std::sync::Arc;

use crate::error::{Error, Result};
use log::*;
use sqlx::types::chrono::NaiveDateTime;
use sqlx::{FromRow, MySql, MySqlPool, Transaction};

pub const DEFAULT_SHARD_NAME: &str = "shard0";
//const MIN_PENDING_WINDOW: u64 = (5 * 60 * 1000) << 18;
const MIN_PENDING_WINDOW: u64 = 0;

#[derive(Debug, Clone, FromRow)]
pub struct CollectionMeta {
    pub tenant_token: String,
    pub collection_name: String,
    pub collection_id: u64,
    pub data_count: Option<i64>,
    pub meta: Option<String>,
    pub created_at: NaiveDateTime,
    pub updated_at: NaiveDateTime,
    pub data_version: u64,
}

impl CollectionMeta {
    pub async fn all(
        ctx: Arc<context::Context>,
        tenant_token: &str,
        pool: &MySqlPool,
    ) -> Result<Vec<CollectionMeta>> {
        match sqlx::query_as::<_, CollectionMeta>(
            r#"SELECT `tenant_token`, `collection_name`, `collection_id`, `data_count`, CAST(`meta` AS CHAR) AS `meta`, `created_at`, `updated_at`,
            `data_version` FROM `collections` WHERE `tenant_token` = ?"#,
        )
        .bind(tenant_token)
        .fetch_all(pool)
        .await {
            Ok(result) => Ok(result),
            Err(e) => {
                error!("[{}] Failed to get all collections for tenant '{}'. {}", ctx.id(), tenant_token, e);
                Err(Error::from(e).ctx(ctx))
            }
        }
    }

    pub async fn first(pool: &MySqlPool) -> Result<Option<CollectionMeta>> {
        match sqlx::query_as::<_, CollectionMeta>(
            r#"SELECT `tenant_token`, `collection_name`, `collection_id`, `data_count`, CAST(`meta` AS CHAR) AS `meta`, `created_at`, `updated_at`,
            `data_version` FROM `collections` ORDER BY `updated_at` DESC LIMIT 1"#,
        )
            .fetch_optional(pool)
            .await {
            Ok(result) => Ok(result),
            Err(e) => {
                error!("Failed to get first collection. {}", e);
                Err(Error::from(e))

            }
        }
    }

    pub async fn get(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
    ) -> Result<Option<CollectionMeta>> {
        match sqlx::query_as::<_, CollectionMeta>(
            r#"SELECT `tenant_token`, `collection_name`, `collection_id`, `data_count`, CAST(`meta` AS CHAR) AS `meta`, `created_at`, `updated_at`,
            `data_version` FROM `collections` WHERE `tenant_token` = ? AND `collection_name` = ?"#,
        )
            .bind(tenant_token.as_ref())
            .bind(collection_name.as_ref())
            .fetch_optional(pool)
            .await {
            Ok(result) => {
                if result.is_none() {
                    warn!("[{}] Collection {}:{} not found in meta-db", ctx.id(), tenant_token.as_ref(), collection_name.as_ref());
                }
                Ok(result)
            },
            Err(e) => {
                error!("[{}] Failed to get collection '{}'. {}", ctx.id(), collection_name.as_ref(), e);
                Err(Error::from(e).ctx(ctx))
            }
        }
    }

    pub async fn get_by_tenants(
        pool: &MySqlPool,
        tenant_tokens: Vec<String>,
    ) -> Result<Option<Vec<CollectionMeta>>> {
        let mut query = String::from("SELECT `tenant_token`, `collection_name`, `collection_id`, `data_count`, CAST(`meta` AS CHAR) AS `meta`, `created_at`, `updated_at`, `data_version` FROM `collections` WHERE `tenant_token` IN (");
        for _ in 0..tenant_tokens.len() - 1 {
            query.push_str("?, ");
        }
        query.push_str("?)");
        let mut query = sqlx::query_as::<_, CollectionMeta>(&query);
        for token in tenant_tokens {
            query = query.bind(token);
        }
        let updated = query.fetch_all(pool).await?;
        Ok(if updated.is_empty() {
            None
        } else {
            Some(updated)
        })
    }

    pub async fn get_by_tenant(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        tenant_token: impl AsRef<str>,
    ) -> Result<Option<Vec<CollectionMeta>>> {
        match sqlx::query_as::<_, CollectionMeta>(
            r#"SELECT `tenant_token`, `collection_name`, `collection_id`, `data_count`, CAST(`meta` AS CHAR) AS `meta`, `created_at`, `updated_at`,
            `data_version` FROM `collections` where `tenant_token` = ?"#,
        )
            .bind(tenant_token.as_ref())
            .fetch_all(pool)
            .await {
                Ok(result) => {
                    if result.is_empty() {
                        warn!("[{}] No collections found for tenant '{}'", ctx.id(), tenant_token.as_ref());
                        Ok(None)
                    } else {
                        Ok(Some(result))
                    }
                },
                Err(e) => {
                    error!("[{}] Failed to get collections for tenant '{}'. {}", ctx.id(), tenant_token.as_ref(), e);
                    Err(Error::from(e).ctx(ctx))
                }
            }
    }

    pub async fn update_data_count(
        db: &MySqlPool,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
        data_count: i64,
    ) -> Result<()> {
        sqlx::query(
            "UPDATE `collections` SET `data_count` = `data_count` + ? WHERE `tenant_token` = ? AND `collection_name` = ?",
        )
        .bind(data_count)
        .bind(tenant_token.as_ref())
        .bind(collection_name.as_ref())
        .execute(db)
        .await?;
        Ok(())
    }

    pub async fn get_meta(
        ctx: Arc<context::Context>,
        pool: &MySqlPool,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
    ) -> Result<Option<String>> {
        match sqlx::query_scalar::<_, String>(
            r#"SELECT CAST(`meta` AS CHAR) AS `meta` FROM `collections` WHERE `tenant_token` = ? AND `collection_name` = ?"#,
        )
            .bind(tenant_token.as_ref())
            .bind(collection_name.as_ref())
            .fetch_optional(pool)
            .await {
            Ok(result) => Ok(result),
            Err(e) => {
                error!(
                    "[{}] Failed to get meta for collection '{}'. {}",
                    ctx.id(),
                    collection_name.as_ref(),
                    e
                );
                Err(Error::from(e).ctx(ctx))
            }
        }
    }

    pub async fn new_index_version(pool: &MySqlPool, current_index_version: u64) -> Result<u64> {
        debug!("Current index version: {}", current_index_version);
        let (new_version, tx) = Self::new_version(pool, None).await?;
        let (max_data_version, ): (Option<u64>, ) = sqlx::query_as(
            "SELECT max(`data_version`) FROM `collections` WHERE `data_version` > ? AND `data_version` <= ?",
        )
        .bind(current_index_version)
        .bind(new_version - MIN_PENDING_WINDOW)
        .fetch_one(pool)
        .await?;
        let new_index_version = max_data_version.unwrap_or(current_index_version);
        debug!("New index version: {}", new_index_version);
        tx.commit().await?;
        Ok(new_index_version)
    }

    pub async fn updated_between(
        pool: &MySqlPool,
        start_version: u64,
        end_version: u64,
    ) -> Result<Option<Vec<CollectionMeta>>> {
        debug!(
            "Get collections updated in [{}, {}]",
            start_version, end_version
        );
        let updated = sqlx::query_as::<_, CollectionMeta>(
            r#"SELECT `tenant_token`, `collection_name`, `collection_id`, `created_at`, `updated_at`,
            `data_version` FROM `collections` WHERE `data_version` > ? AND `data_version` <= ?"#,
        )
        .bind(start_version)
        .bind(end_version)
        .fetch_all(pool)
        .await?;
        Ok(if updated.is_empty() {
            None
        } else {
            Some(updated)
        })
    }

    pub async fn updated_between_all(
        pool: &MySqlPool,
        start_version: u64,
        end_version: u64,
    ) -> Result<Option<Vec<CollectionMeta>>> {
        debug!(
            "Get collections updated in [{}, {}]",
            start_version, end_version
        );
        let updated = sqlx::query_as::<_, CollectionMeta>(
            r#"SELECT `tenant_token`, `collection_name`, `collection_id`, `created_at`, `updated_at`,
            `data_version` FROM `collections` WHERE `data_version` >= ? AND `data_version` <= ?"#,
        )
        .bind(start_version)
        .bind(end_version)
        .fetch_all(pool)
        .await?;
        Ok(if updated.is_empty() {
            None
        } else {
            Some(updated)
        })
    }

    pub async fn bump_version(
        db: &MySqlPool,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
    ) -> Result<u64> {
        let (data_version, mut tx) = Self::new_version(db, None).await?;
        let sql = format!(
            "UPDATE `collections` SET `data_version` = {} WHERE `tenant_token` = ? AND `collection_name` = ?",
            data_version
        );
        sqlx::query(&sql)
            .bind(tenant_token.as_ref())
            .bind(collection_name.as_ref())
            .execute(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(data_version)
    }

    pub async fn version(
        db: &MySqlPool,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
    ) -> Result<u64> {
        let (data_version,) =
            sqlx::query_as("SELECT `data_version` FROM `collections` WHERE `tenant_token` = ? AND `collection_name` = ?")
                .bind(tenant_token.as_ref())
                .bind(collection_name.as_ref())
                .fetch_one(db)
                .await?;
        Ok(data_version)
    }

    pub async fn create<'a>(
        tx: &mut Transaction<'a, MySql>,
        version: u64,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
        meta: impl AsRef<str>,
    ) -> Result<u64> {
        sqlx::query(
            r#"INSERT INTO `collections`(`tenant_token`, `collection_name`, `collection_id`,
            `data_version`, `meta`, `shard_name`, `schema_version`) VALUES (?, ?, ?, ?, ?, ?, 0)"#,
        )
        .bind(tenant_token.as_ref())
        .bind(collection_name.as_ref())
        .bind(version)
        .bind(version)
        .bind(meta.as_ref())
        .bind(DEFAULT_SHARD_NAME)
        .execute(&mut **tx)
        .await?;

        Ok(version)
    }

    pub async fn rename(
        db: &MySqlPool,
        tenant_token: &str,
        from_collection_name: impl AsRef<str>,
        to_collection_name: impl AsRef<str>,
    ) -> Result<()> {
        sqlx::query(
            "UPDATE `collections` SET `collection_name` = ? WHERE `tenant_token` = ? AND `collection_name` = ?",
        )
        .bind(to_collection_name.as_ref())
        .bind(tenant_token)
        .bind(from_collection_name.as_ref())
        .execute(db)
        .await?;
        Ok(())
    }

    pub async fn delete(
        db: &MySqlPool,
        tenant_token: impl AsRef<str>,
        collection_name: impl AsRef<str>,
    ) -> Result<u64> {
        let (version, mut tx) = Self::new_version(db, None).await?;
        let collection_id: Option<u64> = sqlx::query_scalar(
            "SELECT `collection_id` FROM `collections` WHERE `tenant_token` = ? AND `collection_name` = ?",
        )
            .bind(tenant_token.as_ref())
            .bind(collection_name.as_ref())
            .fetch_optional(&mut *tx)
            .await?;
        let collection_id = if let Some(x) = collection_id {
            x
        } else {
            return Ok(0);
        };

        let _ = sqlx::query(
            "DELETE FROM `collections` WHERE `tenant_token` = ? AND `collection_name` = ? AND `collection_id` = ?"
        )
            .bind(tenant_token.as_ref())
            .bind(collection_name.as_ref())
            .bind(collection_id)
            .execute(&mut *tx)
            .await?;

        let _ = sqlx::query(
            "REPLACE INTO `deleted_collections`(`tenant_token`, `collection_name`, `collection_id`, `data_version`, `shard_name`) VALUES(?, ?, ?, ?, ?)",
        )
        .bind(tenant_token.as_ref())
        .bind(collection_name.as_ref())
        .bind(collection_id)
        .bind(version)
        .bind(DEFAULT_SHARD_NAME)
        .execute(&mut *tx)
        .await?;

        tx.commit().await?;
        Ok(collection_id)
    }

    pub async fn new_version<'a>(
        db: &MySqlPool,
        tx: impl Into<Option<Transaction<'a, MySql>>>,
    ) -> Result<(u64, Transaction<'a, MySql>)> {
        let mut tx = match tx.into() {
            Some(tx) => tx,
            None => db.begin().await?,
        };
        let version = sqlx::query_scalar(
            "SELECT CAST(@@tidb_current_ts AS UNSIGNED INTEGER) AS `version` FROM DUAL",
        )
        .fetch_one(&mut *tx)
        .await?;
        Ok((version, tx))
    }

    pub async fn count(pool: &MySqlPool) -> Result<(u64, u64)> {
        let (count, tenant_count): (u64, u64) =
            sqlx::query_as("SELECT CAST(COUNT(*) AS UNSIGNED INTEGER), CAST(COUNT(DISTINCT `tenant_token`) AS UNSIGNED INTEGER) FROM `collections`")
                .fetch_one(pool)
                .await?;
        Ok((count, tenant_count))
    }
}