use std::collections::HashSet;

use anyhow::{Ok, Result};
use log::info;
use scopeguard::defer;
use tici::{
    HandleRange,
    config::{make_default_test_meta_config, make_default_test_worker_config},
    meta::metrics::SHARD_EVENT_COUNTER_VEC,
    worker::mocks::{make_frag_meta, make_frag_meta_by_range},
};
use tici_common::{
    codec::{table_end_key, table_raw_key, table_start_key},
    index::IndexKey,
};
use tici_import::ImportJobMeta;
use tici_test_utils::test_utils::{test_suite::E2ETestSuite, *};

#[test]
fn test_merge_shard_basic() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(true);
    cfg.shard.auto_compaction = false;
    cfg.shard.enable_split = false;
    cfg.shard.enable_merge = false;
    cfg.shard.enable_patrol();

    rt.block_on(async {
        let mut suite = E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true))
            .await
            .unwrap();
        let index_key = IndexKey::new_nullspace(1, 1001);
        suite.must_create_index(&index_key).await.unwrap();
        let index = suite.get_index(&index_key).await;
        let shard = index.select_first_shard().await.unwrap();
        let shard_unwrap = shard.read().await.clone();
        let addr = suite.get_one_worker_node_addr();
        suite
            .must_append_base_frag(
                &addr,
                1,
                &index_key,
                vec![make_frag_meta(&index_key, &mock_frag_path("b"))],
            )
            .await
            .unwrap();
        suite
            .must_append_delta_by_shard(
                &shard_unwrap,
                "g",
                1,
                AppendExtraParam {
                    upload_to_s3: false,
                    do_manifest_check: true,
                    ..Default::default()
                },
            )
            .await;

        // Can not split on error table range.
        let index_key_2 = IndexKey::new_nullspace(2, 0);
        assert!(
            suite
                .must_manual_split_shard(
                    &index_key,
                    shard.read().await.range.sortable_start(),
                    vec![
                        HandleRange::new(
                            table_start_key(&index_key_2).to_vec(),
                            table_raw_key(&index_key_2, b"A"),
                        ),
                        HandleRange::new(
                            table_raw_key(&index_key_2, b"A"),
                            table_end_key(&index_key_2).to_vec(),
                        ),
                    ],
                    0,
                )
                .await
                .is_err()
        );

        suite
            .must_split_shard_simple(&index_key, shard.read().await.range.clone(), Some(b"A"))
            .await
            .unwrap();
        let shards = index.clone_shards().await;
        assert_eq!(shards.len(), 2);
        for (i, shard) in shards.iter().enumerate() {
            let shard_unwrap = shard.read().await.clone();
            let shard_id = shard_unwrap.shard_id;
            let range = shard_unwrap.range.debug_derive_frag_range();
            let addr = suite.get_one_worker_node_addr();
            suite
                .must_append_base_frag(
                    &addr,
                    shard_id,
                    &index_key,
                    vec![make_frag_meta_by_range(
                        &mock_frag_path(format!("f_b_{}", i)),
                        &range,
                    )],
                )
                .await
                .unwrap();
            suite
                .must_append_delta_by_shard(
                    &shard_unwrap,
                    &format!("f_{}", i),
                    i as u64 + 100,
                    AppendExtraParam {
                        upload_to_s3: false,
                        do_manifest_check: true,
                        ..Default::default()
                    },
                )
                .await;
        }
        let wn = suite
            .get_meta_service()
            .worker_node_mgr
            .get_worker_node(&suite.get_worker_node_template_addr())
            .await
            .unwrap();
        // Have overlap in fragments
        suite
            .must_not_merge_shard_for_overlap(shards.as_slice(), wn.clone())
            .await;

        for shard in shards.iter() {
            shard
                .write()
                .await
                .manifest
                .remove_frag(&mock_frag_path("b"));
            shard
                .write()
                .await
                .manifest
                .remove_frag(&mock_frag_path("g"));
        }
        suite
            .must_merge_shard(shards.as_slice(), wn.addr.clone())
            .await;

        let shards = index.clone_shards().await;
        assert_eq!(shards.len(), 1);
        let manifest = shards.first().unwrap().read().await.manifest.clone();
        for f in manifest.get_fragments() {
            assert_eq!(f.p, 1.0);
        }
        assert_eq!(manifest.base_slice().len(), 2);
        assert_eq!(manifest.delta_cloned().len(), 2);
        assert_eq!(
            shards
                .first()
                .unwrap()
                .read()
                .await
                .progress
                .cdc_s3_last_file,
            create_cdc_file_name(100)
        );
        // Wait a round of ShardStatsPatrol
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;

        suite.stop_worker().await.unwrap();
        suite.check_meta_subscriber().await;
        suite.validate_size_histo().await;
        suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_merge_shard_strategy() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(true);
    cfg.shard.auto_compaction = false;
    cfg.shard.enable_split = false;
    cfg.shard.enable_merge = true;
    let patrol_secs = 2;
    cfg.shard.disable_patrol();
    cfg.shard.merge_strategy.write_qps = 0.5; // merge happens when less than 1 write in 2s

    rt.block_on(async {
        let index_key = IndexKey::new_nullspace(1, 1001);
        let mut suite = E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true))
            .await
            .unwrap();
        suite.must_create_index(&index_key).await.unwrap();
        let index = suite.get_index(&index_key).await;
        let shard = index.select_first_shard().await.unwrap();
        let shard_unwrap = shard.read().await.clone();
        suite
            .must_append_delta_by_shard(
                &shard_unwrap,
                "g",
                1,
                AppendExtraParam {
                    do_manifest_check: true,
                    upload_to_s3: false,
                    ..Default::default()
                },
            )
            .await;
        // Do not check index_id.
        let same_table_different_index = IndexKey::new_nullspace(1, 0);
        suite
            .must_manual_split_shard(
                &index_key,
                shard.read().await.range.sortable_start(),
                vec![
                    HandleRange::new(
                        table_start_key(&same_table_different_index).to_vec(),
                        table_raw_key(&same_table_different_index, b"A"),
                    ),
                    HandleRange::new(
                        table_raw_key(&same_table_different_index, b"A"),
                        table_end_key(&same_table_different_index).to_vec(),
                    ),
                ],
                0,
            )
            .await
            .unwrap();
        let shards = index.clone_shards().await;
        assert_eq!(shards.len(), 2);

        // Reject due to overlap
        assert_eq!(
            suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await
                .unwrap()
                .sched_try_merge_count,
            0
        );
        tokio::time::sleep(std::time::Duration::from_secs(patrol_secs * 2)).await;
        assert_eq!(
            suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await
                .unwrap()
                .sched_try_merge_count,
            2
        );

        assert_eq!(
            SHARD_EVENT_COUNTER_VEC
                .get_metric_with_label_values(&["merge_succeed"])
                .iter()
                .next()
                .unwrap()
                .get(),
            0
        );
        must_finally(async || {
            SHARD_EVENT_COUNTER_VEC
                .get_metric_with_label_values(&["merge_failed_overlap"])
                .iter()
                .next()
                .unwrap()
                .get()
                == 1
        })
        .await;

        for s in shards.iter() {
            let shard_id = s.read().await.shard_id;
            suite
                .must_compact_all(&index_key, shard_id, "c_")
                .await
                .unwrap();
        }

        info!("rerun merge");
        assert_eq!(
            suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await
                .unwrap()
                .sched_try_merge_count,
            2
        );
        must_finally(async || {
            SHARD_EVENT_COUNTER_VEC
                .get_metric_with_label_values(&["merge_succeed"])
                .iter()
                .next()
                .unwrap()
                .get()
                == 1
        })
        .await;

        suite.stop_worker().await.unwrap();
        suite.check_meta_subscriber().await;
        suite.validate_size_histo().await;
        suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_merge_shard_suppressed_during_active_import_job() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(true);
    cfg.shard.auto_compaction = false;
    cfg.shard.enable_split = false;
    cfg.shard.enable_merge = true;
    cfg.shard.disable_patrol();
    cfg.shard.merge_strategy.write_qps = 0.5;
    let patrol_secs = 2;

    rt.block_on(async {
        let index_key = IndexKey::new_nullspace(1, 1001);
        let mut suite = E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true))
            .await
            .unwrap();
        suite.must_create_index(&index_key).await.unwrap();
        let index = suite.get_index(&index_key).await;
        let shard = index.select_first_shard().await.unwrap();
        let shard_unwrap = shard.read().await.clone();
        suite
            .must_append_delta_by_shard(
                &shard_unwrap,
                "g",
                1,
                AppendExtraParam {
                    do_manifest_check: true,
                    upload_to_s3: false,
                    ..Default::default()
                },
            )
            .await;

        let same_table_different_index = IndexKey::new_nullspace(1, 0);
        suite
            .must_manual_split_shard(
                &index_key,
                shard.read().await.range.sortable_start(),
                vec![
                    HandleRange::new(
                        table_start_key(&same_table_different_index).to_vec(),
                        table_raw_key(&same_table_different_index, b"A"),
                    ),
                    HandleRange::new(
                        table_raw_key(&same_table_different_index, b"A"),
                        table_end_key(&same_table_different_index).to_vec(),
                    ),
                ],
                0,
            )
            .await
            .unwrap();
        assert_eq!(index.clone_shards().await.len(), 2);

        tokio::time::sleep(std::time::Duration::from_secs(patrol_secs * 2)).await;
        let job_id = suite
            .get_meta_service()
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: "merge-blocked-by-import".to_string(),
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                job_id: 0,
            })
            .await
            .unwrap();
        suite
            .get_meta_service()
            .shard_task_manager
            .set_import_topology_frozen(
                index_key.keyspace_id,
                index_key.table_id,
                &[index_key.index_id],
                true,
            )
            .await;

        assert_eq!(
            suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await
                .unwrap()
                .sched_try_merge_count,
            0
        );

        suite
            .get_meta_service()
            .meta_store
            .finish_import_job(&[job_id])
            .await
            .unwrap();
        suite
            .get_meta_service()
            .shard_task_manager
            .set_import_topology_frozen(
                index_key.keyspace_id,
                index_key.table_id,
                &[index_key.index_id],
                false,
            )
            .await;

        assert_eq!(
            suite
                .get_meta_service()
                .shard_task_manager
                .get_shard_patrol()
                .debug_run_one(0)
                .await
                .unwrap()
                .sched_try_merge_count,
            2
        );

        suite.stop_worker().await.unwrap();
        suite.check_meta_subscriber().await;
        suite.stop_meta().await.unwrap();
    });

    Ok(())
}

