use std::collections::HashMap;

use anyhow::Result;
use log::info;
use scopeguard::defer;
use storage_backend::storage::create_backend_storage_for_test;
use tici::{
    config::{make_default_test_meta_config, make_default_test_worker_config},
    meta::writer::CompactShardReason,
    worker::mocks::{make_frag_meta, make_frag_meta_extended},
};
use tici_common::index::IndexKey;
use tici_shard::utils::frag_path_generator::FragPathGenerator;
use tici_system::readable::ReadableDuration;
use tici_test_utils::test_utils::{
    AppendExtraParam, mock_frag_path, must_finally, test_suite::E2ETestSuite,
};

fn with_runtime() -> Result<tokio::runtime::Runtime> {
    Ok(tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?)
}

#[test]
fn test_uni_compaction_can_schedule_and_apply() -> Result<()> {
    let rt = with_runtime()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    fail::cfg("mock_etcd_leader_term", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
        fail::remove("mock_etcd_leader_term");
    });

    let mut config = make_default_test_meta_config(true);
    config.shard.uni_compaction.min_compaction_frags = 3;
    config.shard.auto_compaction = false;
    config.shard.enable_split = false;
    config.shard.gc_config.gc_worker_check_interval = ReadableDuration::secs(2);
    config.shard.gc_config.safe_gc_ts_diff = ReadableDuration::millis(10);
    config.s3.use_path_style = true;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();

        let table_id = 1;
        let index_id = 1001;
        let index_key = IndexKey::new_nullspace(table_id, index_id);

        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.validate_add_index_result(&index_key).await;

        let shard_id = test_suite.get_first_shard_id().await;
        let shard = test_suite.get_first_shard().await;

        test_suite
            .must_append_delta_by_shard(&shard, "A1", 1, AppendExtraParam::default())
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A2", 2, AppendExtraParam::default())
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A3", 3, AppendExtraParam::default())
            .await;

        let (ind, shard) = test_suite
            .get_shard_and_index(&index_key, shard_id)
            .await
            .unwrap();
        let compaction = test_suite
            .get_meta_service()
            .shard_task_manager
            .debug_maybe_schedule_compaction(ind, &mut *shard.write().await)
            .await
            .unwrap()
            .first()
            .unwrap()
            .clone();

        must_finally(async || {
            test_suite
                .collect_worker_nodes()
                .await
                .first()
                .unwrap()
                .1
                .compaction_len()
                .await
                == 1
        })
        .await;

        info!("compaction={:?}", compaction);

        assert!(
            test_suite
                .do_compact_scheduled(
                    &index_key,
                    shard_id,
                    compaction.input_fragments.clone(),
                    false
                )
                .await
                .unwrap()
                .is_none(),
        );

        let storage = create_backend_storage_for_test(&config.s3).await.unwrap();
        for path in compaction.input_fragments.iter() {
            let path = format!("{}/{}", compaction.upload_location.prefix, path.frag_path);
            assert!(
                !storage
                    .list_all(&config.s3.bucket, &path)
                    .await
                    .unwrap()
                    .is_empty()
            );
        }

        test_suite
            .do_compact_finished(
                shard_id,
                compaction.clone(),
                vec![make_frag_meta(&index_key, &mock_frag_path("compacted_uni"))],
                true,
            )
            .await
            .unwrap();

        tokio::time::sleep(std::time::Duration::from_secs(5)).await;
        for path in compaction.input_fragments.iter() {
            let path = format!("{}/{}", compaction.upload_location.prefix, path.frag_path);
            let files = storage.list_all(&config.s3.bucket, &path).await.unwrap();
            for file in files {
                assert!(!FragPathGenerator::path_is_fragment_file(file.as_str()));
            }
        }

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

    Ok(())
}

