use anyhow::Result;
use mysql_async::{params, prelude::Queryable};
use storage_backend::storage::create_backend_storage_for_test;
use tici::{
    config::{make_default_test_meta_config, make_default_test_worker_config},
    connection_pool::TiDBConnectionPool,
    meta::Bootstrap,
    worker::mocks::make_frag_meta,
};
use tici_common::{index::IndexKey, s3_keys::import_input_prefix};
use tici_grpc::{WorkerNodeHeartbeatRequest, WorkerNodeStatus};
use tici_import::{ImportJobMeta, ImportTaskParams, ImportTaskSummary};
use tici_test_utils::test_utils::{
    AppendExtraParam, meta_scaffold::build_meta_service, mock_frag_path, mock_frag_paths,
    test_suite::E2ETestSuite,
};

async fn debug_insert_external_terminal_task(
    config: &tici::config::MetaServiceConfig,
    task_key: &str,
    state: &str,
) {
    let pools = TiDBConnectionPool::new(&config.tidb_server);
    let mut conn = pools.get_conn().await.unwrap();
    conn.exec_drop(
        "INSERT INTO mysql.tidb_global_task(task_key, type, state)
            VALUES(:task_key, 'import-into', :state)
            ON DUPLICATE KEY UPDATE state=VALUES(state)",
        params! {
            "task_key" => task_key,
            "state" => state,
        },
    )
    .await
    .unwrap();
}

