// TODO(P2) We don't do any refactor until this path is covered with tests.
#![allow(clippy::unwrap_used)]

use std::collections::{HashMap, HashSet};

use anyhow::bail;
use log::{debug, info, warn};
use mysql_async::{Conn, Error, Params, Transaction, TxOpts, prelude::*};
use tici_import::{
    ImportJobID, ImportJobMeta, ImportJobSummary, ImportTask, ImportTaskID, ImportTaskParams,
    ImportTaskSummary, TiDBImportTaskID,
};

use crate::meta::store::metastore::META_DATABASE_NAME;

pub const META_TABLE_NAMES: [&str; 2] = ["tici_import_jobs", "tici_import_jobs_task"];

pub async fn ensure_schema(conn: &mut Conn) -> Result<(), anyhow::Error> {
    info!("Ensure ImportTaskStore table schema");
    let _s: Vec<String> = conn
        .query(format!(
            r"
            CREATE TABLE IF NOT EXISTS {0}.tici_import_jobs (
                `job_id` bigint UNSIGNED NOT NULL AUTO_INCREMENT,
                `tidb_task_id` varchar(256) NOT NULL COMMENT 'The unique task ID from TiDB job',
                `keyspace_id` int unsigned NOT NULL,
                `table_id` bigint NOT NULL,
                `index_ids` mediumtext DEFAULT NULL,
                `status` varchar(64) NOT NULL COMMENT 'init, failed, finished',
                `create_time` timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
                `start_time` timestamp(6) DEFAULT NULL,
                `update_time` timestamp(6) DEFAULT NULL,
                `end_time` timestamp(6) DEFAULT NULL,
                `parameters` json NOT NULL,
                `step` varchar(64) NOT NULL COMMENT '',
                `summary` json DEFAULT NULL,
                `error_message` text DEFAULT NULL,
                PRIMARY KEY (`job_id`) /*T![clustered_index] CLUSTERED */,
                UNIQUE KEY `idx_tidb_task_id` (`tidb_task_id`)
            ) AUTO_ID_CACHE 100;
            CREATE TABLE IF NOT EXISTS {0}.tici_import_jobs_task (
                `task_id` bigint UNSIGNED NOT NULL AUTO_INCREMENT,
                `job_id` bigint NOT NULL COMMENT 'The TiCI job ID this task belongs to',
                `status` varchar(64) NOT NULL COMMENT 'init, running, failed, finished',
                `create_time` timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
                `start_time` timestamp(6) DEFAULT NULL,
                `update_time` timestamp(6) DEFAULT NULL,
                `end_time` timestamp(6) DEFAULT NULL,
                `owner_addr` varchar(256) DEFAULT NULL COMMENT 'The worker node currently owning this task',
                `owner_seq` bigint UNSIGNED NOT NULL DEFAULT 0 COMMENT 'Incremented on each task assignment; finish/fail callbacks must match the assignment sequence',
                `parameters` json NOT NULL,
                `summary` json DEFAULT NULL,
                `error_message` text DEFAULT NULL,
                PRIMARY KEY (`task_id`) /*T![clustered_index] CLUSTERED */,
                KEY `idx_jid_status` (`job_id`, `status`),
                KEY `idx_status_ctime` (`status`, `create_time`)
            ) AUTO_ID_CACHE 100 ;",
            META_DATABASE_NAME
        ))
        .await?;
    info!("Ensure ImportTaskStore table schema finished");
    Ok(())
}

const ERR_DUPLICATED_ENTRY: u16 = 1062;

fn parse_index_ids_csv(raw: &str) -> Vec<i64> {
    raw.split(',')
        .filter_map(|s| s.trim().parse::<i64>().ok())
        .collect()
}

