use std::collections::HashSet;

use async_trait::async_trait;
use log::{info, warn};
use mysql_async::prelude::Queryable;
use tici_common::index::IndexKey;
use tici_grpc::{IndexInfo, TableInfo};
use tici_import::{
    ImportJobID, ImportJobMeta, ImportTask, ImportTaskID, ImportTaskParams, ImportTaskSummary,
    TiDBImportTaskID,
};
use tici_service_discovery::election::leadership::Leadership;

use crate::{
    ShardInfo, ShardInfoUpdate,
    common::connection_pool::TiDBConnectionPool,
    meta::{
        metrics::BACKEND_DELAY_DURATION_HISTOGRAM,
        store::{
            self, IndexMetaProperties,
            import_task::{ImportJobStat, InactiveImportTask},
        },
    },
};

#[derive(Debug, Clone)]
pub struct ImportIndexProgress {
    pub job_id: ImportJobID,
    pub job_status: String,
    pub task_count: i64,
    pub running_task_count: i64,
    pub finished_task_count: i64,
    pub failed_task_count: i64,
}

/// `MetaStore` is to store meta data of meta service.
/// All its methods must be implemented as atomic, and persistent.
#[async_trait]
pub trait MetaStore: Send + Sync {
    // ================
    // Index management methods.

    async fn restore_indexes(
        &self,
    ) -> Result<Vec<(IndexInfo, TableInfo, String, String, IndexMetaProperties)>, anyhow::Error>;
    async fn init_index(
        &self,
        table_info: &TableInfo,
        index_info: &IndexInfo,
        cdc_s3_bucket: String,
        cdc_s3_prefix: String,
    ) -> Result<(), anyhow::Error>;
    async fn finish_index_append_base(&self, index_key: &IndexKey) -> Result<(), anyhow::Error>;
    async fn drop_index(&self, index_key: &IndexKey) -> Result<(), anyhow::Error>;

    // ================
    // Shard management methods.

    async fn restore_shards(&self) -> Result<Vec<ShardInfo>, anyhow::Error>;
    async fn init_shard(&self, new_shard: &ShardInfo) -> Result<u64, anyhow::Error>;
    async fn remove_shard(&self, shard: &ShardInfo) -> Result<(), anyhow::Error>;
    async fn split_shards(
        &self,
        original_shard: &ShardInfo,
        new_shard_infos: Vec<ShardInfo>,
        batch: usize,
    ) -> Result<Vec<ShardInfo>, anyhow::Error>;
    async fn merge_shards(
        &self,
        original_shards: Vec<&ShardInfo>,
        new_shard: ShardInfo,
    ) -> Result<ShardInfo, anyhow::Error>;
    async fn update_shard(
        &self,
        tag: &str,
        shard: &ShardInfo,
        shard_update: ShardInfoUpdate,
    ) -> Result<(), anyhow::Error>;

    // ================
    // Import job management methods.
    // Import job is a job corresponding to a ImportInto job or CreateIndex
    // job.

    /// Add a new job to the meta store.
    async fn create_or_get_import_job(
        &self,
        job_meta: &ImportJobMeta,
    ) -> Result<ImportJobID, anyhow::Error>;
    /// Mark the job as finished uploading.
    async fn finish_import_job_upload(
        &self,
        tidb_task_id: &TiDBImportTaskID,
    ) -> Result<(), anyhow::Error>;
    /// Mark the job as finished building indexes.
    async fn finish_import_job<'a>(&self, job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error>;
    /// Mark the job as failed after all tasks have reached a terminal state.
    async fn fail_import_job<'a>(&self, job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error>;
    async fn get_jobs_should_finish(&self) -> Result<Vec<ImportJobStat>, mysql_async::Error>;
    async fn get_running_import_task_count(&self) -> Result<usize, mysql_async::Error>;
    /// List persisted running import tasks so meta can reconcile ownership
    /// after restart/failover.
    async fn get_running_import_tasks(&self) -> Result<Vec<ImportTask>, mysql_async::Error>;
    async fn get_import_index_progress(
        &self,
        keyspace_id: u32,
        table_id: i64,
        index_id: i64,
    ) -> Result<Option<ImportIndexProgress>, anyhow::Error>;
    async fn get_active_import_jobs(&self) -> Result<Vec<ImportJobMeta>, anyhow::Error>;
    async fn get_external_terminal_task_keys(
        &self,
        task_keys: &[String],
    ) -> Result<HashSet<String>, anyhow::Error>;

