use std::{collections::HashMap, sync::Arc, time::Duration};

use log::{debug, error, info, warn};
use mysql_async::Error as MySQLError;
use storage_backend::storage::BackendStorage;
use tici_common::{
    config::S3Config,
    s3_keys::{import_frags_prefix, import_input_prefix},
};
use tici_import::{ImportJobMeta, ImportTask};
use tokio::sync::mpsc::UnboundedSender;

use crate::{
    config::MetaServiceConfig,
    meta::{
        ImportJobStat, MetaStore,
        writer::{ImportWorkerNodeSelection, ShardTaskManager, WorkerNodeMgr, WorkerNodeTask},
    },
};

pub struct ImportTaskMgr {
    task_gen_handle: tokio::task::JoinHandle<()>,
    job_check_handle: tokio::task::JoinHandle<()>,
}

const IMPORT_TASK_POLL_INTERVAL_SECS: u64 = 10;
const IMPORT_TASK_PENDING_RETRY_INTERVAL_MS: u64 = 100;
const IMPORT_JOB_CHECK_INTERVAL_SECS: u64 = 30;
const ABSENT_OWNER_RECLAIM_CHECK_TOLERANCE: u8 = 3;
const STORAGE_DELETE_BATCH_SIZE: usize = 1000;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ImportTaskDispatchResult {
    Dispatched,
    BlockedByPendingDispatch,
    NoProgress,
}

