use std::sync::Arc;

use tici::{
    HandleRange, ShardInfo, ShardInfoUpdate,
    config::make_default_test_meta_config,
    connection_pool::TiDBConnectionPool,
    meta::{Bootstrap, MetaStore, MetaStoreImplTiDB},
};
use tici_common::keyspace::NULLSPACEID;
use tici_import::{ImportJobMeta, ImportTaskParams, ImportTaskSummary};
use tici_service_discovery::election::leadership::MockLeadership;
use tici_test_utils::test_utils::*;

#[tokio::test]
async fn test_shard_store() {
    let config = make_default_test_meta_config(false);
    let s1 = {
        Bootstrap::new_with_config(&config)
            .clear_all_meta(true)
            .await
            .unwrap();

        let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
        let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());
        let mut s = ShardInfo::new_for_test(HandleRange::new(vec![b'A'], vec![b'B']));
        assert_eq!(meta_store.init_shard(&s).await.unwrap(), 1);
        assert!(meta_store.init_shard(&s).await.is_err());
        s.shard_id = 1;
        let mut s1 = s.clone();
        s1.seq += 1;
        let mut update = ShardInfoUpdate::from(&s1);
        update.progress = Some(s1.progress.clone());
        update.manifest = Some(s1.manifest.clone());
        update.seq = Some(s1.seq);
        meta_store
            .update_shard("commit_append", &s, update)
            .await
            .unwrap();
        s1
    };

    {
        let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
        let meta_store2 = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());
        let v = meta_store2.restore_shards().await.unwrap();
        assert_eq!(v.len(), 1);
        assert_eq!(v[0], s1); // check the shard info restored
    }
}

#[tokio::test]
async fn test_import_store() {
    let config = load_meta_config("../../ci/meta-template.toml").unwrap();
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let tidb_task_id = "tidb_task_1".to_string();
    let job_meta = ImportJobMeta {
        tidb_task_id: tidb_task_id.clone(),
        keyspace_id: NULLSPACEID,
        table_id: 1000,
        index_ids: vec![1],
        job_id: 0, // job_id will be assigned by meta_store
    };

    // Submit an import job
    meta_store
        .create_or_get_import_job(&job_meta)
        .await
        .unwrap();

    // Submit tasks for the import job
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "c".to_string(),
                "s3://bucket/path/to/file_a_c".to_string(),
            ),
        )
        .await
        .unwrap();

    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "c".to_string(),
                "z".to_string(),
                "s3://bucket/path/to/file_c_z".to_string(),
            ),
        )
        .await
        .unwrap();

    // All the tasks are submitted, now TiCI should mark the job as upload
    // finished. And TiCI should wait for all tasks to be built before
    // marking the job as finished.
    meta_store
        .finish_import_job_upload(&tidb_task_id)
        .await
        .unwrap();

    // Fetch the import tasks and mock dispatching them to TiCI worker nodes
    let task0 = meta_store.next_import_task("worker-a").await.unwrap();
    assert!(task0.is_some());
    let task1 = meta_store.next_import_task("worker-b").await.unwrap();
    assert!(task1.is_some());
    // We only have two tasks, so the next one should be None
    assert!(
        meta_store
            .next_import_task("worker-c")
            .await
            .unwrap()
            .is_none()
    );

    // Mock the tasks are finished by the TiCI worker nodes
    let task0 = task0.unwrap();
    meta_store
        .import_task_build_success(
            task0.task_id,
            &task0.owner_addr,
            task0.owner_seq,
            ImportTaskSummary::default(),
        )
        .await
        .unwrap();
    let task1 = task1.unwrap();
    meta_store
        .import_task_build_success(
            task1.task_id,
            &task1.owner_addr,
            task1.owner_seq,
            ImportTaskSummary::default(),
        )
        .await
        .unwrap();

    let jobs_stat = meta_store.get_jobs_should_finish().await.unwrap();
    assert_eq!(jobs_stat.len(), 1);
    assert_eq!(jobs_stat[0].task_count, 2);
    assert_eq!(jobs_stat[0].finished_task_count, 2);

    let jobs_finished = jobs_stat
        .into_iter()
        .filter(|job| job.finished_task_count == job.task_count)
        .map(|job| job.job_id)
        .collect::<Vec<_>>();

    // meta_store mark the import job as finished
    meta_store.finish_import_job(&jobs_finished).await.unwrap();
}