    // ================
    // Import task management methods.
    // A task is generated by ImportInto or CreateIndex, and it contains all
    // the row data within a given key-range.

    /// The task data is stored in S3. Add the task to pending list for
    /// building index.
    async fn add_import_task(
        &self,
        tidb_task_id: &TiDBImportTaskID,
        params: ImportTaskParams,
    ) -> Result<(), anyhow::Error>;
    /// Fetch the next import task to build index.
    async fn next_import_task(
        &self,
        owner_addr: &str,
    ) -> Result<Option<ImportTask>, mysql_async::Error>;
    /// Reset a running import task to `init` when the current owner has lost
    /// it.
    async fn reclaim_import_task(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
    ) -> Result<bool, anyhow::Error>;
    /// Reset all running import tasks owned by the removed worker to `init`.
    async fn reclaim_import_tasks_by_worker(
        &self,
        owner_addr: &str,
    ) -> Result<Vec<ImportTaskID>, anyhow::Error>;
    /// Mark the jobs as failed and reset any still-running tasks under them to
    /// `init` in one transaction.
    async fn fail_import_jobs_and_reclaim_running_tasks(
        &self,
        job_ids: &[ImportJobID],
    ) -> Result<Vec<ImportTaskID>, anyhow::Error>;
    /// Fetch the import tasks that have been inactive for more than
    /// `inactive_secs`.
    async fn inactive_import_tasks(
        &self,
        inactive_secs: u64,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error>;
    /// Reset the inactive import tasks to 'init' status for re-processing.
    async fn reset_inactive_import_tasks(
        &self,
        inactive_secs: u64,
        limit: usize,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error>;
    /// Mark the task as finished, which means index is built successfully.
    async fn import_task_build_success(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
        owner_seq: u64,
        summary: ImportTaskSummary,
    ) -> Result<bool, anyhow::Error>;
    /// Mark the task as failed, which means building index meet unrecoverable
    /// error.
    async fn import_task_build_fail(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
        owner_seq: u64,
        err_msg: &str,
    ) -> Result<bool, anyhow::Error>;

    // Leadership
    fn get_leadership(&self) -> &dyn Leadership;
}

/// Use TiDB as MetaStore's backend.
pub struct MetaStoreImplTiDB<L: Leadership> {
    pools: TiDBConnectionPool,
    leadership: L,
}

// The database name used by the meta service in TiDB.
pub const META_DATABASE_NAME: &str = "tici";

impl<L: Leadership> MetaStoreImplTiDB<L> {
    pub fn new(pools: TiDBConnectionPool, leadership: L) -> Self {
        Self { pools, leadership }
    }

    async fn get_tidb_conn(&self) -> Result<mysql_async::Conn, mysql_async::Error> {
        match self.pools.get_conn().await {
            Ok(conn) => Ok(conn),
            Err(err) => {
                warn!("MetaStore failed to connect to TiDB server: {:?}", err);
                Err(err)
            }
        }
    }
}

fn get_all_meta_table_names() -> impl Iterator<Item = &'static str> {
    store::index::META_TABLE_NAMES
        .iter()
        .chain(store::shard::META_TABLE_NAMES.iter())
        .chain(store::import_task::META_TABLE_NAMES.iter())
        .copied()
}

fn get_missing_tables(all_tables: &HashSet<String>, tables_created: &[String]) -> Vec<String> {
    all_tables
        .difference(&tables_created.iter().cloned().collect())
        .cloned()
        .collect()
}

fn get_drop_statements() -> Vec<String> {
    store::index::META_TABLE_NAMES
            .iter()
            .chain(store::shard::META_TABLE_NAMES.iter())
            .chain(store::import_task::META_TABLE_NAMES.iter())
            .rev() // Drop meta tables in reverse order of creation
            .map(|&s| format!("DROP TABLE IF EXISTS {}.{}", META_DATABASE_NAME, s))
            .collect()
}

fn get_term_for_write<L: Leadership>(leadership: &L, action: &str) -> Result<i64, anyhow::Error> {
    leadership
        .try_get_leader_term()
        .ok_or_else(|| anyhow::anyhow!("meta instance is not leader when calling {action}"))
}

fn get_term_for_write_mysql<L: Leadership>(
    leadership: &L,
    action: &str,
) -> Result<i64, mysql_async::Error> {
    get_term_for_write(leadership, action)
        .map_err(|err| mysql_async::Error::Other(err.to_string().into()))
}

// Implement bootstrap related methods for MetaStoreImplTiDB.
impl<L: Leadership> MetaStoreImplTiDB<L> {
    pub async fn is_bootstrapped(pools: &TiDBConnectionPool) -> Result<bool, anyhow::Error> {
        info!("Check if MetaStore is bootstrapped");
        let mut conn = pools.get_conn().await?;

        let all_tables: HashSet<String> =
            get_all_meta_table_names().map(|s| s.to_string()).collect();
        let tables_created: Vec<String> = conn.query(format!(
            "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = '{}' AND TABLE_NAME IN ({})",
            META_DATABASE_NAME,
            all_tables
                .iter()
                .map(|s| format!("'{}'", s))
                .collect::<Vec<_>>()
                .join(",")
        )).await?;

        let tables_missing = get_missing_tables(&all_tables, &tables_created);
        if tables_missing.is_empty() {
            info!("MetaStore is bootstrapped");
            Ok(true)
        } else {
            info!(
                "MetaStore is not bootstrapped, missing tables={:?}",
                tables_missing
            );
            Ok(false)
        }
    }