pub async fn create_or_get_job(
    mut conn: Conn,
    job_meta: &ImportJobMeta,
) -> Result<ImportJobID, anyhow::Error> {
    debug_assert!(
        job_meta.job_id == 0,
        "job_meta.job_id should be 0 when creating a new job, tidb_task_id={}",
        job_meta.tidb_task_id
    );

    let backoff_max_retries = 3;
    let mut retries = 0;
    loop {
        let mut tx = conn.start_transaction(TxOpts::default()).await?;
        if let Some(created_job) =
            inner_get_import_job_meta(&mut tx, &job_meta.tidb_task_id).await?
        {
            return Ok(created_job.job_id);
        }

        let index_ids_str = itertools::join(job_meta.index_ids.iter(), ",");
        match tx
            .exec_iter(
                format!(
                    "INSERT INTO {}.tici_import_jobs
                (tidb_task_id, keyspace_id, table_id, index_ids, parameters, status, step)
                VALUES(:tidb_task_id,:keyspace_id,:table_id,:index_ids,:param,:status,:step)",
                    META_DATABASE_NAME
                ),
                params! {
                    "tidb_task_id" => &job_meta.tidb_task_id,
                    "keyspace_id" => job_meta.keyspace_id,
                    "table_id" => job_meta.table_id,
                    "index_ids" => index_ids_str,
                    "param" => "{}",
                    "status" => "init",
                    "step" => "init",
                },
            )
            .await
        {
            Err(err) => {
                if let mysql_async::Error::Server(ref e) = err {
                    if e.code == ERR_DUPLICATED_ENTRY {
                        // Another thread may create the same job concurrently.
                        // We can retry to get the job_id.
                        tx.rollback().await?;
                        if retries < backoff_max_retries {
                            retries += 1;
                            info!(
                                "Import - import job meet duplicated entry, retry. tidb_task_id={} retries={}",
                                job_meta.tidb_task_id, retries
                            );
                            continue;
                        } else {
                            // Expected to be success after the second try. If not, return error to
                            // avoid infinite loop.
                            return Err(anyhow::format_err!(
                                "Failed to create or get import job, tidb_task_id={} retries={} error={}",
                                job_meta.tidb_task_id,
                                retries,
                                err
                            ));
                        }
                    }
                }
                // Not ServerError, or the ServerError code is not duplicated entry
                tx.rollback().await?;
                return Err(anyhow::Error::from(err));
            }
            Ok(res) => {
                if let Some(job_id) = res.last_insert_id() {
                    tx.commit().await?; // TODO: Check the result of commit
                    info!(
                        "Import - import job added, job_id={} from_job_meta={:?}",
                        job_id, job_meta
                    );
                    return Ok(job_id);
                }
                // `last_insert_id` is not set, query the job again.
                match inner_get_import_job_meta(&mut tx, &job_meta.tidb_task_id).await {
                    Ok(Some(added_job_meta)) => {
                        tx.commit().await?; // TODO: Check the result of commit
                        info!(
                            "Import - import job added, job_id={} added_job_meta={:?}",
                            added_job_meta.job_id, added_job_meta
                        );
                        return Ok(added_job_meta.job_id);
                    }
                    Ok(None) => {
                        tx.rollback().await?;
                        return Err(anyhow::format_err!(
                            "Import - failed to get import job after create, tidb_task_id={}",
                            job_meta.tidb_task_id
                        ));
                    }
                    Err(e) => {
                        tx.rollback().await?;
                        return Err(e);
                    }
                }
            }
        }
    }
}

async fn inner_get_import_job_meta<'t>(
    tx: &mut Transaction<'t>,
    tidb_task_id: &TiDBImportTaskID,
) -> Result<Option<ImportJobMeta>, anyhow::Error> {
    // Find job_id for the task
    // ensure that the job status is not upload_finished or finished or failed
    let job_id_res: Vec<ImportJobMeta> = tx
        .exec_iter(
            format!(
                "SELECT job_id, keyspace_id, table_id, index_ids
                FROM {}.tici_import_jobs
                WHERE tidb_task_id=:tidb_task_id
                  AND status NOT IN ('finished', 'failed')
                ",
                META_DATABASE_NAME
            ),
            params! {
                "tidb_task_id" => tidb_task_id,
            },
        )
        .await?
        .map(|row| {
            let (job_id, keyspace_id, table_id, index_ids): (ImportJobID, u32, i64, String) =
                mysql_async::from_row(row);
            ImportJobMeta {
                tidb_task_id: tidb_task_id.clone(),
                keyspace_id,
                table_id,
                index_ids: parse_index_ids_csv(&index_ids),
                job_id,
            }
        })
        .await?;
    if job_id_res.len() > 1 {
        Err(anyhow::format_err!(
            "Failed to find job_meta for task, tidb_task_id={:?}, found={:?}",
            tidb_task_id,
            job_id_res
        ))
    } else if job_id_res.is_empty() {
        Ok(None)
    } else {
        Ok(Some(job_id_res[0].clone()))
    }
}

pub async fn finish_job_upload(
    mut conn: Conn,
    tidb_task_id: &TiDBImportTaskID,
) -> Result<(), anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let _ = tx
        .exec_iter(
            format!(
                "UPDATE {}.tici_import_jobs
                SET step='upload_finished', update_time=CURRENT_TIMESTAMP(6)
                WHERE tidb_task_id=:tidb_task_id",
                META_DATABASE_NAME
            ),
            params! {
                "tidb_task_id" => tidb_task_id,
            },
        )
        .await?;
    tx.commit().await?; // TODO: Check the result of commit
    info!(
        "Import - import job finish upload, tidb_task_id={:?}",
        tidb_task_id
    );
    Ok(())
}