#[test]
fn test_uni_compaction_chain_after_compaction() -> Result<()> {
    let rt = with_runtime()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    fail::cfg("mock_etcd_leader_term", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
        fail::remove("mock_etcd_leader_term");
    });

    let mut config = make_default_test_meta_config(true);
    config.shard.auto_compaction = false;
    config.shard.uni_compaction.max_largest_ratio = 0.5;
    config.shard.uni_compaction.min_compaction_frags = 2;
    config.shard.uni_compaction.max_compaction_iterations = 2;
    config.shard.enable_split = false;
    config.shard.gc_config.gc_worker_check_interval = ReadableDuration::secs(2);
    config.shard.gc_config.safe_gc_ts_diff = ReadableDuration::millis(10);
    config.s3.use_path_style = true;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();

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

        let shard_id = test_suite.get_first_shard_id().await;
        let shard = test_suite.get_first_shard().await;

        let big_frag = AppendExtraParam {
            size: Some(10),
            count: Some(1),
            segs_count: Some(1),
            ..Default::default()
        };
        test_suite
            .must_append_delta_by_shard(&shard, "A1", 1, big_frag)
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A2", 2, AppendExtraParam::default())
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A3", 3, AppendExtraParam::default())
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A4", 4, AppendExtraParam::default())
            .await;

        let (ind, shard) = test_suite
            .get_shard_and_index(&index_key, shard_id)
            .await
            .unwrap();
        let compaction = test_suite
            .get_meta_service()
            .shard_task_manager
            .debug_maybe_schedule_compaction(ind, &mut *shard.write().await)
            .await
            .unwrap()
            .first()
            .unwrap()
            .clone();

        must_finally(async || {
            test_suite
                .collect_worker_nodes()
                .await
                .first()
                .unwrap()
                .1
                .compaction_len()
                .await
                == 1
        })
        .await;

        let mut changes = HashMap::new();
        changes.insert("shard.auto_compaction".to_string(), "true".to_string());
        test_suite
            .get_meta_service()
            .update_config(changes)
            .unwrap();

        let range = {
            let shard = test_suite
                .get_shard_from_index(&index_key, shard_id)
                .await
                .unwrap();
            shard.read().await.range.clone()
        };
        let output =
            make_frag_meta_extended(&mock_frag_path("compacted_chain_1"), &range, 10, 1, 1);
        test_suite
            .do_compact_finished(shard_id, compaction.clone(), vec![output], true)
            .await
            .unwrap();

        must_finally(async || {
            let shard = test_suite
                .get_shard_from_index(&index_key, shard_id)
                .await
                .unwrap();
            let states = shard.read().await.progress.collect_compaction_states();
            states.len() == 1 && states[0].reason == CompactShardReason::UniChain
        })
        .await;

        let chain_state = {
            let shard = test_suite
                .get_shard_from_index(&index_key, shard_id)
                .await
                .unwrap();
            shard
                .read()
                .await
                .progress
                .collect_compaction_states()
                .first()
                .cloned()
                .unwrap()
        };
        assert_eq!(chain_state.remaining_iterations, 1);

        test_suite
            .do_compact_finished(
                shard_id,
                chain_state,
                vec![make_frag_meta(
                    &index_key,
                    &mock_frag_path("compacted_chain_2"),
                )],
                true,
            )
            .await
            .unwrap();

        must_finally(async || {
            let shard = test_suite
                .get_shard_from_index(&index_key, shard_id)
                .await
                .unwrap();
            shard.read().await.progress.compaction_state.is_empty()
        })
        .await;

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

    Ok(())
}

#[test]
fn test_uni_compaction_scheduler_respects_policy() -> Result<()> {
    let rt = with_runtime()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();
    fail::cfg("mock_etcd_leader_term", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
        fail::remove("mock_etcd_leader_term");
    });

    let mut config = make_default_test_meta_config(true);
    config.shard.uni_compaction.min_compaction_frags = 3;
    config.shard.enable_split = false;
    config.shard.gc_config.gc_worker_check_interval = ReadableDuration::secs(2);
    config.shard.gc_config.safe_gc_ts_diff = ReadableDuration::millis(10);
    config.s3.use_path_style = true;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();

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

        let shard_id = test_suite.get_first_shard_id().await;
        let shard = test_suite.get_first_shard().await;

        test_suite
            .must_append_delta_by_shard(&shard, "A1", 1, AppendExtraParam::default())
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A2", 2, AppendExtraParam::default())
            .await;
        test_suite
            .must_append_delta_by_shard(&shard, "A3", 3, AppendExtraParam::default())
            .await;

        must_finally(async || {
            let shard = test_suite
                .get_shard_from_index(&index_key, shard_id)
                .await
                .unwrap();
            shard.read().await.progress.compaction_state.len() == 1
        })
        .await;

        let shard = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        let state = shard
            .read()
            .await
            .progress
            .collect_compaction_states()
            .first()
            .cloned()
            .unwrap();

        test_suite
            .do_compact_finished(
                shard_id,
                state,
                vec![make_frag_meta(
                    &index_key,
                    &mock_frag_path("compacted_uni_auto"),
                )],
                true,
            )
            .await
            .unwrap();

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

    Ok(())
}