    pub async fn ensure_meta_schema(pools: &TiDBConnectionPool) -> Result<(), anyhow::Error> {
        // By default, `max-index-length` is 3072, which is the maximum length of the
        // index key. And according to TiDB, the `start_key` already includes
        // table_id and index_id.
        info!("Ensure meta table schema");
        let mut conn = pools.get_conn().await?;
        store::index::ensure_schema(&mut conn).await?;
        store::shard::ensure_schema(&mut conn).await?;
        store::import_task::ensure_schema(&mut conn).await?;

        info!("Ensure meta table schema finished");
        Ok(())
    }

    /// Dangerous method to clear all meta tables.
    pub async fn clear_all_meta(pools: &TiDBConnectionPool) -> Result<(), anyhow::Error> {
        warn!("Clear meta table schema");
        let mut conn = pools.get_conn().await?;
        let statement = get_drop_statements().join(";\n");
        let _: Vec<String> = conn.query(statement).await?;
        warn!("Clear meta table schema done");
        Ok(())
    }

    fn get_term(&self, action: &str) -> Result<i64, anyhow::Error> {
        get_term_for_write(&self.leadership, action)
    }

    fn get_term_mysql(&self, action: &str) -> Result<i64, mysql_async::Error> {
        get_term_for_write_mysql(&self.leadership, action)
    }
}

// MySQL error code for "table not found"
const ERR_TABLE_NOT_FOUND: u16 = 1146;

#[async_trait]
impl<L: Leadership> MetaStore for MetaStoreImplTiDB<L> {
    // Index management methods.

    async fn restore_indexes(
        &self,
    ) -> Result<Vec<(IndexInfo, TableInfo, String, String, IndexMetaProperties)>, anyhow::Error>
    {
        info!("Restoring indexes from meta table");
        let mut conn = self.get_tidb_conn().await?;
        store::index::restore(&mut conn).await
    }

    async fn init_index(
        &self,
        table_info: &TableInfo,
        index_info: &IndexInfo,
        cdc_s3_bucket: String,
        cdc_s3_prefix: String,
    ) -> Result<(), anyhow::Error> {
        self.get_term(&format!(
            "init_index, index=({}, {}, {})",
            index_info.keyspace_id, index_info.table_id, index_info.index_id
        ))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["index_init"])
            .start_timer();
        store::index::init(
            &mut conn,
            table_info,
            index_info,
            cdc_s3_bucket,
            cdc_s3_prefix,
        )
        .await
    }