pub async fn finish_job<'a>(
    mut conn: Conn,
    job_ids: &'a [ImportJobID],
) -> Result<(), anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    // TODO: support "status=failed" and fill the "error_message"

    if job_ids.is_empty() {
        tx.rollback().await?;
        return Err(anyhow::format_err!("No job IDs provided to finish."));
    }
    let job_summaries = collect_job_summaries(&mut tx, job_ids).await?;
    let mut params_list = Vec::with_capacity(job_ids.len());
    for job_id in job_ids {
        let summary = job_summaries.get(job_id).cloned().unwrap_or_default();
        params_list.push(params! {
            "job_id" => *job_id,
            "summary" => serde_json::to_string(&summary)?,
        });
    }
    tx.exec_batch(
        format!(
            "UPDATE {}.tici_import_jobs
                SET status='finished', step='finish', summary=:summary, update_time=CURRENT_TIMESTAMP(6), end_time=CURRENT_TIMESTAMP(6)
                WHERE job_id=:job_id",
            META_DATABASE_NAME
        ),
        params_list,
    )
    .await?;

    tx.commit().await?; // TODO: Check the result of commit
    info!("Import - import job finished, job_ids={:?}", job_ids);
    Ok(())
}

pub async fn fail_job<'a>(mut conn: Conn, job_ids: &'a [ImportJobID]) -> Result<(), anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    if job_ids.is_empty() {
        tx.rollback().await?;
        return Err(anyhow::format_err!("No job IDs provided to fail."));
    }

    tx.exec_batch(
        format!(
            "UPDATE {}.tici_import_jobs
                SET status='failed', step='finish', update_time=CURRENT_TIMESTAMP(6), end_time=CURRENT_TIMESTAMP(6)
                WHERE job_id=:job_id AND status != 'finished'",
            META_DATABASE_NAME
        ),
        job_ids.iter().map(|job_id| {
            params! {
                "job_id" => *job_id,
            }
        }),
    )
    .await?;

    tx.commit().await?;
    info!("Import - import job failed, job_ids={:?}", job_ids);
    Ok(())
}

pub async fn fail_jobs_and_reclaim_running_tasks<'a>(
    mut conn: Conn,
    job_ids: &'a [ImportJobID],
) -> Result<Vec<ImportTaskID>, anyhow::Error> {
    if job_ids.is_empty() {
        return Err(anyhow::format_err!(
            "No job IDs provided to fail and reclaim running tasks."
        ));
    }

    let placeholders = vec!["?"; job_ids.len()].join(",");
    let params = || {
        Params::Positional(
            job_ids
                .iter()
                .map(|job_id| (*job_id).into())
                .collect::<Vec<_>>(),
        )
    };
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let running_task_ids: Vec<ImportTaskID> = tx
        .exec_iter(
            format!(
                "SELECT task_id
                    FROM {}.tici_import_jobs_task
                    WHERE job_id IN ({})
                      AND status='running'",
                META_DATABASE_NAME, placeholders
            ),
            params(),
        )
        .await?
        .map(mysql_async::from_row::<ImportTaskID>)
        .await?;
    tx.exec_batch(
        format!(
            "UPDATE {}.tici_import_jobs
                SET status='failed', step='finish', update_time=CURRENT_TIMESTAMP(6), end_time=CURRENT_TIMESTAMP(6)
                WHERE job_id=:job_id AND status != 'finished'",
            META_DATABASE_NAME
        ),
        job_ids.iter().map(|job_id| {
            params! {
                "job_id" => *job_id,
            }
        }),
    )
    .await?;
    if !running_task_ids.is_empty() {
        let _ = tx
            .exec_iter(
                format!(
                    "UPDATE {}.tici_import_jobs_task
                        SET status='init',
                            owner_addr=NULL,
                            update_time=CURRENT_TIMESTAMP(6),
                            start_time=NULL,
                            end_time=NULL,
                            summary=NULL,
                            error_message=NULL
                        WHERE job_id IN ({})
                          AND status='running'",
                    META_DATABASE_NAME, placeholders
                ),
                params(),
            )
            .await?;
    }
    tx.commit().await?;
    info!(
        "Import - import jobs failed and running tasks reclaimed, job_ids={:?}, running_task_ids={:?}",
        job_ids, running_task_ids
    );
    Ok(running_task_ids)
}