fn import_task_gen_wait_duration(result: ImportTaskDispatchResult) -> Option<Duration> {
    match result {
        ImportTaskDispatchResult::Dispatched => None,
        ImportTaskDispatchResult::BlockedByPendingDispatch => {
            Some(Duration::from_millis(IMPORT_TASK_PENDING_RETRY_INTERVAL_MS))
        }
        ImportTaskDispatchResult::NoProgress => {
            Some(Duration::from_secs(IMPORT_TASK_POLL_INTERVAL_SECS))
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ImportJobCleanupTarget {
    job_id: u64,
    keyspace_id: u32,
    table_id: i64,
    index_ids: Vec<i64>,
}

impl From<&ImportJobStat> for ImportJobCleanupTarget {
    fn from(value: &ImportJobStat) -> Self {
        Self {
            job_id: value.job_id,
            keyspace_id: value.keyspace_id,
            table_id: value.table_id,
            index_ids: value.index_ids.clone(),
        }
    }
}

impl From<&ImportJobMeta> for ImportJobCleanupTarget {
    fn from(value: &ImportJobMeta) -> Self {
        Self {
            job_id: value.job_id,
            keyspace_id: value.keyspace_id,
            table_id: value.table_id,
            index_ids: value.index_ids.clone(),
        }
    }
}

fn is_supported_external_task(task_key: &str) -> bool {
    match task_key.split('/').collect::<Vec<_>>().as_slice() {
        ["ImportInto", job_id] => !job_id.is_empty(),
        [keyspace, "ImportInto", job_id] => !keyspace.is_empty() && !job_id.is_empty(),
        ["ddl", "backfill", job_id] => !job_id.is_empty(),
        [keyspace, "ddl", "backfill", job_id] => !keyspace.is_empty() && !job_id.is_empty(),
        _ => false,
    }
}

fn build_import_input_cleanup_prefix(s3_prefix: &str, job: &ImportJobCleanupTarget) -> String {
    let import_prefix =
        import_input_prefix(job.keyspace_id, job.table_id, &job.index_ids, job.job_id);
    if s3_prefix.is_empty() {
        import_prefix
    } else {
        format!("{s3_prefix}/{import_prefix}")
    }
}

fn build_import_frag_gc_paths(s3_prefix: &str, job: &ImportJobCleanupTarget) -> Vec<String> {
    job.index_ids
        .iter()
        .map(|index_id| {
            let frag_prefix =
                import_frags_prefix(job.keyspace_id, job.table_id, *index_id, job.job_id);
            if s3_prefix.is_empty() {
                frag_prefix
            } else {
                format!("{s3_prefix}/{frag_prefix}")
            }
        })
        .collect()
}

fn filter_import_gc_paths(keys: Vec<String>, gc_paths: &[String]) -> Vec<String> {
    keys.into_iter()
        .filter(|key| {
            !gc_paths
                .iter()
                .any(|path| key == path || key.starts_with(&format!("{path}/")))
        })
        .collect()
}

async fn cleanup_import_job_input(
    backend_storage: &dyn BackendStorage,
    s3_config: &S3Config,
    job: &ImportJobCleanupTarget,
) -> Result<(), anyhow::Error> {
    let prefix = build_import_input_cleanup_prefix(&s3_config.prefix, job);
    let gc_paths = build_import_frag_gc_paths(&s3_config.prefix, job);
    let keys = backend_storage
        .list_all(&s3_config.bucket, &prefix)
        .await
        .map_err(anyhow::Error::from)?;
    if keys.is_empty() {
        info!(
            "ImportMgr - no import inputs to clean, job_id={} prefix={}",
            job.job_id, prefix
        );
        return Ok(());
    }
    let keys = filter_import_gc_paths(keys, &gc_paths);
    if keys.is_empty() {
        info!(
            "ImportMgr - no import inputs to clean after preserving fragments, job_id={} prefix={}",
            job.job_id, prefix
        );
        return Ok(());
    }
    info!(
        "ImportMgr - cleaning import inputs, job_id={} prefix={} key_count={}",
        job.job_id,
        prefix,
        keys.len()
    );
    for chunk in keys.chunks(STORAGE_DELETE_BATCH_SIZE) {
        backend_storage
            .delete_objects(&s3_config.bucket, chunk)
            .await
            .map_err(anyhow::Error::from)?;
    }
    Ok(())
}

async fn num_ongoing_tasks(meta_store: Arc<dyn MetaStore>) -> Result<usize, MySQLError> {
    // Ongoing tasks should only reflect tasks already dispatched to workers.
    meta_store.get_running_import_task_count().await
}

async fn import_task_recovery_check(
    meta_store: Arc<dyn MetaStore>,
    worker_node_mgr: Arc<WorkerNodeMgr>,
    absent_owner_checks: &mut HashMap<String, u8>,
) -> bool {
    let running_tasks = match meta_store.get_running_import_tasks().await {
        Ok(tasks) => tasks,
        Err(e) => {
            warn!("Error fetching running import tasks for recovery, err={e}");
            return false;
        }
    };
    if running_tasks.is_empty() {
        absent_owner_checks.clear();
        return false;
    }

    let mut tasks_by_owner: HashMap<String, Vec<ImportTask>> = HashMap::new();
    for task in running_tasks {
        if task.owner_addr.is_empty() {
            warn!(
                "ImportMgr - skip recovery for running import task without owner_addr, task_id={}",
                task.task_id
            );
            continue;
        }
        tasks_by_owner
            .entry(task.owner_addr.clone())
            .or_default()
            .push(task);
    }

    let mut recovered = false;
    absent_owner_checks.retain(|owner_addr, _| tasks_by_owner.contains_key(owner_addr));
    for (owner_addr, tasks) in tasks_by_owner {
        if let Some(worker) = worker_node_mgr.get_worker_node(&owner_addr).await {
            absent_owner_checks.remove(&owner_addr);
            let mut seeded_task_ids = Vec::new();
            for task in tasks {
                let task_id = task.task_id;
                if worker.remember_import_task(task).await {
                    seeded_task_ids.push(task_id);
                }
            }
            if !seeded_task_ids.is_empty() {
                recovered = true;
                info!(
                    "ImportMgr - restored running import tasks into worker state, owner_addr={}, task_ids={:?}",
                    owner_addr, seeded_task_ids
                );
            }
            continue;
        }

        let missing_checks = absent_owner_checks.entry(owner_addr.clone()).or_default();
        *missing_checks = missing_checks.saturating_add(1);
        if *missing_checks <= ABSENT_OWNER_RECLAIM_CHECK_TOLERANCE {
            debug!(
                "ImportMgr - defer reclaiming running import tasks for absent owner, owner_addr={}, missing_checks={}/{}",
                owner_addr, missing_checks, ABSENT_OWNER_RECLAIM_CHECK_TOLERANCE
            );
            continue;
        }

        match meta_store.reclaim_import_tasks_by_worker(&owner_addr).await {
            Ok(task_ids) => {
                if !task_ids.is_empty() {
                    recovered = true;
                    warn!(
                        "ImportMgr - reclaimed running import tasks for absent owner, owner_addr={}, missing_checks={}, task_ids={:?}",
                        owner_addr, missing_checks, task_ids
                    );
                }
                absent_owner_checks.remove(&owner_addr);
            }
            Err(e) => {
                warn!(
                    "ImportMgr - failed to reclaim running import tasks for absent owner, owner_addr={}, err={:?}",
                    owner_addr, e
                );
            }
        }
    }

    recovered
}

// TODO: make max_concurrency dynamic configurable
async fn import_task_gen_loop(
    meta_store: Arc<dyn MetaStore>,
    worker_node_mgr: Arc<WorkerNodeMgr>,
    shard_task_tx: UnboundedSender<WorkerNodeTask>,
    max_concurrency: usize,
) {
    let mut absent_owner_checks = HashMap::new();
    loop {
        // TODO: Add metrics for monitoring the time taken to acquire and process tasks
        // TODO: Fetch tasks in batch to reduce DB calls

        let _ = import_task_recovery_check(
            meta_store.clone(),
            worker_node_mgr.clone(),
            &mut absent_owner_checks,
        )
        .await;

        let mut dispatch_result = ImportTaskDispatchResult::NoProgress;
        let num_ongoing_tasks = match num_ongoing_tasks(meta_store.clone()).await {
            Ok(num_ongoing_tasks) => num_ongoing_tasks,
            Err(e) => {
                warn!("Error checking running import tasks, err={}", e);
                tokio::time::sleep(Duration::from_secs(IMPORT_TASK_POLL_INTERVAL_SECS)).await;
                continue;
            }
        };
        if num_ongoing_tasks < max_concurrency {
            dispatch_result = dispatch_next_import_task(
                meta_store.clone(),
                worker_node_mgr.clone(),
                shard_task_tx.clone(),
            )
            .await;
        } else {
            info!(
                "ImportMgr - Too many ongoing import jobs, skip acquiring new tasks. ongoing_jobs={} max_concurrency={}",
                num_ongoing_tasks, max_concurrency
            );
        }

        if let Some(wait_duration) = import_task_gen_wait_duration(dispatch_result) {
            if dispatch_result == ImportTaskDispatchResult::BlockedByPendingDispatch {
                tokio::select! {
                    _ = worker_node_mgr.wait_import_dispatch_progress() => {}
                    _ = tokio::time::sleep(wait_duration) => {}
                }
            } else {
                // No tasks to process, too many ongoing tasks, or no durable worker candidate.
                tokio::time::sleep(wait_duration).await;
            }
        }
    }
}

async fn dispatch_next_import_task(
    meta_store: Arc<dyn MetaStore>,
    worker_node_mgr: Arc<WorkerNodeMgr>,
    shard_task_tx: UnboundedSender<WorkerNodeTask>,
) -> ImportTaskDispatchResult {
    let owner_addr = match worker_node_mgr.select_import_worker_node().await {
        ImportWorkerNodeSelection::Selected(owner_addr) => owner_addr,
        ImportWorkerNodeSelection::BlockedByPendingDispatch => {
            debug!("All import worker nodes are blocked by pending dispatch.");
            return ImportTaskDispatchResult::BlockedByPendingDispatch;
        }
        ImportWorkerNodeSelection::NoAvailableWorker => {
            debug!("No worker node available for import task.");
            return ImportTaskDispatchResult::NoProgress;
        }
    };

    match meta_store.next_import_task(&owner_addr).await {
        Ok(Some(sub_task)) => {
            info!("ImportMgr - Processing import task={:?}", sub_task);

            let Some(worker) = worker_node_mgr.get_worker_node(&owner_addr).await else {
                warn!(
                    "ImportMgr - selected worker disappeared before dispatch, owner_addr={}, task_id={}",
                    owner_addr, sub_task.task_id
                );
                if let Err(e) = meta_store
                    .reclaim_import_task(sub_task.task_id, &sub_task.owner_addr)
                    .await
                {
                    warn!(
                        "ImportMgr - failed to reclaim task for disappeared worker, owner_addr={}, task_id={}, err={:?}",
                        owner_addr, sub_task.task_id, e
                    );
                }
                return ImportTaskDispatchResult::NoProgress;
            };

            if !worker.reserve_import_task(sub_task.clone()).await {
                warn!(
                    "ImportMgr - skip duplicate in-memory import task reservation, owner_addr={}, task_id={}, owner_seq={}",
                    owner_addr, sub_task.task_id, sub_task.owner_seq
                );
                return ImportTaskDispatchResult::NoProgress;
            }

            if let Err(e) =
                shard_task_tx.send(WorkerNodeTask::AddImportFragBuildTask(sub_task.clone()))
            {
                let _ = worker.remove_import_task(sub_task.task_id).await;
                if let Err(reclaim_err) = meta_store
                    .reclaim_import_task(sub_task.task_id, &sub_task.owner_addr)
                    .await
                {
                    warn!(
                        "ImportMgr - failed to reclaim task after worker dispatch send failure, owner_addr={}, task_id={}, err={:?}",
                        owner_addr, sub_task.task_id, reclaim_err
                    );
                }
                error!("Failed to send import task to worker, err={}", e);
                return ImportTaskDispatchResult::NoProgress;
            }
            ImportTaskDispatchResult::Dispatched
        }
        Ok(None) => {
            debug!("No import task available.");
            ImportTaskDispatchResult::NoProgress
        }
        Err(e) => {
            warn!("Error fetching import task, err={}", e);
            ImportTaskDispatchResult::NoProgress
        }
    }
}

async fn import_job_check(
    meta_store: Arc<dyn MetaStore>,
    shard_task_mgr: Arc<ShardTaskManager>,
    backend_storage: Arc<dyn BackendStorage>,
    s3_config: S3Config,
) -> bool {
    let jobs_stat = match meta_store.get_jobs_should_finish().await {
        Ok(stat) => stat,
        Err(e) => {
            warn!("Error checking import jobs, err={}", e);
            return false;
        }
    };

    let (terminal_jobs, ongoing_jobs): (Vec<_>, Vec<_>) = jobs_stat.into_iter().partition(|job| {
        job.running_task_count == 0
            && job.finished_task_count + job.failed_task_count == job.task_count
    });
    if terminal_jobs.is_empty() {
        if !ongoing_jobs.is_empty() {
            info!("ImportMgr - on-going jobs={:?}", ongoing_jobs);
        }
        return false;
    }

    let (failed_jobs, finished_jobs): (Vec<_>, Vec<_>) = terminal_jobs
        .into_iter()
        .partition(|job| job.failed_task_count > 0);
    let cleanup_jobs = finished_jobs
        .iter()
        .chain(failed_jobs.iter())
        .map(ImportJobCleanupTarget::from)
        .collect::<Vec<_>>();
    for job in &finished_jobs {
        if let Err(e) = shard_task_mgr
            .finish_import_append_base(job.keyspace_id, job.table_id, &job.index_ids)
            .await
        {
            warn!(
                "Failed to finish append base for import job_id={}, keyspace_id={}, table_id={}, index_ids={:?}, err={:?}",
                job.job_id, job.keyspace_id, job.table_id, job.index_ids, e
            );
            return false;
        }
    }
    let finished_job_ids = finished_jobs
        .iter()
        .map(|job| job.job_id)
        .collect::<Vec<_>>();
    let failed_job_ids = failed_jobs.iter().map(|job| job.job_id).collect::<Vec<_>>();
    for job in &cleanup_jobs {
        if let Err(e) = cleanup_import_job_input(backend_storage.as_ref(), &s3_config, job).await {
            warn!(
                "ImportMgr - failed to clean import inputs for terminal job (best-effort), job_id={} err={:?}",
                job.job_id, e
            );
        }
    }
    if !finished_job_ids.is_empty() {
        if let Err(e) = meta_store.finish_import_job(&finished_job_ids).await {
            warn!("Failed to finish import jobs={:?}", e);
            return false;
        }
    }
    if !failed_job_ids.is_empty() {
        if let Err(e) = meta_store.fail_import_job(&failed_job_ids).await {
            warn!("Failed to fail import jobs={:?}", e);
            return false;
        }
    }
    // Keep topology frozen until every task in the import job has reached a
    // terminal state so late base-frag appends cannot race with split/merge.
    for job in finished_jobs.iter().chain(failed_jobs.iter()) {
        shard_task_mgr
            .set_import_topology_frozen(job.keyspace_id, job.table_id, &job.index_ids, false)
            .await;
    }
    info!(
        "ImportMgr - Finished import jobs={:?}, failed jobs={:?}, on-going jobs={:?}",
        finished_job_ids, failed_job_ids, ongoing_jobs
    );
    true
}

async fn import_cleanup_check(
    meta_store: Arc<dyn MetaStore>,
    shard_task_mgr: Arc<ShardTaskManager>,
    backend_storage: Arc<dyn BackendStorage>,
    s3_config: S3Config,
) -> bool {
    let active_jobs = match meta_store.get_active_import_jobs().await {
        Ok(jobs) => jobs,
        Err(e) => {
            warn!("Error fetching active import jobs for cleanup, err={e:?}");
            return false;
        }
    };
    if active_jobs.is_empty() {
        return false;
    }

    let mut supported_jobs = Vec::new();
    for job in active_jobs {
        if is_supported_external_task(&job.tidb_task_id) {
            supported_jobs.push(job);
        } else {
            warn!(
                "ImportMgr - skip external cleanup for task_key without known prefix, task_key={}",
                job.tidb_task_id
            );
        }
    }
    if supported_jobs.is_empty() {
        return false;
    }

    let known_task_keys = supported_jobs
        .iter()
        .map(|job| job.tidb_task_id.clone())
        .collect::<Vec<_>>();
    let terminal_task_keys = match meta_store
        .get_external_terminal_task_keys(&known_task_keys)
        .await
    {
        Ok(states) => states,
        Err(e) => {
            warn!("Error fetching external terminal task states, err={e:?}");
            return false;
        }
    };
    let cleanup_jobs = supported_jobs
        .into_iter()
        .filter(|job| terminal_task_keys.contains(&job.tidb_task_id))
        .collect::<Vec<_>>();
    if cleanup_jobs.is_empty() {
        return false;
    }

    for job in &cleanup_jobs {
        let cleanup_target = ImportJobCleanupTarget::from(job);
        if let Err(e) =
            cleanup_import_job_input(backend_storage.as_ref(), &s3_config, &cleanup_target).await
        {
            warn!(
                "ImportMgr - failed to clean import inputs for externally terminated job (best-effort), job_id={} err={:?}",
                job.job_id, e
            );
        }
    }

    let failed_job_ids = cleanup_jobs
        .iter()
        .map(|job| job.job_id)
        .collect::<Vec<_>>();
    let reclaimed_task_ids = match meta_store
        .fail_import_jobs_and_reclaim_running_tasks(&failed_job_ids)
        .await
    {
        Ok(task_ids) => task_ids,
        Err(e) => {
            warn!("Failed to fail/reclaim externally terminated import jobs={e:?}");
            return false;
        }
    };
    if !reclaimed_task_ids.is_empty() {
        info!(
            "ImportMgr - reclaimed running tasks for externally terminated jobs, job_ids={:?}, task_ids={:?}",
            failed_job_ids, reclaimed_task_ids
        );
    }
    for job in &cleanup_jobs {
        shard_task_mgr
            .set_import_topology_frozen_for_job(job, false)
            .await;
    }
    info!(
        "ImportMgr - cleaned externally terminated import jobs={:?}",
        failed_job_ids
    );
    true
}

#[cfg(any(test, feature = "testexport"))]
pub(crate) async fn run_import_cleanup_check_once_for_test(
    meta_store: Arc<dyn MetaStore>,
    shard_task_mgr: Arc<ShardTaskManager>,
    backend_storage: Arc<dyn BackendStorage>,
    s3_config: S3Config,
) -> bool {
    import_cleanup_check(meta_store, shard_task_mgr, backend_storage, s3_config).await
}

#[cfg(any(test, feature = "testexport"))]
pub(crate) async fn run_import_recovery_once_for_test(
    meta_store: Arc<dyn MetaStore>,
    worker_node_mgr: Arc<WorkerNodeMgr>,
) -> bool {
    let mut absent_owner_checks = HashMap::new();
    import_task_recovery_check(meta_store, worker_node_mgr, &mut absent_owner_checks).await
}

#[cfg(any(test, feature = "testexport"))]
pub(crate) async fn run_import_job_check_once_for_test(
    meta_store: Arc<dyn MetaStore>,
    shard_task_mgr: Arc<ShardTaskManager>,
    backend_storage: Arc<dyn BackendStorage>,
    s3_config: S3Config,
) -> bool {
    import_job_check(meta_store, shard_task_mgr, backend_storage, s3_config).await
}

async fn import_job_check_loop(
    meta_store: Arc<dyn MetaStore>,
    shard_task_mgr: Arc<ShardTaskManager>,
    backend_storage: Arc<dyn BackendStorage>,
    s3_config: S3Config,
) {
    loop {
        let cleaned_external_jobs = import_cleanup_check(
            meta_store.clone(),
            shard_task_mgr.clone(),
            backend_storage.clone(),
            s3_config.clone(),
        )
        .await;
        let finished_jobs = import_job_check(
            meta_store.clone(),
            shard_task_mgr.clone(),
            backend_storage.clone(),
            s3_config.clone(),
        )
        .await;
        if !cleaned_external_jobs && !finished_jobs {
            // Sleep for a while before checking again
            tokio::time::sleep(tokio::time::Duration::from_secs(
                IMPORT_JOB_CHECK_INTERVAL_SECS,
            ))
            .await;
        }
    }
}

impl ImportTaskMgr {
    pub fn new(
        config: &MetaServiceConfig,
        meta_store: Arc<dyn MetaStore>,
        shard_task_mgr: Arc<ShardTaskManager>,
        worker_node_mgr: Arc<WorkerNodeMgr>,
        shard_task_tx: UnboundedSender<WorkerNodeTask>,
        backend_storage: Arc<dyn BackendStorage>,
    ) -> Self {
        let task_gen_handle = tokio::spawn(import_task_gen_loop(
            meta_store.clone(),
            worker_node_mgr,
            shard_task_tx,
            config.import.max_concurrency,
        ));
        let job_check_handle = tokio::spawn(import_job_check_loop(
            meta_store.clone(),
            shard_task_mgr.clone(),
            backend_storage.clone(),
            config.s3.clone(),
        ));

        Self {
            task_gen_handle,
            job_check_handle,
        }
    }

    pub fn stop(&self) {
        self.job_check_handle.abort();
        self.task_gen_handle.abort();
        info!("ImportTaskMgr stopped.");
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, pin::Pin, sync::Arc, time::Duration};

    use arc_swap::ArcSwap;
    use bytes::Bytes;
    use futures::Stream;
    use storage_backend::storage::{
        BackendStorage, ObjectMeta, ReadRange, Result as StorageResult, StorageError,
    };
    use tici_grpc::{ImportWriteStatus, WorkerNodeHeartbeatRequest, WorkerNodeStatus};
    use tici_import::{ImportTask, ImportTaskParams};
    use tokio::io::AsyncRead;

    use super::{
        ABSENT_OWNER_RECLAIM_CHECK_TOLERANCE, IMPORT_TASK_POLL_INTERVAL_SECS,
        ImportTaskDispatchResult, build_import_frag_gc_paths, build_import_input_cleanup_prefix,
        cleanup_import_job_input, dispatch_next_import_task, filter_import_gc_paths,
        import_task_gen_wait_duration, import_task_recovery_check, is_supported_external_task,
        run_import_recovery_once_for_test,
    };
    use crate::{
        config::MetaServiceConfig,
        meta::{
            MetaStore,
            store::test_utils::RecordingMetaStore,
            writer::{WorkerNodeMgr, WorkerNodeTask},
        },
    };

    #[test]
    fn test_build_import_input_cleanup_prefix() {
        let job = super::ImportJobCleanupTarget {
            job_id: 7,
            keyspace_id: 0,
            table_id: 42,
            index_ids: vec![9],
        };
        assert_eq!(
            build_import_input_cleanup_prefix("root", &job),
            "root/ks_0/t_42/i_9/import/7"
        );
        assert_eq!(
            build_import_input_cleanup_prefix("", &job),
            "ks_0/t_42/i_9/import/7"
        );
    }

    #[test]
    fn test_is_supported_external_task() {
        assert!(is_supported_external_task("ImportInto/1"));
        assert!(is_supported_external_task("ks1/ImportInto/1"));
        assert!(is_supported_external_task("ddl/backfill/42"));
        assert!(is_supported_external_task("ks1/ddl/backfill/42"));
        assert!(!is_supported_external_task("1"));
        assert!(!is_supported_external_task("AddIndex/42"));
        assert!(!is_supported_external_task("ImportInto/"));
        assert!(!is_supported_external_task("ks1/ImportInto/1/extra"));
    }

    #[test]
    fn test_build_import_frag_gc_paths() {
        let job = super::ImportJobCleanupTarget {
            job_id: 7,
            keyspace_id: 0,
            table_id: 42,
            index_ids: vec![9, 11],
        };
        assert_eq!(
            build_import_frag_gc_paths("root", &job),
            vec![
                "root/ks_0/t_42/i_9/import/7/frags".to_string(),
                "root/ks_0/t_42/i_11/import/7/frags".to_string(),
            ]
        );
    }

    #[test]
    fn test_filter_import_gc_paths() {
        let gc_paths = vec!["root/ks_0/t_42/i_9/import/7/frags".to_string()];
        let filtered = filter_import_gc_paths(
            vec![
                "root/ks_0/t_42/i_9/import/7/frags/t_42/i_9/f_1/multi.0".to_string(),
                "root/ks_0/t_42/i_9/import/7/00124676-af60-45dc-8450-3698805ece0b".to_string(),
            ],
            &gc_paths,
        );
        assert_eq!(
            filtered,
            vec!["root/ks_0/t_42/i_9/import/7/00124676-af60-45dc-8450-3698805ece0b".to_string()]
        );
    }

    #[derive(Default)]
    struct RecordingBackendStorage {
        listed_keys: std::sync::Mutex<Vec<String>>,
        deleted_keys: std::sync::Mutex<Vec<String>>,
        list_error_prefixes: std::sync::Mutex<Vec<String>>,
        delete_error_keys: std::sync::Mutex<Vec<String>>,
    }

    #[async_trait::async_trait]
    impl BackendStorage for RecordingBackendStorage {
        async fn put_object(
            &self,
            _bucket: &str,
            _key: &str,
            _body: Pin<
                Box<
                    dyn Stream<Item = std::result::Result<Bytes, std::io::Error>>
                        + Send
                        + Sync
                        + 'static,
                >,
            >,
            _content_len: u64,
        ) -> StorageResult<()> {
            Err(StorageError::Unknown("not used in test".to_string()))
        }

        async fn put_data(&self, _bucket: &str, _key: &str, _data: Vec<u8>) -> StorageResult<()> {
            Err(StorageError::Unknown("not used in test".to_string()))
        }

        async fn get_object(
            &self,
            _bucket: &str,
            _key: &str,
            _range: Option<ReadRange>,
        ) -> StorageResult<(Box<dyn AsyncRead + Unpin + Send>, ObjectMeta)> {
            Err(StorageError::Unknown("not used in test".to_string()))
        }

        async fn head_object(&self, _bucket: &str, _key: &str) -> StorageResult<ObjectMeta> {
            Err(StorageError::Unknown("not used in test".to_string()))
        }

        async fn list(
            &self,
            _bucket: &str,
            _prefix: &str,
            _start_after: Option<&str>,
        ) -> StorageResult<Vec<String>> {
            Err(StorageError::Unknown("not used in test".to_string()))
        }

        async fn list_all(&self, _bucket: &str, prefix: &str) -> StorageResult<Vec<String>> {
            let list_error_prefixes = match self.list_error_prefixes.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            if list_error_prefixes
                .iter()
                .any(|error_prefix| error_prefix == prefix)
            {
                return Err(StorageError::Unknown(format!(
                    "list failed for prefix {prefix}"
                )));
            }
            drop(list_error_prefixes);

            let guard = match self.listed_keys.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            Ok(guard
                .iter()
                .filter(|key| key.starts_with(prefix))
                .cloned()
                .collect())
        }

        async fn list_with_delimiter(
            &self,
            _bucket: &str,
            _prefix: &str,
            _start_after: Option<&str>,
            _delimiter: &str,
        ) -> StorageResult<Vec<String>> {
            Err(StorageError::Unknown("not used in test".to_string()))
        }

        async fn delete_objects(&self, _bucket: &str, keys: &[String]) -> StorageResult<()> {
            let delete_error_keys = match self.delete_error_keys.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            if keys.iter().any(|key| delete_error_keys.contains(key)) {
                return Err(StorageError::Unknown(format!(
                    "delete failed for keys {keys:?}"
                )));
            }
            drop(delete_error_keys);

            let mut guard = match self.deleted_keys.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            guard.extend_from_slice(keys);
            Ok(())
        }
    }

    fn test_s3_config() -> tici_common::config::S3Config {
        tici_common::config::S3Config {
            bucket: "bucket".to_string(),
            prefix: "root".to_string(),
            ..Default::default()
        }
    }

    fn cleanup_target(job_id: u64) -> super::ImportJobCleanupTarget {
        super::ImportJobCleanupTarget {
            job_id,
            keyspace_id: 0,
            table_id: 42,
            index_ids: vec![9],
        }
    }

    fn import_input_key(job: &super::ImportJobCleanupTarget, suffix: &str) -> String {
        format!(
            "{}/{}",
            build_import_input_cleanup_prefix("root", job),
            suffix
        )
    }

    fn import_frag_key(job: &super::ImportJobCleanupTarget, suffix: &str) -> String {
        format!("{}/{}", build_import_frag_gc_paths("root", job)[0], suffix)
    }

    #[tokio::test]
    async fn cleanup_import_job_input_preserves_frag_objects() {
        let job = cleanup_target(7);
        let input_key = import_input_key(&job, "00124676-af60-45dc-8450-3698805ece0b");
        let backend_storage = Arc::new(RecordingBackendStorage {
            listed_keys: std::sync::Mutex::new(vec![
                input_key.clone(),
                import_frag_key(&job, "t_42/i_9/f_123/multi.0"),
                import_frag_key(&job, "t_42/i_9/f_123/multi.1"),
            ]),
            ..Default::default()
        });
        let s3_config = test_s3_config();

        cleanup_import_job_input(backend_storage.as_ref(), &s3_config, &job)
            .await
            .unwrap();

        let deleted_keys = match backend_storage.deleted_keys.lock() {
            Ok(guard) => guard.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        };
        assert_eq!(deleted_keys, vec![input_key]);
    }

    #[tokio::test]
    async fn cleanup_import_job_input_allows_per_job_best_effort_loop() {
        let list_error_job = cleanup_target(7);
        let delete_error_job = cleanup_target(8);
        let successful_job = cleanup_target(9);
        let delete_error_key =
            import_input_key(&delete_error_job, "00124676-af60-45dc-8450-3698805ece0b");
        let successful_key =
            import_input_key(&successful_job, "00124676-af60-45dc-8450-3698805ece0b");
        let backend_storage = Arc::new(RecordingBackendStorage {
            listed_keys: std::sync::Mutex::new(vec![
                delete_error_key.clone(),
                successful_key.clone(),
            ]),
            list_error_prefixes: std::sync::Mutex::new(vec![build_import_input_cleanup_prefix(
                "root",
                &list_error_job,
            )]),
            delete_error_keys: std::sync::Mutex::new(vec![delete_error_key.clone()]),
            ..Default::default()
        });
        let s3_config = test_s3_config();
        let jobs = vec![list_error_job, delete_error_job, successful_job];

        let mut failed_job_ids = Vec::new();
        for job in &jobs {
            if cleanup_import_job_input(backend_storage.as_ref(), &s3_config, job)
                .await
                .is_err()
            {
                failed_job_ids.push(job.job_id);
            }
        }

        assert_eq!(failed_job_ids, vec![7, 8]);
        let deleted_keys = match backend_storage.deleted_keys.lock() {
            Ok(guard) => guard.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        };
        assert_eq!(deleted_keys, vec![successful_key]);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn recovery_check_reclaims_tasks_for_absent_owner() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        {
            let mut tasks = match store.running_import_tasks.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            tasks.push(ImportTask {
                task_id: 1,
                job_id: 10,
                owner_addr: "missing-worker".to_string(),
                owner_seq: 1,
                parameters: ImportTaskParams::new(
                    "a".to_string(),
                    "z".to_string(),
                    "mock://import/missing-worker".to_string(),
                ),
            });
        }

        let store_dyn: Arc<dyn MetaStore> = store.clone();
        let mgr = WorkerNodeMgr::new(
            shard_task_tx,
            wn_task_tx,
            wn_task_rx,
            config,
            store_dyn.clone(),
        );

        let mut absent_owner_checks = HashMap::new();
        for _ in 0..ABSENT_OWNER_RECLAIM_CHECK_TOLERANCE {
            assert!(
                !import_task_recovery_check(
                    store_dyn.clone(),
                    mgr.clone(),
                    &mut absent_owner_checks,
                )
                .await
            );
        }
        assert!(import_task_recovery_check(store_dyn, mgr, &mut absent_owner_checks).await);

        let reclaimed = match store.reclaimed_import_workers.lock() {
            Ok(guard) => guard.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        };
        assert_eq!(reclaimed, vec!["missing-worker".to_string()]);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn recovery_check_seeds_worker_state_then_heartbeat_reclaims_missing_task() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        let task = ImportTask {
            task_id: 7,
            job_id: 11,
            owner_addr: "worker-a".to_string(),
            owner_seq: 1,
            parameters: ImportTaskParams::new(
                "a".to_string(),
                "z".to_string(),
                "mock://import/worker-a".to_string(),
            ),
        };
        {
            let mut tasks = match store.running_import_tasks.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            tasks.push(task.clone());
        }

        let store_dyn: Arc<dyn MetaStore> = store.clone();
        let mgr = WorkerNodeMgr::new(
            shard_task_tx,
            wn_task_tx,
            wn_task_rx,
            config,
            store_dyn.clone(),
        );
        mgr.report_worker_node_status(WorkerNodeHeartbeatRequest {
            addr: task.owner_addr.clone(),
            status: Some(WorkerNodeStatus::default()),
            ..Default::default()
        })
        .await
        .unwrap();
        let worker = mgr.get_worker_node(&task.owner_addr).await.unwrap();
        assert!(!worker.has_import_task(task.task_id).await);

        assert!(run_import_recovery_once_for_test(store_dyn.clone(), mgr.clone()).await);
        assert!(worker.has_import_task(task.task_id).await);

        for _ in 0..3 {
            mgr.report_worker_node_status(WorkerNodeHeartbeatRequest {
                addr: task.owner_addr.clone(),
                status: Some(WorkerNodeStatus::default()),
                ..Default::default()
            })
            .await
            .unwrap();
        }

        let reclaimed = match store.reclaimed_import_tasks.lock() {
            Ok(guard) => guard.clone(),
            Err(poisoned) => poisoned.into_inner().clone(),
        };
        assert_eq!(reclaimed, vec![(task.task_id, task.owner_addr.clone())]);
        assert!(!worker.has_import_task(task.task_id).await);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn recovery_check_clears_seeded_untracked_import_task() -> anyhow::Result<()> {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        let task = ImportTask {
            task_id: 9,
            job_id: 12,
            owner_addr: "worker-recovery-untracked".to_string(),
            owner_seq: 4,
            parameters: ImportTaskParams::new(
                "a".to_string(),
                "z".to_string(),
                "mock://import/worker-recovery-untracked".to_string(),
            ),
        };
        {
            let mut tasks = match store.running_import_tasks.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            tasks.push(task.clone());
        }

        let store_dyn: Arc<dyn MetaStore> = store.clone();
        let mgr = WorkerNodeMgr::new(
            shard_task_tx,
            wn_task_tx,
            wn_task_rx,
            config,
            store_dyn.clone(),
        );
        mgr.report_worker_node_status(WorkerNodeHeartbeatRequest {
            addr: task.owner_addr.clone(),
            status: Some(WorkerNodeStatus::default()),
            import_writes: vec![ImportWriteStatus {
                task_id: task.task_id,
                owner_seq: task.owner_seq,
                ..Default::default()
            }],
            ..Default::default()
        })
        .await?;

        let worker = mgr
            .get_worker_node(&task.owner_addr)
            .await
            .ok_or_else(|| anyhow::anyhow!("missing worker {}", task.owner_addr))?;
        assert_eq!(worker.untracked_import_tasks_len().await, 1);
        assert_eq!(mgr.select_import_worker_node_addr().await.as_deref(), None);

        assert!(run_import_recovery_once_for_test(store_dyn, mgr.clone()).await);
        assert!(worker.has_import_task(task.task_id).await);
        assert_eq!(worker.untracked_import_tasks_len().await, 0);
        assert_eq!(
            mgr.select_import_worker_node_addr().await.as_deref(),
            Some(task.owner_addr.as_str())
        );
        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn dispatch_next_import_task_reserves_worker_slot_before_queue_consumption() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        {
            let mut queued = match store.next_import_tasks.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            queued.push(ImportTask {
                task_id: 99,
                job_id: 18,
                owner_addr: String::new(),
                owner_seq: 1,
                parameters: ImportTaskParams::new(
                    "a".to_string(),
                    "z".to_string(),
                    "mock://import/reserve-before-queue".to_string(),
                ),
            });
        }

        let store_dyn: Arc<dyn MetaStore> = store.clone();
        let mgr = WorkerNodeMgr::new(
            shard_task_tx,
            wn_task_tx,
            wn_task_rx,
            config,
            store_dyn.clone(),
        );
        mgr.create_worker_node(store_dyn.clone(), "worker-a".to_string(), None)
            .await;
        mgr.create_worker_node(store_dyn, "worker-b".to_string(), None)
            .await;
        let (dispatch_tx, _dispatch_rx) = tokio::sync::mpsc::unbounded_channel();

        assert_eq!(
            dispatch_next_import_task(store.clone(), mgr.clone(), dispatch_tx).await,
            ImportTaskDispatchResult::Dispatched
        );

        let worker_a = mgr.get_worker_node("worker-a").await.unwrap();
        let worker_b = mgr.get_worker_node("worker-b").await.unwrap();
        let reserved_owner = if worker_a.has_import_task(99).await {
            "worker-a"
        } else if worker_b.has_import_task(99).await {
            "worker-b"
        } else {
            panic!("expected dispatch_next_import_task to reserve task 99 on one worker");
        };

        assert_ne!(
            mgr.select_import_worker_node_addr().await.as_deref(),
            Some(reserved_owner)
        );
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn dispatch_next_import_task_does_not_stack_pending_dispatches_on_worker() {
        let (shard_task_tx, _shard_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let (wn_task_tx, wn_task_rx) = tokio::sync::mpsc::unbounded_channel();
        let config = Arc::new(ArcSwap::from_pointee(MetaServiceConfig::default()));
        let store = Arc::new(RecordingMetaStore::default());
        {
            let mut queued = match store.next_import_tasks.lock() {
                Ok(guard) => guard,
                Err(poisoned) => poisoned.into_inner(),
            };
            queued.push(ImportTask {
                task_id: 100,
                job_id: 19,
                owner_addr: String::new(),
                owner_seq: 1,
                parameters: ImportTaskParams::new(
                    "a".to_string(),
                    "m".to_string(),
                    "mock://import/pending-stack-1".to_string(),
                ),
            });
            queued.push(ImportTask {
                task_id: 101,
                job_id: 19,
                owner_addr: String::new(),
                owner_seq: 1,
                parameters: ImportTaskParams::new(
                    "m".to_string(),
                    "z".to_string(),
                    "mock://import/pending-stack-2".to_string(),
                ),
            });
        }

        let store_dyn: Arc<dyn MetaStore> = store.clone();
        let mgr = WorkerNodeMgr::new(
            shard_task_tx,
            wn_task_tx,
            wn_task_rx,
            config,
            store_dyn.clone(),
        );
        let worker = mgr
            .create_worker_node(store_dyn, "worker-a".to_string(), None)
            .await;
        let (dispatch_tx, mut dispatch_rx) = tokio::sync::mpsc::unbounded_channel();

        assert_eq!(
            dispatch_next_import_task(store.clone(), mgr.clone(), dispatch_tx.clone()).await,
            ImportTaskDispatchResult::Dispatched
        );
        assert_eq!(worker.pending_import_tasks_len().await, 1);

        let blocked =
            dispatch_next_import_task(store.clone(), mgr.clone(), dispatch_tx.clone()).await;
        assert_eq!(blocked, ImportTaskDispatchResult::BlockedByPendingDispatch);
        let pending_wait = import_task_gen_wait_duration(blocked)
            .expect("pending dispatch block should use a retry wait");
        assert!(pending_wait < Duration::from_secs(IMPORT_TASK_POLL_INTERVAL_SECS));

        assert_eq!(worker.import_tasks_len().await, 1);
        assert_eq!(worker.pending_import_tasks_len().await, 1);

        let queued = match store.next_import_tasks.lock() {
            Ok(guard) => guard.len(),
            Err(poisoned) => poisoned.into_inner().len(),
        };
        assert_eq!(queued, 1);

        let pending_task = match dispatch_rx.try_recv() {
            Ok(WorkerNodeTask::AddImportFragBuildTask(task)) => task,
            Ok(other) => panic!("expected import dispatch task, got {other:?}"),
            Err(e) => panic!("expected one import dispatch task, err={e:?}"),
        };
        assert!(
            worker
                .mark_import_task_running(
                    pending_task.task_id,
                    &pending_task.owner_addr,
                    pending_task.owner_seq,
                )
                .await
        );
        assert_eq!(worker.pending_import_tasks_len().await, 0);
        assert_eq!(
            dispatch_next_import_task(store.clone(), mgr, dispatch_tx).await,
            ImportTaskDispatchResult::Dispatched
        );
    }
}