/// Verifies that when worker 1 has shard 1 and worker 2 has shard 2, after
/// merge succeeds, both worker 1 and worker 2 should have shard 1 and shard 2
/// removed (i.e., deletable).
#[test]
fn test_merge_shard_two_workers_removes_old_shards() -> Result<()> {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let worker_addr_1 = "127.0.0.1:8510".to_string();
    let worker_addr_2 = "127.0.0.1:8556".to_string();

    let mut cfg = make_default_test_meta_config(true);
    cfg.shard.auto_compaction = false;
    cfg.shard.enable_split = false;
    cfg.shard.enable_merge = false;
    cfg.shard.writer_scheduling_strategy = "round-robin".to_string();
    cfg.shard.enable_patrol();

    rt.block_on(async {
        let mut suite = E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true))
            .await
            .unwrap();
        suite
            .add_mocked_worker_node(Some(worker_addr_2.clone()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        suite.must_create_index(&index_key).await.unwrap();
        let index = suite.get_index(&index_key).await;
        let shard = index.select_first_shard().await.unwrap();
        suite
            .must_split_shard_simple(&index_key, shard.read().await.range.clone(), Some(b"A"))
            .await
            .unwrap();
        let shards = index.clone_shards().await;
        assert_eq!(shards.len(), 2);

        let mut shards_info: Vec<_> = Vec::new();
        for s in shards.iter() {
            shards_info.push(s.read().await.clone());
        }
        shards_info.sort_by(|a, b| a.range.sortable_start().cmp(&b.range.sortable_start()));
        let left_shard = shards_info[0].clone();
        let right_shard = shards_info[1].clone();
        let shard_id_1 = left_shard.shard_id;
        let shard_id_2 = right_shard.shard_id;
        info!(
            "shard1 writer={}, shard2 writer={}",
            left_shard.writer, right_shard.writer
        );

        suite
            .must_merge_shard(shards.as_slice(), worker_addr_1.clone())
            .await;

        let shards_after = index.clone_shards().await;
        info!(
            "finished merge: {}+{} => {}",
            shard_id_1,
            shard_id_2,
            shards_after.first().unwrap().read().await.shard_id
        );
        assert_eq!(shards_after.len(), 1);

        // Wait for heartbeat so workers process remove_shard.
        tokio::time::sleep(std::time::Duration::from_secs(2)).await;

        let worker_1_shards = suite.collect_shards_by_worker_node(&worker_addr_1).await;
        let worker_2_shards = suite.collect_shards_by_worker_node(&worker_addr_2).await;

        let worker_1_shard_ids: HashSet<u64> = worker_1_shards.iter().map(|(id, _)| *id).collect();
        let worker_2_shard_ids: HashSet<u64> = worker_2_shards.iter().map(|(id, _)| *id).collect();

        assert!(
            !worker_1_shard_ids.contains(&shard_id_1) && !worker_2_shard_ids.contains(&shard_id_2),
            "worker 1 should have removed shard {}/{} after merge, worker_1_shards={:?}",
            shard_id_1,
            shard_id_2,
            worker_1_shard_ids
        );
        assert!(
            !worker_2_shard_ids.contains(&shard_id_2),
            "worker 2 should have removed shard {}/{} after merge, worker_2_shards={:?}",
            shard_id_1,
            shard_id_2,
            worker_2_shard_ids
        );

        suite.stop_worker().await.unwrap();
        suite.check_meta_subscriber().await;
        suite.stop_meta().await.unwrap();
    });

    Ok(())
}