async fn collect_job_summaries<'a>(
    tx: &mut Transaction<'a>,
    job_ids: &[ImportJobID],
) -> Result<HashMap<ImportJobID, ImportJobSummary>, anyhow::Error> {
    if job_ids.is_empty() {
        return Ok(HashMap::new());
    }
    let placeholders = vec!["?"; job_ids.len()].join(",");
    let sql = format!(
        "SELECT job_id, summary
            FROM {}.tici_import_jobs_task
            WHERE job_id IN ({}) AND status='finish'",
        META_DATABASE_NAME, placeholders
    );
    let params = Params::Positional(
        job_ids
            .iter()
            .map(|job_id| (*job_id).into())
            .collect::<Vec<_>>(),
    );
    let rows: Vec<(ImportJobID, Option<String>)> = tx
        .exec_iter(sql, params)
        .await?
        .map(|row| {
            let (job_id, summary): (ImportJobID, Option<String>) = mysql_async::from_row(row);
            (job_id, summary)
        })
        .await?;
    let mut summaries = HashMap::new();
    for (job_id, summary_json) in rows {
        let task_summary = summary_json
            .as_deref()
            .and_then(|raw| serde_json::from_str::<ImportTaskSummary>(raw).ok())
            .unwrap_or_default();
        let entry = summaries
            .entry(job_id)
            .or_insert_with(ImportJobSummary::default);
        entry.task_count = entry.task_count.saturating_add(1);
        entry.rows_processed = entry
            .rows_processed
            .saturating_add(task_summary.rows_processed);
        entry.bytes_processed = entry
            .bytes_processed
            .saturating_add(task_summary.bytes_processed);
        // The size calculation for deduplicated files is not required here.
        // Because tici-worker really need to read the file multiple times for
        // multi-index IMPORT jobs now.
        entry.input_file_size = entry
            .input_file_size
            .saturating_add(task_summary.input_file_size);
    }
    Ok(summaries)
}

#[derive(Debug)]
pub struct ImportJobStat {
    pub job_id: ImportJobID,
    pub keyspace_id: u32,
    pub table_id: i64,
    pub index_ids: Vec<i64>,
    pub task_count: i64,
    pub running_task_count: i64,
    pub finished_task_count: i64,
    pub failed_task_count: i64,
}
pub async fn get_jobs_should_finish(
    mut conn: Conn,
) -> Result<Vec<ImportJobStat>, mysql_async::Error> {
    let jobs = conn
        .query_iter(format!(
            "SELECT
                  j.job_id,
                  j.keyspace_id,
                  j.table_id,
                  j.index_ids,
                  count(1) AS tsk_cnt,
                  SUM(IF(t.status='running', 1, 0)) AS tsk_running_cnt,
                  SUM(IF(t.status='finish', 1, 0)) AS tsk_finish_cnt,
                  SUM(IF(t.status='failed', 1, 0)) AS tsk_failed_cnt
                FROM {0}.tici_import_jobs as j, {0}.tici_import_jobs_task as t
                WHERE j.job_id=t.job_id
                  AND j.step='upload_finished' AND j.status='running'
                GROUP BY j.job_id, j.keyspace_id, j.table_id, j.index_ids",
            META_DATABASE_NAME
        ))
        .await?
        .map(|row| {
            let (
                job_id,
                keyspace_id,
                table_id,
                index_ids,
                tsk_cnt,
                tsk_running_cnt,
                tsk_finish_cnt,
                tsk_failed_cnt,
            ): (ImportJobID, u32, i64, String, i64, i64, i64, i64) = mysql_async::from_row(row);
            ImportJobStat {
                job_id,
                keyspace_id,
                table_id,
                index_ids: parse_index_ids_csv(&index_ids),
                task_count: tsk_cnt,
                running_task_count: tsk_running_cnt,
                finished_task_count: tsk_finish_cnt,
                failed_task_count: tsk_failed_cnt,
            }
        })
        .await?;

    Ok(jobs)
}

pub async fn get_running_tasks_count(mut conn: Conn) -> Result<usize, mysql_async::Error> {
    let running_task_count = conn
        .query_first::<u64, _>(format!(
            "SELECT COUNT(*)
                FROM {}.tici_import_jobs_task
                WHERE status='running'",
            META_DATABASE_NAME
        ))
        .await?
        .unwrap_or_default();
    Ok(running_task_count as usize)
}