#[test]
fn test_import_job_check_finishes_append_base_and_switches_to_delta() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut suite = E2ETestSuite::new_with_mock_cfg(
            make_default_test_meta_config(false),
            make_default_test_worker_config(false),
        )
        .await
        .unwrap();

        let index_key = IndexKey::new_nullspace(1, 1005);
        suite.must_create_index(&index_key).await.unwrap();
        suite.validate_add_index_result(&index_key).await;

        let shard = suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let worker_addr = suite.get_one_worker_node_addr();
        let meta_service = suite.get_meta_service();

        let tidb_task_id = "import-job-check-switches-base-to-delta".to_string();
        meta_service
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: tidb_task_id.clone(),
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                job_id: 0,
            })
            .await
            .unwrap();
        meta_service
            .shard_task_manager
            .set_import_topology_frozen(
                index_key.keyspace_id,
                index_key.table_id,
                &[index_key.index_id],
                true,
            )
            .await;

        meta_service
            .meta_store
            .add_import_task(
                &tidb_task_id,
                ImportTaskParams::new(
                    "a".to_string(),
                    "z".to_string(),
                    "mock://import/task".to_string(),
                ),
            )
            .await
            .unwrap();
        meta_service
            .meta_store
            .finish_import_job_upload(&tidb_task_id)
            .await
            .unwrap();
        let task = meta_service
            .meta_store
            .next_import_task("worker-a")
            .await
            .unwrap()
            .unwrap();
        meta_service
            .meta_store
            .import_task_build_success(
                task.task_id,
                &task.owner_addr,
                task.owner_seq,
                ImportTaskSummary::default(),
            )
            .await
            .unwrap();

        suite
            .must_append_base_frag(
                &worker_addr,
                shard_id,
                &index_key,
                vec![make_frag_meta(&index_key, &mock_frag_path("base_0"))],
            )
            .await
            .unwrap();
        suite
            .must_append_base_frag(
                &worker_addr,
                shard_id,
                &index_key,
                vec![make_frag_meta(&index_key, &mock_frag_path("base_1"))],
            )
            .await
            .unwrap();

        let shard_before_check = suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        let shard_before_check = shard_before_check.read().await;
        assert_eq!(shard_before_check.manifest.base_slice().len(), 2);
        assert_eq!(shard_before_check.manifest.delta_slice().len(), 0);
        drop(shard_before_check);

        assert!(meta_service.run_import_job_check_once_for_test().await);

        let index = suite.get_index(&index_key).await;
        assert!(index.is_finished_append_base());
        assert!(!index.is_topology_frozen());

        let shard_after_check = suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        let shard_after_check = shard_after_check.read().await;
        assert_eq!(shard_after_check.manifest.base_slice().len(), 0);
        assert_eq!(shard_after_check.manifest.delta_slice().len(), 2);
        assert_eq!(
            shard_after_check
                .manifest
                .delta_slice()
                .iter()
                .map(|frag| frag.frag_path.clone())
                .collect::<Vec<_>>(),
            mock_frag_paths(&["base_1", "base_0"])
        );
        drop(shard_after_check);

        let shard_after_switch = suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        let shard_after_switch = shard_after_switch.read().await.clone();
        suite
            .must_append_delta_by_shard(
                &shard_after_switch,
                "delta_0",
                1,
                AppendExtraParam::default(),
            )
            .await;

        let shard_after_delta = suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        let shard_after_delta = shard_after_delta.read().await;
        assert_eq!(shard_after_delta.manifest.base_slice().len(), 0);
        assert_eq!(shard_after_delta.manifest.delta_slice().len(), 3);
        assert_eq!(
            shard_after_delta
                .manifest
                .delta_slice()
                .iter()
                .map(|frag| frag.frag_path.clone())
                .collect::<Vec<_>>(),
            mock_frag_paths(&["base_1", "base_0", "delta_0"])
        );

        suite.stop().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_import_job_check_cleans_import_inputs_for_finished_job() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut suite = E2ETestSuite::new_with_mock_cfg(
            make_default_test_meta_config(false),
            make_default_test_worker_config(false),
        )
        .await
        .unwrap();

        let index_key = IndexKey::new_nullspace(1, 1005);
        suite.must_create_index(&index_key).await.unwrap();
        suite.validate_add_index_result(&index_key).await;

        let meta_service = suite.get_meta_service();
        let tidb_task_id = "ImportInto/finished-cleanup".to_string();
        let job_id = meta_service
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: tidb_task_id.clone(),
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                job_id: 0,
            })
            .await
            .unwrap();
        meta_service
            .shard_task_manager
            .set_import_topology_frozen(
                index_key.keyspace_id,
                index_key.table_id,
                &[index_key.index_id],
                true,
            )
            .await;
        meta_service
            .meta_store
            .add_import_task(
                &tidb_task_id,
                ImportTaskParams::new(
                    "a".to_string(),
                    "z".to_string(),
                    "mock://import/task".to_string(),
                ),
            )
            .await
            .unwrap();
        meta_service
            .meta_store
            .finish_import_job_upload(&tidb_task_id)
            .await
            .unwrap();
        let task = meta_service
            .meta_store
            .next_import_task("worker-a")
            .await
            .unwrap()
            .unwrap();
        meta_service
            .meta_store
            .import_task_build_success(
                task.task_id,
                &task.owner_addr,
                task.owner_seq,
                ImportTaskSummary::default(),
            )
            .await
            .unwrap();

        let mut s3_cfg = meta_service.config.load().s3.clone();
        s3_cfg.use_path_style = true;
        let storage = create_backend_storage_for_test(&s3_cfg).await.unwrap();
        let import_prefix = import_input_prefix(
            index_key.keyspace_id,
            index_key.table_id,
            &[index_key.index_id],
            job_id,
        );
        let object_prefix = if s3_cfg.prefix.is_empty() {
            import_prefix
        } else {
            format!("{}/{}", s3_cfg.prefix, import_prefix)
        };
        let object_key = format!("{object_prefix}/part-0000");
        storage
            .put_data(&s3_cfg.bucket, &object_key, b"payload".to_vec())
            .await
            .unwrap();
        assert_eq!(
            storage
                .list_all(&s3_cfg.bucket, &object_prefix)
                .await
                .unwrap()
                .len(),
            1
        );

        assert!(meta_service.run_import_job_check_once_for_test().await);

        assert!(
            storage
                .list_all(&s3_cfg.bucket, &object_prefix)
                .await
                .unwrap()
                .is_empty()
        );

        suite.stop().await.unwrap();
    });

    Ok(())
}