#[tokio::test]
async fn test_import_store_multi_index_uses_requested_index_only() -> anyhow::Result<()> {
    let config = load_meta_config("../../ci/meta-template.toml")?;
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await?;

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let tidb_task_id = "tidb_task_multi_index_requested".to_string();
    meta_store
        .create_or_get_import_job(&ImportJobMeta {
            tidb_task_id: tidb_task_id.clone(),
            keyspace_id: NULLSPACEID,
            table_id: 1000,
            index_ids: vec![11, 22],
            job_id: 0,
        })
        .await?;

    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "c".to_string(),
                "s3://bucket/path/to/file_a_c".to_string(),
            )
            .with_index_id(11),
        )
        .await?;

    let task = meta_store
        .next_import_task("worker-a")
        .await?
        .ok_or_else(|| anyhow::anyhow!("expected one import task"))?;
    assert_eq!(task.parameters.location, "s3://bucket/path/to/file_a_c");
    assert_eq!(task.parameters.index_id, 11);
    assert!(meta_store.next_import_task("worker-b").await?.is_none());

    Ok(())
}

#[tokio::test]
async fn test_import_store_rejects_invalid_and_unsafe_legacy_index_id() -> anyhow::Result<()> {
    let config = load_meta_config("../../ci/meta-template.toml")?;
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await?;

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    for (tidb_task_id, index_id, expected_error) in [
        (
            "tidb_task_multi_index_invalid",
            33,
            "is not in import job index_ids",
        ),
        (
            "tidb_task_multi_index_legacy",
            0,
            "index_id is required for multi-index import job",
        ),
    ] {
        let tidb_task_id = tidb_task_id.to_string();
        meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: tidb_task_id.clone(),
                keyspace_id: NULLSPACEID,
                table_id: 1000,
                index_ids: vec![11, 22],
                job_id: 0,
            })
            .await?;

        let err = match meta_store
            .add_import_task(
                &tidb_task_id,
                ImportTaskParams::new(
                    "a".to_string(),
                    "c".to_string(),
                    "s3://bucket/path/to/file_a_c".to_string(),
                )
                .with_index_id(index_id),
            )
            .await
        {
            Ok(()) => anyhow::bail!("expected add_import_task to reject index_id={index_id}"),
            Err(err) => err,
        };
        assert!(
            err.to_string().contains(expected_error),
            "err={err:#}, expected={expected_error}"
        );
    }

    assert!(meta_store.next_import_task("worker-a").await?.is_none());
    Ok(())
}

#[tokio::test]
async fn test_get_running_import_task_count() {
    let config = load_meta_config("../../ci/meta-template.toml").unwrap();
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let tidb_task_id = "tidb_task_running_count".to_string();
    let job_meta = ImportJobMeta {
        tidb_task_id: tidb_task_id.clone(),
        keyspace_id: NULLSPACEID,
        table_id: 3000,
        index_ids: vec![1],
        job_id: 0,
    };

    meta_store
        .create_or_get_import_job(&job_meta)
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "c".to_string(),
                "s3://bucket/path/to/file_a_c".to_string(),
            ),
        )
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "c".to_string(),
                "z".to_string(),
                "s3://bucket/path/to/file_c_z".to_string(),
            ),
        )
        .await
        .unwrap();

    assert_eq!(meta_store.get_running_import_task_count().await.unwrap(), 0);

    let task0 = meta_store.next_import_task("worker-a").await.unwrap();
    assert!(task0.is_some());
    assert_eq!(meta_store.get_running_import_task_count().await.unwrap(), 1);

    // The running-task count should not depend on whether the job already
    // reaches the upload_finished step.
    meta_store
        .finish_import_job_upload(&tidb_task_id)
        .await
        .unwrap();

    let task1 = meta_store.next_import_task("worker-b").await.unwrap();
    assert!(task1.is_some());
    assert_eq!(meta_store.get_running_import_task_count().await.unwrap(), 2);

    let task0 = task0.unwrap();
    meta_store
        .import_task_build_success(
            task0.task_id,
            &task0.owner_addr,
            task0.owner_seq,
            ImportTaskSummary::default(),
        )
        .await
        .unwrap();
    assert_eq!(meta_store.get_running_import_task_count().await.unwrap(), 1);

    let task1 = task1.unwrap();
    meta_store
        .import_task_build_success(
            task1.task_id,
            &task1.owner_addr,
            task1.owner_seq,
            ImportTaskSummary::default(),
        )
        .await
        .unwrap();
    assert_eq!(meta_store.get_running_import_task_count().await.unwrap(), 0);
}