pub async fn get_running_tasks(mut conn: Conn) -> Result<Vec<ImportTask>, mysql_async::Error> {
    let rows: Vec<(ImportTaskID, ImportJobID, String, u64, String)> = conn
        .query_iter(format!(
            "SELECT t.task_id, t.job_id, t.owner_addr, t.owner_seq, t.parameters
                FROM {}.tici_import_jobs_task AS t
                INNER JOIN {}.tici_import_jobs AS j
                    ON j.job_id = t.job_id
                WHERE t.status='running'
                  AND t.owner_addr IS NOT NULL
                  AND t.owner_addr != ''
                  AND j.status NOT IN ('finished', 'failed')",
            META_DATABASE_NAME, META_DATABASE_NAME
        ))
        .await?
        .map(mysql_async::from_row)
        .await?;

    let mut tasks = Vec::with_capacity(rows.len());
    for (task_id, job_id, owner_addr, owner_seq, params_str) in rows {
        let parameters = serde_json::from_str(&params_str)
            .map_err(|err| mysql_async::Error::Other(err.to_string().into()))?;
        tasks.push(ImportTask {
            task_id,
            job_id,
            owner_addr,
            owner_seq,
            parameters,
        });
    }

    Ok(tasks)
}

pub async fn get_index_progress(
    mut conn: Conn,
    keyspace_id: u32,
    table_id: i64,
    index_id: i64,
) -> Result<Option<super::metastore::ImportIndexProgress>, mysql_async::Error> {
    let progress_rows: Vec<super::metastore::ImportIndexProgress> = conn
        .exec_iter(
            format!(
                r"SELECT
                  j.job_id,
                  j.status,
                  COUNT(t.task_id) AS task_count,
                  COALESCE(SUM(IF(t.status='running', 1, 0)), 0) AS task_running_cnt,
                  COALESCE(SUM(IF(t.status='finish', 1, 0)), 0) AS task_finish_cnt,
                  COALESCE(SUM(IF(t.status='failed', 1, 0)), 0) AS task_failed_cnt
                FROM {0}.tici_import_jobs AS j
                LEFT JOIN {0}.tici_import_jobs_task AS t
                  ON j.job_id = t.job_id
                 AND CAST(JSON_UNQUOTE(JSON_EXTRACT(t.parameters, '$.index_id')) AS SIGNED) = :index_id
                WHERE j.keyspace_id = :keyspace_id
                  AND j.table_id = :table_id
                  AND FIND_IN_SET(:index_id, j.index_ids) > 0
                GROUP BY j.job_id, j.status
                ORDER BY
                  CASE WHEN COUNT(t.task_id) > 0 THEN 1 ELSE 0 END DESC,
                  j.job_id DESC
                LIMIT 1",
                META_DATABASE_NAME
            ),
            params! {
                "keyspace_id" => keyspace_id,
                "table_id" => table_id,
                "index_id" => index_id,
            },
        )
        .await?
        .map(|row| {
            let (
                job_id,
                status,
                task_count,
                running_task_count,
                finished_task_count,
                failed_task_count,
            ): (ImportJobID, String, i64, i64, i64, i64) = mysql_async::from_row(row);
            super::metastore::ImportIndexProgress {
                job_id,
                job_status: status,
                task_count,
                running_task_count,
                finished_task_count,
                failed_task_count,
            }
        })
        .await?;

    Ok(progress_rows.into_iter().next())
}

pub async fn get_active_jobs(mut conn: Conn) -> Result<Vec<ImportJobMeta>, mysql_async::Error> {
    conn.query_iter(format!(
        r"SELECT job_id, tidb_task_id, keyspace_id, table_id, index_ids
            FROM {}.tici_import_jobs
            WHERE status NOT IN ('finished', 'failed')",
        META_DATABASE_NAME
    ))
    .await?
    .map(|row| {
        let (job_id, tidb_task_id, keyspace_id, table_id, index_ids): (
            ImportJobID,
            TiDBImportTaskID,
            u32,
            i64,
            String,
        ) = mysql_async::from_row(row);
        ImportJobMeta {
            tidb_task_id,
            keyspace_id,
            table_id,
            index_ids: parse_index_ids_csv(&index_ids),
            job_id,
        }
    })
    .await
}

pub async fn add_task(
    mut conn: Conn,
    tidb_task_id: &TiDBImportTaskID,
    params: ImportTaskParams,
) -> Result<(), anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    let job_meta = match inner_get_import_job_meta(&mut tx, tidb_task_id).await {
        Ok(Some(j)) => j,
        Ok(None) => bail!(
            "Import - Failed to find job_meta for add_task, tidb_task_id={}",
            tidb_task_id
        ),
        Err(e) => bail!(
            "Import - Failed to find job_meta for add_task, tidb_task_id={}, error={}",
            tidb_task_id,
            e
        ),
    };

    let params = params.split_with_job_info(
        job_meta.keyspace_id,
        job_meta.table_id,
        &job_meta.index_ids,
        job_meta.job_id,
    )?;

    tx.exec_batch(
        format!(
            "INSERT INTO {}.tici_import_jobs_task
                (job_id, parameters, status)
                VALUES(:job_id,:param,:status)",
            META_DATABASE_NAME
        ),
        params.iter().map(|p| {
            params! {
                "job_id" => job_meta.job_id,
                "param" => serde_json::to_string(p).unwrap(),
                "status" => "init",
            }
        }),
    )
    .await?;

    tx.commit().await?; // TODO: Check the result of commit
    info!("Import - import job task added, job_meta={:?}", job_meta);
    Ok(())
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct InactiveImportTask {
    pub task_id: ImportTaskID,
    pub job_id: ImportJobID,
    pub update_time: String,
    pub elapsed_secs: u64,
}