fn run_external_terminal_cleanup_blocks_running_task_finish(
    terminal_state: &str,
    table_id: i64,
) -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let mut suite = E2ETestSuite::new_with_mock_cfg(
            make_default_test_meta_config(false),
            make_default_test_worker_config(false),
        )
        .await
        .unwrap();

        let index_key = IndexKey::new_nullspace(1, table_id);
        suite.must_create_index(&index_key).await.unwrap();
        suite.validate_add_index_result(&index_key).await;

        let meta_service = suite.get_meta_service();
        let worker_addr = suite.get_one_worker_node_addr();
        let tidb_task_id = format!("ImportInto/external-terminal-cleanup-{terminal_state}");
        let job_id = meta_service
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: tidb_task_id.clone(),
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                job_id: 0,
            })
            .await
            .unwrap();
        meta_service
            .meta_store
            .add_import_task(
                &tidb_task_id,
                ImportTaskParams::new(
                    "a".to_string(),
                    "z".to_string(),
                    format!("mock://import/external-terminal-{terminal_state}"),
                ),
            )
            .await
            .unwrap();
        let task = meta_service
            .meta_store
            .next_import_task(&worker_addr)
            .await
            .unwrap()
            .unwrap();

        meta_service
            .shard_task_manager
            .set_import_topology_frozen(
                index_key.keyspace_id,
                index_key.table_id,
                &[index_key.index_id],
                true,
            )
            .await;
        debug_insert_external_terminal_task(
            &suite.get_meta_config(),
            &tidb_task_id,
            terminal_state,
        )
        .await;

        assert!(meta_service.run_import_cleanup_check_once_for_test().await);
        assert_eq!(
            meta_service
                .meta_store
                .get_running_import_task_count()
                .await
                .unwrap(),
            0
        );
        assert!(
            meta_service
                .meta_store
                .get_active_import_jobs()
                .await
                .unwrap()
                .is_empty()
        );
        let progress = meta_service
            .meta_store
            .get_import_index_progress(
                index_key.keyspace_id,
                index_key.table_id,
                index_key.index_id,
            )
            .await
            .unwrap()
            .unwrap();
        assert_eq!(progress.job_id, job_id);
        assert_eq!(progress.job_status, "failed");
        assert!(
            !meta_service
                .meta_store
                .import_task_build_success(
                    task.task_id,
                    &task.owner_addr,
                    task.owner_seq,
                    ImportTaskSummary::default(),
                )
                .await
                .unwrap()
        );

        let index = suite.get_index(&index_key).await;
        assert!(!index.is_topology_frozen());

        suite.stop().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_failed_external_terminal_cleanup_blocks_running_task_finish() -> Result<()> {
    run_external_terminal_cleanup_blocks_running_task_finish("failed", 1006)
}

#[test]
fn test_reverted_external_terminal_cleanup_blocks_running_task_finish() -> Result<()> {
    run_external_terminal_cleanup_blocks_running_task_finish("reverted", 1007)
}

#[test]
fn test_recovery_seeds_running_task_and_heartbeat_reclaims_it() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;

    rt.block_on(async {
        let config = make_default_test_meta_config(false);
        Bootstrap::new_with_config(&config)
            .clear_all_meta(true)
            .await
            .unwrap();
        let meta_service = build_meta_service(&config).await;
        meta_service.start().await.unwrap();

        let worker_addr = "recovery-heartbeat-worker".to_string();
        meta_service
            .worker_node_mgr
            .block_workers(std::slice::from_ref(&worker_addr))
            .await;
        meta_service
            .worker_node_mgr
            .report_worker_node_status(WorkerNodeHeartbeatRequest {
                addr: worker_addr.clone(),
                status: Some(WorkerNodeStatus::default()),
                ..Default::default()
            })
            .await
            .unwrap();

        let tidb_task_id = "ImportInto/recovery-heartbeat-miss".to_string();
        meta_service
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: tidb_task_id.clone(),
                keyspace_id: 0,
                table_id: 5102,
                index_ids: vec![12],
                job_id: 0,
            })
            .await
            .unwrap();
        meta_service
            .meta_store
            .add_import_task(
                &tidb_task_id,
                ImportTaskParams::new(
                    "a".to_string(),
                    "z".to_string(),
                    "mock://import/recovery-heartbeat-miss".to_string(),
                ),
            )
            .await
            .unwrap();
        let task = meta_service
            .meta_store
            .next_import_task(&worker_addr)
            .await
            .unwrap()
            .unwrap();

        let worker = meta_service
            .worker_node_mgr
            .get_worker_node(&worker_addr)
            .await
            .unwrap();
        assert!(!worker.has_import_task(task.task_id).await);

        assert!(meta_service.run_import_recovery_once_for_test().await);
        assert!(worker.has_import_task(task.task_id).await);

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

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

        let recovered = meta_service
            .meta_store
            .next_import_task("recovery-reassigned-worker")
            .await
            .unwrap()
            .unwrap();
        assert_eq!(recovered.task_id, task.task_id);

        meta_service.stop().await.unwrap();
    });

    Ok(())
}