#[tokio::test]
async fn test_get_import_index_progress() {
    let config = make_default_test_meta_config(false);
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let table_id = 2000;
    let index_id = 10;
    let tidb_task_id = "tidb_task_progress_1".to_string();
    let job_meta = ImportJobMeta {
        tidb_task_id: tidb_task_id.clone(),
        keyspace_id: NULLSPACEID,
        table_id,
        index_ids: vec![index_id],
        job_id: 0,
    };
    meta_store
        .create_or_get_import_job(&job_meta)
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "c".to_string(),
                "s3://bucket/path/to/file_a_c".to_string(),
            ),
        )
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "c".to_string(),
                "z".to_string(),
                "s3://bucket/path/to/file_c_z".to_string(),
            ),
        )
        .await
        .unwrap();
    meta_store
        .finish_import_job_upload(&tidb_task_id)
        .await
        .unwrap();

    let progress = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_id)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(progress.task_count, 2);
    assert_eq!(progress.finished_task_count, 0);
    assert_eq!(progress.failed_task_count, 0);

    let task0 = meta_store.next_import_task("worker-a").await.unwrap();
    assert!(task0.is_some());
    let task1 = meta_store.next_import_task("worker-b").await.unwrap();
    assert!(task1.is_some());

    let task0 = task0.unwrap();
    meta_store
        .import_task_build_success(
            task0.task_id,
            &task0.owner_addr,
            task0.owner_seq,
            ImportTaskSummary::default(),
        )
        .await
        .unwrap();
    let task1 = task1.unwrap();
    meta_store
        .import_task_build_success(
            task1.task_id,
            &task1.owner_addr,
            task1.owner_seq,
            ImportTaskSummary::default(),
        )
        .await
        .unwrap();

    let progress = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_id)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(progress.task_count, 2);
    assert_eq!(progress.finished_task_count, 2);
    assert_eq!(progress.failed_task_count, 0);

    meta_store
        .finish_import_job(&[progress.job_id])
        .await
        .unwrap();

    let progress = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_id)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(progress.job_status, "finished");

    let index_id_failed = 11;
    let tidb_task_id_failed = "tidb_task_progress_2".to_string();
    let job_meta_failed = ImportJobMeta {
        tidb_task_id: tidb_task_id_failed.clone(),
        keyspace_id: NULLSPACEID,
        table_id,
        index_ids: vec![index_id_failed],
        job_id: 0,
    };
    meta_store
        .create_or_get_import_job(&job_meta_failed)
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id_failed,
            ImportTaskParams::new(
                "a".to_string(),
                "b".to_string(),
                "s3://bucket/path/to/file_a_b".to_string(),
            ),
        )
        .await
        .unwrap();
    meta_store
        .finish_import_job_upload(&tidb_task_id_failed)
        .await
        .unwrap();

    let task_failed = meta_store.next_import_task("worker-a").await.unwrap();
    assert!(task_failed.is_some());
    let task_failed = task_failed.unwrap();
    meta_store
        .import_task_build_fail(
            task_failed.task_id,
            &task_failed.owner_addr,
            task_failed.owner_seq,
            "boom",
        )
        .await
        .unwrap();

    let progress = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_id_failed)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(progress.job_status, "running");
    assert_eq!(progress.task_count, 1);
    assert_eq!(progress.failed_task_count, 1);
    let failed_jobs = meta_store.get_jobs_should_finish().await.unwrap();
    assert_eq!(failed_jobs.len(), 1);
    assert_eq!(failed_jobs[0].job_id, progress.job_id);
    assert_eq!(failed_jobs[0].running_task_count, 0);
    assert_eq!(failed_jobs[0].finished_task_count, 0);
    assert_eq!(failed_jobs[0].failed_task_count, 1);
    assert!(
        meta_store
            .get_active_import_jobs()
            .await
            .unwrap()
            .iter()
            .any(|job| job.job_id == progress.job_id)
    );

    meta_store
        .fail_import_job(&[progress.job_id])
        .await
        .unwrap();

    let progress = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_id_failed)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(progress.job_status, "failed");
    assert!(
        meta_store
            .get_active_import_jobs()
            .await
            .unwrap()
            .is_empty()
    );

    let other_keyspace_id = 1;
    let tidb_task_id_other_keyspace = "tidb_task_progress_3".to_string();
    let job_meta_other_keyspace = ImportJobMeta {
        tidb_task_id: tidb_task_id_other_keyspace.clone(),
        keyspace_id: other_keyspace_id,
        table_id,
        index_ids: vec![index_id],
        job_id: 0,
    };
    meta_store
        .create_or_get_import_job(&job_meta_other_keyspace)
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id_other_keyspace,
            ImportTaskParams::new(
                "z".to_string(),
                "zz".to_string(),
                "s3://bucket/path/to/file_z_zz".to_string(),
            ),
        )
        .await
        .unwrap();
    meta_store
        .finish_import_job_upload(&tidb_task_id_other_keyspace)
        .await
        .unwrap();

    let nullspace_progress = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_id)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(nullspace_progress.task_count, 2);
    assert_eq!(nullspace_progress.job_status, "finished");

    let other_keyspace_progress = meta_store
        .get_import_index_progress(other_keyspace_id, table_id, index_id)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(other_keyspace_progress.task_count, 1);
    assert_ne!(other_keyspace_progress.job_id, nullspace_progress.job_id);
    assert_eq!(other_keyspace_progress.job_status, "init");
}