pub async fn inner_inactive_tasks<'a>(
    tx: &mut Transaction<'a>,
    inactive_secs: u64,
) -> Result<Vec<InactiveImportTask>, Error> {
    let sql = format!(
        r"SELECT
                task_id, job_id,
                update_time, TIMESTAMPDIFF(SECOND, update_time, NOW()) AS elapsed
              FROM {}.tici_import_jobs_task
              WHERE status = 'running'
                AND update_time IS NOT NULL
                AND TIMESTAMPDIFF(SECOND, update_time, NOW()) > :inactive_secs;",
        META_DATABASE_NAME
    );
    let tasks = tx
        .exec_iter(
            sql,
            params! {
                "inactive_secs" => inactive_secs,
            },
        )
        .await?
        .map(|row| {
            let (task_id, job_id, update_time, elapsed): (ImportTaskID, ImportJobID, String, u64) =
                mysql_async::from_row(row);
            InactiveImportTask {
                task_id,
                job_id,
                update_time,
                elapsed_secs: elapsed,
            }
        })
        .await?;

    Ok(tasks)
}

pub async fn inactive_tasks(
    mut conn: Conn,
    inactive_secs: u64,
) -> Result<Vec<InactiveImportTask>, Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    inner_inactive_tasks(&mut tx, inactive_secs).await
}

pub async fn reset_inactive_tasks(
    mut conn: Conn,
    inactive_secs: u64,
    limit: usize,
) -> Result<Vec<InactiveImportTask>, Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let inactive_tasks = inner_inactive_tasks(&mut tx, inactive_secs).await?;
    // Avoid resetting too many tasks at once that lead to txn failure. Split the
    // tasks by the given limit.
    let (reset_tasks, left_tasks) = if inactive_tasks.len() >= limit {
        inactive_tasks.split_at(limit)
    } else {
        (&inactive_tasks[..], &[] as &[InactiveImportTask])
    };
    // for each inactive task, reset its status to 'init'
    for task in reset_tasks {
        let _ = tx
            .exec_iter(
                format!(
                    "UPDATE {}.tici_import_jobs_task
                SET status='init',
                    owner_addr=NULL,
                    update_time=CURRENT_TIMESTAMP(6),
                    start_time=NULL,
                    end_time=NULL,
                    summary=NULL,
                    error_message=NULL
                WHERE task_id=:task_id",
                    META_DATABASE_NAME,
                ),
                params! {
                    "task_id" => task.task_id,
                },
            )
            .await?;
    }
    tx.commit().await?; // TODO: Check the result of commit

    info!(
        "Import - reset inactive import task to init, inactive_secs={} limit={} reset_tasks={:?} left_tasks={:?}",
        inactive_secs, limit, reset_tasks, left_tasks
    );
    Ok(reset_tasks.to_vec())
}

async fn query_external_terminal_task_keys_from_table<'a>(
    tx: &mut Transaction<'a>,
    table_name: &str,
    task_keys: &[String],
) -> Result<HashSet<String>, Error> {
    if task_keys.is_empty() {
        return Ok(HashSet::new());
    }

    let placeholders = vec!["?"; task_keys.len()].join(",");
    let sql = format!(
        "SELECT task_key
            FROM {table_name}
            WHERE task_key IN ({placeholders})
              AND LOWER(state) IN ('failed', 'canceled', 'cancelled', 'reverted')"
    );
    let params = Params::Positional(
        task_keys
            .iter()
            .map(|task_key| task_key.clone().into())
            .collect::<Vec<_>>(),
    );
    let rows: Vec<String> = tx
        .exec_iter(sql, params)
        .await?
        .map(mysql_async::from_row)
        .await?;

    Ok(rows.into_iter().collect())
}