    async fn finish_index_append_base(&self, index_key: &IndexKey) -> Result<(), anyhow::Error> {
        self.get_term(&format!(
            "finish_index_append_base, index_key={:?}",
            index_key
        ))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["index_finish_append_base"])
            .start_timer();
        store::index::finish_append_base(&mut conn, index_key).await
    }

    async fn drop_index(&self, index_key: &IndexKey) -> Result<(), anyhow::Error> {
        self.get_term(&format!("drop_index, index_key={:?}", index_key))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["index_drop"])
            .start_timer();
        store::index::drop(&mut conn, index_key).await
    }

    // Shard management methods.

    async fn restore_shards(&self) -> Result<Vec<ShardInfo>, anyhow::Error> {
        let mut conn = self.get_tidb_conn().await?;
        store::shard::restore(&mut conn).await
    }

    async fn init_shard(&self, new_shard: &ShardInfo) -> Result<u64, anyhow::Error> {
        let leader_term = self.get_term(&format!(
            "init_shard, shard.index_key={:?}",
            new_shard.index_key
        ))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["shard_init"])
            .start_timer();
        store::shard::init(&mut conn, new_shard, leader_term).await
    }

    async fn split_shards(
        &self,
        original_shard: &ShardInfo,
        new_shard_infos: Vec<ShardInfo>,
        batch: usize,
    ) -> Result<Vec<ShardInfo>, anyhow::Error> {
        let leader_term = self.get_term(&format!(
            "split_shard, ori_shard_id={}",
            original_shard.shard_id
        ))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["shard_split"])
            .start_timer();
        store::shard::split_shards(
            &mut conn,
            original_shard,
            new_shard_infos,
            leader_term,
            batch,
        )
        .await
    }

    async fn update_shard(
        &self,
        tag: &str,
        shard: &ShardInfo,
        shard_update: ShardInfoUpdate,
    ) -> Result<(), anyhow::Error> {
        self.get_term(&format!(
            "update_shard({tag}), shard_id={}, index_key={:?}",
            shard.shard_id, shard.index_key
        ))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["shard_update"])
            .start_timer();
        store::shard::update_shard(&mut conn, tag, shard, shard_update).await
    }

    async fn merge_shards(
        &self,
        original_shards: Vec<&ShardInfo>,
        new_shard: ShardInfo,
    ) -> Result<ShardInfo, anyhow::Error> {
        let leader_term =
            self.get_term(&format!("merge_shard, ori_shards:{:?}", original_shards))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["shard_merge"])
            .start_timer();
        store::shard::merge_shards(&mut conn, original_shards, new_shard, leader_term).await
    }

    async fn remove_shard(&self, shard: &ShardInfo) -> Result<(), anyhow::Error> {
        self.get_term(&format!(
            "remove_shard, shard_id={}, index_key={:?}",
            shard.shard_id, shard.index_key
        ))?;
        let mut conn = self.get_tidb_conn().await?;
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["shard_remove"])
            .start_timer();
        store::shard::remove_shard(&mut conn, shard).await
    }

    // Import job/task management methods.

    async fn create_or_get_import_job(
        &self,
        job_meta: &ImportJobMeta,
    ) -> Result<ImportJobID, anyhow::Error> {
        self.get_term(&format!(
            "create_or_get_import_job, tidb_task_id={}",
            job_meta.tidb_task_id
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::create_or_get_job(conn, job_meta).await
    }

    async fn finish_import_job_upload(
        &self,
        tidb_task_id: &TiDBImportTaskID,
    ) -> Result<(), anyhow::Error> {
        self.get_term(&format!(
            "finish_import_job_upload, tidb_task_id={tidb_task_id}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::finish_job_upload(conn, tidb_task_id).await
    }

    async fn finish_import_job<'a>(&self, job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error> {
        self.get_term(&format!("finish_import_job, job_ids={job_ids:?}"))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::finish_job(conn, job_ids).await
    }

    async fn fail_import_job<'a>(&self, job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error> {
        self.get_term(&format!("fail_import_job, job_ids={job_ids:?}"))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::fail_job(conn, job_ids).await
    }

    async fn get_jobs_should_finish(&self) -> Result<Vec<ImportJobStat>, mysql_async::Error> {
        let conn = self.get_tidb_conn().await?;
        store::import_task::get_jobs_should_finish(conn)
            .await
            .or_else(|err| {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_TABLE_NOT_FOUND {
                        Ok(vec![])
                    } else {
                        Err(err)
                    }
                } else {
                    Err(err)
                }
            })
    }

    async fn get_running_import_task_count(&self) -> Result<usize, mysql_async::Error> {
        let conn = self.get_tidb_conn().await?;
        store::import_task::get_running_tasks_count(conn)
            .await
            .or_else(|err| {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_TABLE_NOT_FOUND {
                        Ok(0)
                    } else {
                        Err(err)
                    }
                } else {
                    Err(err)
                }
            })
    }

    async fn get_running_import_tasks(&self) -> Result<Vec<ImportTask>, mysql_async::Error> {
        let conn = self.get_tidb_conn().await?;
        store::import_task::get_running_tasks(conn)
            .await
            .or_else(|err| {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_TABLE_NOT_FOUND {
                        Ok(vec![])
                    } else {
                        Err(err)
                    }
                } else {
                    Err(err)
                }
            })
    }

    async fn get_import_index_progress(
        &self,
        keyspace_id: u32,
        table_id: i64,
        index_id: i64,
    ) -> Result<Option<ImportIndexProgress>, anyhow::Error> {
        let _timer = BACKEND_DELAY_DURATION_HISTOGRAM
            .with_label_values(&["import_get_index_progress"])
            .start_timer();
        let conn = self.get_tidb_conn().await?;
        store::import_task::get_index_progress(conn, keyspace_id, table_id, index_id)
            .await
            .or_else(|err| {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_TABLE_NOT_FOUND {
                        Ok(None)
                    } else {
                        Err(err)
                    }
                } else {
                    Err(err)
                }
            })
            .map_err(anyhow::Error::from)
    }

    async fn get_active_import_jobs(&self) -> Result<Vec<ImportJobMeta>, anyhow::Error> {
        let conn = self.get_tidb_conn().await?;
        store::import_task::get_active_jobs(conn)
            .await
            .or_else(|err| {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_TABLE_NOT_FOUND {
                        Ok(vec![])
                    } else {
                        Err(err)
                    }
                } else {
                    Err(err)
                }
            })
            .map_err(anyhow::Error::from)
    }

    async fn get_external_terminal_task_keys(
        &self,
        task_keys: &[String],
    ) -> Result<HashSet<String>, anyhow::Error> {
        let conn = self.get_tidb_conn().await?;
        store::import_task::get_external_terminal_task_keys(conn, task_keys).await
    }

    async fn add_import_task(
        &self,
        tidb_task_id: &TiDBImportTaskID,
        params: ImportTaskParams,
    ) -> Result<(), anyhow::Error> {
        self.get_term(&format!("add_import_task, tidb_task_id={tidb_task_id}"))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::add_task(conn, tidb_task_id, params).await
    }

    async fn next_import_task(
        &self,
        owner_addr: &str,
    ) -> Result<Option<ImportTask>, mysql_async::Error> {
        self.get_term_mysql("next_import_task")?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::next_task(conn, owner_addr)
            .await
            .or_else(|err| {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_TABLE_NOT_FOUND {
                        Ok(None)
                    } else {
                        Err(err)
                    }
                } else {
                    Err(err)
                }
            })
    }

    async fn reclaim_import_task(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
    ) -> Result<bool, anyhow::Error> {
        self.get_term(&format!(
            "reclaim_import_task, task_id={task_id}, owner_addr={owner_addr}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::reclaim_task(conn, task_id, owner_addr).await
    }

    async fn reclaim_import_tasks_by_worker(
        &self,
        owner_addr: &str,
    ) -> Result<Vec<ImportTaskID>, anyhow::Error> {
        self.get_term(&format!(
            "reclaim_import_tasks_by_worker, owner_addr={owner_addr}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::reclaim_tasks_by_worker(conn, owner_addr).await
    }

    async fn fail_import_jobs_and_reclaim_running_tasks(
        &self,
        job_ids: &[ImportJobID],
    ) -> Result<Vec<ImportTaskID>, anyhow::Error> {
        self.get_term(&format!(
            "fail_import_jobs_and_reclaim_running_tasks, job_ids={job_ids:?}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::fail_jobs_and_reclaim_running_tasks(conn, job_ids).await
    }

    async fn inactive_import_tasks(
        &self,
        inactive_secs: u64,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error> {
        let conn = self.get_tidb_conn().await?;
        store::import_task::inactive_tasks(conn, inactive_secs).await
    }

    async fn reset_inactive_import_tasks(
        &self,
        inactive_secs: u64,
        limit: usize,
    ) -> Result<Vec<InactiveImportTask>, mysql_async::Error> {
        self.get_term_mysql(&format!(
            "reset_inactive_import_tasks, inactive_secs={inactive_secs}, limit={limit}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::reset_inactive_tasks(conn, inactive_secs, limit).await
    }

    async fn import_task_build_success(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
        owner_seq: u64,
        summary: ImportTaskSummary,
    ) -> Result<bool, anyhow::Error> {
        self.get_term(&format!(
            "import_task_build_success, task_id={task_id}, owner_addr={owner_addr}, owner_seq={owner_seq}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::task_build_success(conn, task_id, owner_addr, owner_seq, summary).await
    }

    async fn import_task_build_fail(
        &self,
        task_id: ImportTaskID,
        owner_addr: &str,
        owner_seq: u64,
        err_msg: &str,
    ) -> Result<bool, anyhow::Error> {
        self.get_term(&format!(
            "import_task_build_fail, task_id={task_id}, owner_addr={owner_addr}, owner_seq={owner_seq}, err_msg={err_msg}"
        ))?;
        let conn = self.get_tidb_conn().await?;
        store::import_task::task_build_failed(conn, task_id, owner_addr, owner_seq, err_msg).await
    }

    fn get_leadership(&self) -> &dyn Leadership {
        &self.leadership
    }
}

#[cfg(test)]
mod tests {
    use tici_service_discovery::election::leadership::Leadership;

    use super::*;

    struct NoLeadership;

    #[async_trait]
    impl Leadership for NoLeadership {
        async fn stop_campaign(&self) -> Result<(), anyhow::Error> {
            Ok(())
        }

        fn try_get_leader_term(&self) -> Option<i64> {
            None
        }
    }

    struct FixedLeadership(i64);

    #[async_trait]
    impl Leadership for FixedLeadership {
        async fn stop_campaign(&self) -> Result<(), anyhow::Error> {
            Ok(())
        }

        fn try_get_leader_term(&self) -> Option<i64> {
            Some(self.0)
        }
    }

    #[test]
    fn test_get_all_meta_table_names() {
        let names: Vec<String> = get_all_meta_table_names().map(|s| s.to_string()).collect();
        assert_eq!(
            names,
            vec![
                "tici_index_meta".to_string(),
                "tici_shard_meta".to_string(),
                "tici_import_jobs".to_string(),
                "tici_import_jobs_task".to_string()
            ]
        );
    }

    #[test]
    fn test_get_missing_tables() {
        let all_tables: HashSet<String> =
            get_all_meta_table_names().map(|s| s.to_string()).collect();
        // Only index_meta is created.
        assert!(!get_missing_tables(&all_tables, &["tici_index_meta".to_string()]).is_empty());
        // Only shard_meta is created.
        assert!(!get_missing_tables(&all_tables, &["tici_shard_meta".to_string()]).is_empty());
        // Only import_jobs is created.
        assert!(!get_missing_tables(&all_tables, &["tici_import_jobs".to_string()]).is_empty());
        assert!(
            !get_missing_tables(
                &all_tables,
                &[
                    "tici_import_jobs".to_string(),
                    "tici_import_jobs_task".to_string()
                ]
            )
            .is_empty()
        );

        // No missing tables.
        assert!(
            get_missing_tables(
                &all_tables,
                &[
                    "tici_index_meta".to_string(),
                    "tici_shard_meta".to_string(),
                    "tici_import_jobs".to_string(),
                    "tici_import_jobs_task".to_string()
                ]
            )
            .is_empty()
        );
    }

    #[test]
    fn test_get_drop_statements() {
        let drop_statements = get_drop_statements();
        // The order of drop statements should be
        // import_jobs -> shard_meta -> index_meta
        assert_eq!(
            drop_statements,
            vec![
                "DROP TABLE IF EXISTS tici.tici_import_jobs_task".to_string(),
                "DROP TABLE IF EXISTS tici.tici_import_jobs".to_string(),
                "DROP TABLE IF EXISTS tici.tici_shard_meta".to_string(),
                "DROP TABLE IF EXISTS tici.tici_index_meta".to_string(),
            ]
        );
    }

    #[test]
    fn test_get_term_for_write_rejects_non_leader() {
        let err = get_term_for_write(&NoLeadership, "update_shard").unwrap_err();
        assert_eq!(
            err.to_string(),
            "meta instance is not leader when calling update_shard"
        );
    }

    #[test]
    fn test_get_term_for_write_returns_leader_term() {
        let term = get_term_for_write(&FixedLeadership(7), "update_shard").unwrap();
        assert_eq!(term, 7);
    }

    #[test]
    fn test_get_term_for_write_mysql_rejects_non_leader() {
        let err = get_term_for_write_mysql(&NoLeadership, "next_import_task").unwrap_err();
        assert_eq!(
            err.to_string(),
            "Other error: meta instance is not leader when calling next_import_task"
        );
    }
}