#[tokio::test]
async fn test_get_import_index_progress_filters_tasks_by_index_id() -> anyhow::Result<()> {
    let config = load_meta_config("../../ci/meta-template.toml")?;
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await?;

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let table_id = 2100;
    let index_a = 11;
    let index_b = 22;
    let tidb_task_id = "tidb_task_progress_multi_index".to_string();
    meta_store
        .create_or_get_import_job(&ImportJobMeta {
            tidb_task_id: tidb_task_id.clone(),
            keyspace_id: NULLSPACEID,
            table_id,
            index_ids: vec![index_a, index_b],
            job_id: 0,
        })
        .await?;

    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "m".to_string(),
                "s3://bucket/path/to/file_index_a".to_string(),
            )
            .with_index_id(index_a),
        )
        .await?;
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "m".to_string(),
                "z".to_string(),
                "s3://bucket/path/to/file_index_b".to_string(),
            )
            .with_index_id(index_b),
        )
        .await?;

    let task_a = meta_store
        .next_import_task("worker-a")
        .await?
        .ok_or_else(|| anyhow::anyhow!("expected index_a import task"))?;
    assert_eq!(task_a.parameters.index_id, index_a);
    meta_store
        .import_task_build_success(
            task_a.task_id,
            &task_a.owner_addr,
            task_a.owner_seq,
            ImportTaskSummary::default(),
        )
        .await?;

    let progress_a = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_a)
        .await?
        .ok_or_else(|| anyhow::anyhow!("expected progress for index_a"))?;
    assert_eq!(progress_a.task_count, 1);
    assert_eq!(progress_a.finished_task_count, 1);
    assert_eq!(progress_a.running_task_count, 0);

    let progress_b = meta_store
        .get_import_index_progress(NULLSPACEID, table_id, index_b)
        .await?
        .ok_or_else(|| anyhow::anyhow!("expected progress for index_b"))?;
    assert_eq!(progress_b.task_count, 1);
    assert_eq!(progress_b.finished_task_count, 0);
    assert_eq!(progress_b.running_task_count, 0);

    Ok(())
}