pub async fn get_external_terminal_task_keys(
    mut conn: Conn,
    task_keys: &[String],
) -> Result<HashSet<String>, anyhow::Error> {
    if task_keys.is_empty() {
        return Ok(HashSet::new());
    }

    let unique_task_keys = task_keys
        .iter()
        .cloned()
        .collect::<HashSet<_>>()
        .into_iter()
        .collect::<Vec<_>>();
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let mut terminal_task_keys = query_external_terminal_task_keys_from_table(
        &mut tx,
        "mysql.tidb_global_task",
        &unique_task_keys,
    )
    .await?;
    let unresolved_task_keys = unique_task_keys
        .iter()
        .filter(|task_key| !terminal_task_keys.contains(*task_key))
        .cloned()
        .collect::<Vec<_>>();
    if !unresolved_task_keys.is_empty() {
        match query_external_terminal_task_keys_from_table(
            &mut tx,
            "mysql.tidb_global_task_history",
            &unresolved_task_keys,
        )
        .await
        {
            Ok(history_task_keys) => terminal_task_keys.extend(history_task_keys),
            Err(e) => {
                warn!(
                    "Import - failed to query external task history terminal states, use current table results, unresolved_task_keys={:?}, err={:?}",
                    unresolved_task_keys, e
                );
            }
        }
    }
    tx.rollback().await?;
    Ok(terminal_task_keys)
}

pub async fn next_task(mut conn: Conn, owner_addr: &str) -> Result<Option<ImportTask>, Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let task_ids: Vec<ImportTask> = tx
        .query_iter(format!(
            r"SELECT t.task_id, t.job_id, t.owner_seq, t.parameters
                FROM {}.tici_import_jobs_task AS t
                INNER JOIN {}.tici_import_jobs AS j
                    ON j.job_id = t.job_id
                WHERE t.status = 'init'
                  AND j.status NOT IN ('finished', 'failed')
                ORDER BY t.create_time ASC
                LIMIT 1;",
            META_DATABASE_NAME, META_DATABASE_NAME
        ))
        .await?
        .map(|row| {
            let (task_id, job_id, owner_seq, params_str): (ImportTaskID, ImportJobID, u64, String) =
                mysql_async::from_row(row);
            ImportTask {
                task_id,
                job_id,
                owner_addr: String::new(),
                owner_seq,
                parameters: serde_json::from_str(&params_str).unwrap(),
            }
        })
        .await?;

    if task_ids.is_empty() {
        tx.rollback().await?;
        return Ok(None);
    }

    let task = task_ids[0].clone();
    // Try to update the job status
    let update_job = tx
        .exec_iter(
            format!(
                "UPDATE {}.tici_import_jobs
                SET status='running', start_time=CURRENT_TIMESTAMP(6), update_time=CURRENT_TIMESTAMP(6)
                WHERE job_id=:job_id and status='init'",
                META_DATABASE_NAME
            ),
            params! {"job_id" => task.job_id},
        )
        .await?;
    if update_job.affected_rows() != 1 {
        debug!(
            "Failed to update job status to processing, affected_rows={}, task={:?}",
            update_job.affected_rows(),
            task
        );
    }
    // Update the task status to running
    let update_task = tx
        .exec_iter(
            format!(
                "UPDATE {}.tici_import_jobs_task AS t
                    INNER JOIN {}.tici_import_jobs AS j
                        ON j.job_id = t.job_id
                    SET t.status='running',
                        t.start_time=CURRENT_TIMESTAMP(6),
                        t.update_time=CURRENT_TIMESTAMP(6),
                        t.owner_addr=:owner_addr,
                        t.owner_seq=t.owner_seq+1
                    WHERE t.task_id=:task_id
                      AND t.status='init'
                      AND j.status NOT IN ('finished', 'failed')",
                META_DATABASE_NAME, META_DATABASE_NAME
            ),
            params! {
                "task_id" => task.task_id,
                "owner_addr" => owner_addr,
            },
        )
        .await?;
    if update_task.affected_rows() != 1 {
        tx.rollback().await?;
        return Ok(None);
    }
    tx.commit().await?; // TODO: Check the result of commit
    let mut task = task;
    task.owner_addr = owner_addr.to_string();
    task.owner_seq += 1;
    info!(
        "Import - import job task to process={:?}, set status to running",
        task
    );

    Ok(Some(task))
}

pub async fn reclaim_task(
    mut conn: Conn,
    task_id: ImportTaskID,
    owner_addr: &str,
) -> Result<bool, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let affected_rows = reset_running_import_tasks(
        &mut tx,
        "task_id=:task_id AND owner_addr=:owner_addr",
        params! {
            "task_id" => task_id,
            "owner_addr" => owner_addr,
        },
    )
    .await?;
    tx.commit().await?;
    Ok(affected_rows == 1)
}

async fn reset_running_import_tasks<'t>(
    tx: &mut Transaction<'t>,
    condition: &str,
    params: Params,
) -> Result<u64, Error> {
    let update = tx
        .exec_iter(
            format!(
                "UPDATE {}.tici_import_jobs_task
                    SET status='init',
                        owner_addr=NULL,
                        update_time=CURRENT_TIMESTAMP(6),
                        start_time=NULL,
                        end_time=NULL,
                        summary=NULL,
                        error_message=NULL
                    WHERE status='running' AND {condition}",
                META_DATABASE_NAME
            ),
            params,
        )
        .await?;
    Ok(update.affected_rows())
}

pub async fn reclaim_tasks_by_worker(
    mut conn: Conn,
    owner_addr: &str,
) -> Result<Vec<ImportTaskID>, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let task_ids: Vec<ImportTaskID> = tx
        .exec_iter(
            format!(
                "SELECT task_id FROM {}.tici_import_jobs_task
                    WHERE status='running' AND owner_addr=:owner_addr",
                META_DATABASE_NAME
            ),
            params! {
                "owner_addr" => owner_addr,
            },
        )
        .await?
        .map(mysql_async::from_row::<ImportTaskID>)
        .await?;
    if task_ids.is_empty() {
        tx.rollback().await?;
        return Ok(vec![]);
    }
    let _ = reset_running_import_tasks(
        &mut tx,
        "owner_addr=:owner_addr",
        params! {
            "owner_addr" => owner_addr,
        },
    )
    .await?;
    tx.commit().await?;
    Ok(task_ids)
}

pub async fn task_build_success(
    mut conn: Conn,
    task_id: ImportTaskID,
    owner_addr: &str,
    owner_seq: u64,
    summary: ImportTaskSummary,
) -> Result<bool, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;
    let summary = serde_json::to_string(&summary)?;

    let update = tx
        .exec_iter(
            format!(
                "UPDATE {}.tici_import_jobs_task AS t
                    INNER JOIN {}.tici_import_jobs AS j
                        ON j.job_id = t.job_id
                SET t.status='finish',
                    t.summary=:summary,
                    t.end_time=CURRENT_TIMESTAMP(6),
                    t.update_time=CURRENT_TIMESTAMP(6)
                WHERE t.task_id=:task_id
                  AND t.status='running'
                  AND t.owner_addr=:owner_addr
                  AND t.owner_seq=:owner_seq
                  AND j.status NOT IN ('finished', 'failed')",
                META_DATABASE_NAME, META_DATABASE_NAME
            ),
            params! {
                "task_id" => task_id,
                "owner_addr" => owner_addr,
                "owner_seq" => owner_seq,
                "summary" => summary,
            },
        )
        .await?;
    let affected_rows = update.affected_rows();
    tx.commit().await?; // TODO: Check the result of commit
    if affected_rows != 1 {
        info!(
            "Import - skip stale import task finish, task_id={} owner_addr={} owner_seq={}",
            task_id, owner_addr, owner_seq
        );
        return Ok(false);
    }
    info!(
        "Import - import job task finish build, task_id={:?}",
        task_id
    );

    Ok(true)
}

pub async fn task_build_failed(
    mut conn: Conn,
    task_id: ImportTaskID,
    owner_addr: &str,
    owner_seq: u64,
    err_msg: &str,
) -> Result<bool, anyhow::Error> {
    let mut tx = conn.start_transaction(TxOpts::default()).await?;

    // update the status and fill the "error_message" if the task meet
    // unrecoverable error. For example, the input file is broken.
    let update = tx
        .exec_iter(
            format!(
                "UPDATE {}.tici_import_jobs_task AS t
                    INNER JOIN {}.tici_import_jobs AS j
                        ON j.job_id = t.job_id
                SET t.status='failed',
                    t.error_message=:error_message,
                    t.end_time=CURRENT_TIMESTAMP(6),
                    t.update_time=CURRENT_TIMESTAMP(6)
                WHERE t.task_id=:task_id
                  AND t.status='running'
                  AND t.owner_addr=:owner_addr
                  AND t.owner_seq=:owner_seq
                  AND j.status NOT IN ('finished', 'failed')",
                META_DATABASE_NAME, META_DATABASE_NAME
            ),
            params! {
                "task_id" => task_id,
                "owner_addr" => owner_addr,
                "owner_seq" => owner_seq,
                "error_message" => err_msg,
            },
        )
        .await?;
    let affected_rows = update.affected_rows();
    tx.commit().await?; // TODO: Check the result of commit
    if affected_rows != 1 {
        info!(
            "Import - skip stale import task fail, task_id={} owner_addr={} owner_seq={}",
            task_id, owner_addr, owner_seq
        );
        return Ok(false);
    }
    info!(
        "Import - import job task build failed, task_id={} err={}",
        task_id, err_msg
    );

    Ok(true)
}