#[tokio::test]
async fn test_get_import_index_progress_prefers_task_job_over_empty_duplicate() -> anyhow::Result<()>
{
    let config = load_meta_config("../../ci/meta-template.toml")?;
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await?;

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let table_id = 2200;
    let index_a = 11;
    let index_b = 22;
    let actual_tidb_task_id = "ImportInto/progress_duplicate".to_string();
    let actual_job_id = meta_store
        .create_or_get_import_job(&ImportJobMeta {
            tidb_task_id: actual_tidb_task_id.clone(),
            keyspace_id: NULLSPACEID,
            table_id,
            index_ids: vec![index_a, index_b],
            job_id: 0,
        })
        .await?;

    for (index_id, location) in [
        (index_a, "s3://bucket/path/to/file_index_a"),
        (index_b, "s3://bucket/path/to/file_index_b"),
    ] {
        meta_store
            .add_import_task(
                &actual_tidb_task_id,
                ImportTaskParams::new("a".to_string(), "m".to_string(), location.to_string())
                    .with_index_id(index_id),
            )
            .await?;
    }
    meta_store
        .finish_import_job_upload(&actual_tidb_task_id)
        .await?;

    for worker in ["worker-a", "worker-b"] {
        let task = meta_store
            .next_import_task(worker)
            .await?
            .ok_or_else(|| anyhow::anyhow!("expected import task"))?;
        meta_store
            .import_task_build_success(
                task.task_id,
                &task.owner_addr,
                task.owner_seq,
                ImportTaskSummary::default(),
            )
            .await?;
    }
    meta_store.finish_import_job(&[actual_job_id]).await?;

    let duplicate_job_id = meta_store
        .create_or_get_import_job(&ImportJobMeta {
            tidb_task_id: "progress_duplicate_presplit".to_string(),
            keyspace_id: NULLSPACEID,
            table_id,
            index_ids: vec![index_a, index_b],
            job_id: 0,
        })
        .await?;
    assert!(duplicate_job_id > actual_job_id);

    for index_id in [index_a, index_b] {
        let progress = meta_store
            .get_import_index_progress(NULLSPACEID, table_id, index_id)
            .await?
            .ok_or_else(|| anyhow::anyhow!("expected progress for index_id={index_id}"))?;
        assert_eq!(progress.job_id, actual_job_id);
        assert_eq!(progress.job_status, "finished");
        assert_eq!(progress.task_count, 1);
        assert_eq!(progress.finished_task_count, 1);
    }

    Ok(())
}

#[tokio::test]
async fn test_concurrent_create_or_get_job() {
    let config = load_meta_config("../../ci/meta-template.toml").unwrap();
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = Arc::new(MetaStoreImplTiDB::new(
        tidb_pools,
        MockLeadership::default(),
    ));

    let tidb_task_id = "tidb_task_1".to_string();
    let job_meta = ImportJobMeta {
        tidb_task_id: tidb_task_id.clone(),
        keyspace_id: NULLSPACEID,
        table_id: 1000,
        index_ids: vec![1],
        job_id: 0, // job_id will be assigned by meta_store
    };

    let mut handles = vec![];
    for _ in 0..10 {
        let meta_store = meta_store.clone();
        let job_meta = job_meta.clone();
        handles.push(tokio::spawn(async move {
            meta_store
                .create_or_get_import_job(&job_meta)
                .await
                .unwrap()
        }));
    }

    let mut job_ids = vec![];
    for h in handles {
        let job_id = h.await.unwrap();
        job_ids.push(job_id);
    }

    // All the concurrent calls should return the same job_id
    let first_job_id = job_ids[0];
    for &job_id in &job_ids[1..] {
        assert_eq!(job_id, first_job_id);
    }
}

#[tokio::test]
async fn test_import_store_invalid_task() {
    let config = load_meta_config("../../ci/meta-template.toml").unwrap();
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let tidb_task_id = "tidb_task_1".to_string();
    // Submit task without valid import job
    assert!(
        meta_store
            .add_import_task(
                &tidb_task_id,
                ImportTaskParams::new(
                    "a".to_string(),
                    "c".to_string(),
                    "s3://bucket/path/to/file_a_c".to_string(),
                )
            )
            .await
            .is_err()
    );
    // No valid task is submitted
    assert!(
        meta_store
            .next_import_task("worker-a")
            .await
            .unwrap()
            .is_none()
    );

    // Submit an import job and related task
    let job_meta = ImportJobMeta {
        tidb_task_id: tidb_task_id.clone(),
        keyspace_id: NULLSPACEID,
        table_id: 1000,
        index_ids: vec![1],
        job_id: 0, // job_id will be assigned by meta_store
    };
    meta_store
        .create_or_get_import_job(&job_meta)
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "c".to_string(),
                "s3://bucket/path/to/file_a_c".to_string(),
            ),
        )
        .await
        .unwrap();

    // Submit task without valid import job
    let tidb_task_id_2 = "tidb_task_2".to_string();
    assert!(
        meta_store
            .add_import_task(
                &tidb_task_id_2,
                ImportTaskParams::new(
                    "c".to_string(),
                    "z".to_string(),
                    "s3://bucket/path/to/file_c_z".to_string(),
                )
            )
            .await
            .is_err()
    );

    // Fetch only 1 task
    let task0 = meta_store.next_import_task("worker-a").await.unwrap();
    assert!(task0.is_some());
    assert!(
        meta_store
            .next_import_task("worker-b")
            .await
            .unwrap()
            .is_none()
    );
}

#[tokio::test]
async fn test_import_store_table_not_exist() {
    let config = load_meta_config("../../ci/meta-template.toml").unwrap();
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());
    // meta_store.start() is not called here

    // Calling `next_import_task` without starting the store should return None
    assert!(
        meta_store
            .next_import_task("worker-a")
            .await
            .unwrap()
            .is_none()
    );

    // Calling `get_jobs_should_finish` without starting the store should return an
    // empty vector
    assert!(
        meta_store
            .get_jobs_should_finish()
            .await
            .unwrap()
            .is_empty()
    );
    assert_eq!(meta_store.get_running_import_task_count().await.unwrap(), 0);
}

#[tokio::test]
async fn test_import_task_reclaim_owner_fencing_and_terminal_parent_guard() {
    let config = load_meta_config("../../ci/meta-template.toml").unwrap();
    Bootstrap::new_with_config(&config)
        .clear_all_meta(true)
        .await
        .unwrap();

    let tidb_pools = TiDBConnectionPool::new(&config.tidb_server);
    let meta_store = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default());

    let tidb_task_id = "tidb_task_reclaim".to_string();
    let job_meta = ImportJobMeta {
        tidb_task_id: tidb_task_id.clone(),
        keyspace_id: NULLSPACEID,
        table_id: 4321,
        index_ids: vec![9],
        job_id: 0,
    };
    let job_id = meta_store
        .create_or_get_import_job(&job_meta)
        .await
        .unwrap();
    meta_store
        .add_import_task(
            &tidb_task_id,
            ImportTaskParams::new(
                "a".to_string(),
                "z".to_string(),
                "s3://bucket/path/to/file".to_string(),
            ),
        )
        .await
        .unwrap();
    meta_store
        .finish_import_job_upload(&tidb_task_id)
        .await
        .unwrap();

    let task = meta_store
        .next_import_task("worker-a")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(
        meta_store
            .reclaim_import_tasks_by_worker("worker-a")
            .await
            .unwrap(),
        vec![task.task_id]
    );
    assert!(
        !meta_store
            .import_task_build_success(
                task.task_id,
                &task.owner_addr,
                task.owner_seq,
                ImportTaskSummary::default(),
            )
            .await
            .unwrap()
    );

    let reassigned = meta_store
        .next_import_task("worker-a")
        .await
        .unwrap()
        .unwrap();
    assert_eq!(reassigned.task_id, task.task_id);
    assert_eq!(reassigned.owner_addr, task.owner_addr);
    assert_eq!(reassigned.owner_seq, task.owner_seq + 1);
    assert!(
        !meta_store
            .import_task_build_success(
                task.task_id,
                &task.owner_addr,
                task.owner_seq,
                ImportTaskSummary::default(),
            )
            .await
            .unwrap()
    );
    assert!(
        !meta_store
            .import_task_build_success(
                reassigned.task_id,
                "worker-stale",
                reassigned.owner_seq,
                ImportTaskSummary::default(),
            )
            .await
            .unwrap()
    );
    meta_store.fail_import_job(&[job_id]).await.unwrap();
    assert!(
        !meta_store
            .import_task_build_success(
                reassigned.task_id,
                &reassigned.owner_addr,
                reassigned.owner_seq,
                ImportTaskSummary::default(),
            )
            .await
            .unwrap()
    );
    assert_eq!(
        meta_store
            .reclaim_import_tasks_by_worker("worker-a")
            .await
            .unwrap(),
        vec![reassigned.task_id]
    );
    assert!(
        meta_store
            .next_import_task("worker-c")
            .await
            .unwrap()
            .is_none()
    );
}