use anyhow::{Ok, Result};
use log::info;
use scopeguard::defer;
use storage_backend::storage::create_backend_storage_for_test;
use tici::{
    SplittableFrag,
    config::{make_default_test_meta_config, make_default_test_worker_config},
    meta::{metrics::SHARD_EVENT_COUNTER_VEC, writer::build_compact_frag_req},
    worker::mocks::{make_frag_meta, make_frag_meta_by_range},
};
use tici_common::index::IndexKey;
use tici_grpc::{
    ErrorCode, WorkerNodeCompactionStatus, WorkerNodeHeartbeatRequest, WorkerNodeStatus,
};
use tici_shard::utils::frag_path_generator::FragPathGenerator;
use tici_system::readable::ReadableDuration;
use tici_test_utils::test_utils::{create_cdc_file_name, test_suite::E2ETestSuite, *};

async fn prepare_compact_shard(
    test_suite: &mut E2ETestSuite,
    index_key: &IndexKey,
    shard_id: u64,
) -> Result<Vec<SplittableFrag>> {
    let _ = index_key;
    let _ = shard_id;
    let shard_id = test_suite.get_first_shard_id().await;
    let shard = test_suite.get_first_shard().await;

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

    Ok(vec![frag1, frag2, frag3]
        .into_iter()
        .map(|f| SplittableFrag::new(f, shard_id, 1))
        .collect::<Vec<_>>())
}

// Test if a compaction can be scheduled and applied, and if old fragments can
// be deleted.
fn test_compaction(drop_add_shard: bool) -> 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();

    if drop_add_shard {
        fail::cfg("worker_service_mock_drop_add_shard", "return").unwrap();
    }

    fail::cfg("mock_etcd_leader_term", "return").unwrap();
    defer!({
        fail::remove("worker_service_mock_meta_service_url");
        fail::remove("to_add_shard_must_be_empty");
        fail::remove("worker_service_mock_drop_add_shard");
        fail::remove("mock_etcd_leader_term");
    });

    info!(
        "start test_compaction drop_add_shard={} ...",
        drop_add_shard
    );

    let mut config = make_default_test_meta_config(true);
    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;
    config.shard.uni_compaction.max_compaction_iterations = 1;
    config.shard.uni_compaction.max_largest_ratio = 1.0;
    config.shard.uni_compaction.min_compaction_frags = 1;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();
        if !drop_add_shard {
            test_suite
                .get_worker_node_handle("127.0.0.1:8510")
                .shard_write_manager
                .as_mock()
                .check_to_add_shard
                .store(true, std::sync::atomic::Ordering::Relaxed);
        }

        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 v = prepare_compact_shard(&mut test_suite, &index_key, shard_id)
            .await
            .unwrap();

        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 scheduled
        must_finally(async || {
            test_suite
                .collect_worker_nodes()
                .await
                .first()
                .unwrap()
                .1
                .compaction_len()
                .await
                == 1
        })
        .await;

        info!("compaction={:?}", compaction);
        // Concurrent compaction limit
        assert!(
            test_suite
                .do_compact_scheduled(&index_key, shard_id, v, 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()
            );
        }

        let shard = test_suite
            .collect_shards_from_restore()
            .await
            .first()
            .cloned()
            .unwrap();
        // CompactionState is not persisted.
        assert_eq!(shard.progress.compaction_state.len(), 0);

        test_suite
            .do_compact_finished(
                shard_id,
                compaction.clone(),
                vec![
                    make_frag_meta(&index_key, &mock_frag_path("compacted")),
                    make_frag_meta(&index_key, &mock_frag_path("compacted2")),
                ],
                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()));
            }
        }

        let shard = test_suite
            .collect_shards_from_restore()
            .await
            .first()
            .cloned()
            .unwrap();
        assert_eq!(shard.progress.compaction_state.len(), 0);

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

    Ok(())
}

#[test]
fn test_abort_compaction() -> 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();

    // Use these failpoints to make the worker node returns no compactions in
    // heartbeat. So the meta service will consider the compaction is lost, and
    // abort it.
    fail::cfg("worker_service_mock_drop_compact_shard", "return").unwrap();
    fail::cfg("worker_service_mock_drop_compaction", "return").unwrap();

    defer!({
        fail::remove("worker_service_mock_drop_compact_shard");
        fail::remove("worker_service_mock_meta_service_url");
        fail::remove("worker_service_mock_drop_compaction");
    });

    let mut config = make_default_test_meta_config(true);
    config.shard.uni_compaction.max_compaction_iterations = 1;
    config.shard.uni_compaction.max_largest_ratio = 1.0;
    config.shard.uni_compaction.min_compaction_frags = 1;
    config.shard.auto_compaction = false;
    config.shard.enable_split = false;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, 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 v = prepare_compact_shard(&mut test_suite, &index_key, shard_id)
            .await
            .unwrap();

        let (ind, shard) = test_suite
            .get_shard_and_index(&index_key, shard_id)
            .await
            .unwrap();
        let compaction_aborted = shard.read().await.stats.get_compaction_abort();
        let _ = test_suite
            .get_meta_service()
            .shard_task_manager
            .debug_maybe_schedule_compaction(ind, &mut *shard.write().await)
            .await
            .unwrap()
            .first()
            .unwrap();

        must_finally(async || {
            shard.read().await.stats.get_compaction_abort() == compaction_aborted + 1
        })
        .await;

        test_suite.check_compact_aborted(shard_id, v).await.unwrap();

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

#[test]
fn test_compaction_f() -> Result<()> {
    test_compaction(false).unwrap();
    Ok(())
}

#[test]
fn test_compaction_t() -> Result<()> {
    test_compaction(true).unwrap();
    Ok(())
}

// If worker node is dead, compaction will be aborted on meta's side.
#[test]
fn test_remove_node() -> 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 config = make_default_test_meta_config(true);
    config.shard.uni_compaction.max_compaction_iterations = 1;
    config.shard.uni_compaction.max_largest_ratio = 1.0;
    config.shard.uni_compaction.min_compaction_frags = 1;
    config.shard.auto_compaction = false;
    config.shard.enable_split = false;
    // Do not reschedule shard
    config.shard.shard_no_heartbeat_tick = 999999;
    config.shard.enable_patrol();
    config.worker_node.worker_node_down_timeout = ReadableDuration::secs(3);
    config.worker_node.health_check_interval = ReadableDuration::millis(500);
    let wn_cfg = make_default_test_worker_config(true);

    rt.block_on(async {
        let mut test_suite = E2ETestSuite::new_with_mock_cfg(config.clone(), wn_cfg)
            .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 v = prepare_compact_shard(&mut test_suite, &index_key, shard_id)
            .await
            .unwrap();

        let (ind, shard) = test_suite
            .get_shard_and_index(&index_key, shard_id)
            .await
            .unwrap();
        let compaction_aborted = shard.read().await.stats.get_compaction_abort();
        let _ = test_suite
            .get_meta_service()
            .shard_task_manager
            .debug_maybe_schedule_compaction(ind, &mut *shard.write().await)
            .await
            .unwrap()
            .first()
            .unwrap();

        test_suite
            .get_worker_node_handle("127.0.0.1:8510")
            .shard_write_manager
            .as_mock()
            .skip_heartbeat
            .store(true, std::sync::atomic::Ordering::Relaxed);

        tokio::time::sleep(std::time::Duration::from_millis(
            config.worker_node.worker_node_down_timeout.as_millis() * 2,
        ))
        .await;

        must_finally(async || {
            shard.read().await.stats.get_compaction_abort() == compaction_aborted + 1
        })
        .await;

        test_suite.check_compact_aborted(shard_id, v).await.unwrap();

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

#[test]
fn test_abort_one_compaction_then_finish_other() -> 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;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg, 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 = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;

        for i in 0..6 {
            test_suite
                .must_append_delta_by_shard(
                    &shard,
                    &format!("L0_{}", i),
                    1,
                    AppendExtraParam::default(),
                )
                .await;
        }

        let delta_frags = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap()
            .read()
            .await
            .manifest
            .delta_cloned();
        assert_eq!(delta_frags.len(), 6);

        let input_first = delta_frags[..3].to_vec();
        let input_second = delta_frags[3..].to_vec();

        let _state_first = test_suite
            .do_compact_scheduled(&index_key, shard_id, input_first, true)
            .await
            .unwrap()
            .unwrap();
        let state_second = test_suite
            .do_compact_scheduled(&index_key, shard_id, input_second, true)
            .await
            .unwrap()
            .unwrap();

        let shard = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        let compaction_aborted = shard.read().await.stats.get_compaction_abort();

        let index = test_suite.get_index(&index_key).await;
        let shard_info = test_suite.get_shard_from_restore(shard_id).await.unwrap();
        let shard_header = shard_info.get_header();

        // Mock aborting compaction 1
        let compact_req_second = build_compact_frag_req(
            shard_header,
            &state_second,
            String::new(),
            index.table_info.clone(),
            index.index_info.clone(),
        );
        let heartbeat = WorkerNodeHeartbeatRequest {
            addr: state_second.worker.clone(),
            status: Some(WorkerNodeStatus {
                cpu_usage: 0.0,
                memory_usage: 0,
            }),
            shards: vec![],
            compactions: vec![WorkerNodeCompactionStatus {
                req: Some(compact_req_second),
            }],
            import_writes: vec![],
        };

        for _ in 0..4 {
            test_suite
                .get_meta_service()
                .worker_node_mgr
                .report_worker_node_status(heartbeat.clone())
                .await
                .unwrap();
        }

        must_finally(async || {
            shard.read().await.progress.compaction_state.len() == 1
                && shard.read().await.stats.get_compaction_abort() == compaction_aborted + 1
        })
        .await;

        let output_second = vec![make_frag_meta(&index_key, &mock_frag_path("L1_second"))];
        // The second compaction can still splice in place even when the first one
        // is aborted, because tiered compaction operates on a contiguous window.
        test_suite
            .do_compact_finished(shard_id, state_second, output_second, true)
            .await
            .unwrap();

        let shard = test_suite.get_shard_from_restore(shard_id).await.unwrap();
        assert_eq!(shard.progress.compaction_state.len(), 0);

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

/// Test that overlapping compactions are properly detected and rejected at
/// scheduling time, and that apply fails when fragments have changed (race
/// condition).
#[test]
fn test_overlapping_compactions_rejected_and_fail_on_apply() -> 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 config = make_default_test_meta_config(true);
    config.shard.auto_compaction = false;
    config.shard.enable_split = false;

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(config, 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 = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;

        // Create 4 fragments
        for i in 0..4 {
            test_suite
                .must_append_delta_by_shard(
                    &shard,
                    &format!("Frag_{}", i),
                    1,
                    AppendExtraParam::default(),
                )
                .await;
        }

        // Get all fragments from manifest (level-agnostic)
        let all_frags = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap()
            .read()
            .await
            .manifest
            .get_fragments()
            .to_vec();
        assert_eq!(all_frags.len(), 4, "Should have 4 fragments initially");

        // Step 1: Schedule compaction C1 with first 2 fragments
        let input_c1 = all_frags[0..2].to_vec();
        let state_c1 = test_suite
            .do_compact_scheduled(&index_key, shard_id, input_c1.clone(), true)
            .await
            .unwrap()
            .unwrap();

        // Step 2: Try to schedule overlapping compaction C2
        // C2 overlaps with C1 because it includes fragment at index 1 (endpoint)
        let input_c2 = vec![all_frags[1].clone(), all_frags[2].clone()];
        let result_c2 = test_suite
            .do_compact_scheduled(&index_key, shard_id, input_c2, false)
            .await
            .unwrap();

        // Step 3: Verify C2 was rejected due to overlap (endpoint detection works)
        assert!(
            result_c2.is_none(),
            "C2 should be rejected due to overlap with C1"
        );

        // Step 4: Verify only C1 exists in compaction_state
        let shard_snap = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        assert_eq!(shard_snap.read().await.progress.compaction_state.len(), 1);

        // Step 5: Complete C1 successfully
        let output_c1 = vec![make_frag_meta(&index_key, &mock_frag_path("compacted_0_1"))];
        test_suite
            .do_compact_finished(shard_id, state_c1, output_c1, true)
            .await
            .unwrap();

        // Verify C1 is removed from compaction_state
        let shard_snap = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        assert_eq!(shard_snap.read().await.progress.compaction_state.len(), 0);

        // Step 6: Get updated fragment list after C1 completed
        // Manifest now has 3 fragments: compaction output + remaining 2 fragments
        let new_frags = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap()
            .read()
            .await
            .manifest
            .get_fragments()
            .to_vec();

        assert_eq!(
            new_frags.len(),
            3,
            "Should have 3 fragments after C1 completed"
        );

        // Step 7: Schedule C2 with current valid fragments
        let input_c2_new = new_frags[1..3].to_vec();
        let state_c2 = test_suite
            .do_compact_scheduled(&index_key, shard_id, input_c2_new.clone(), true)
            .await
            .unwrap()
            .unwrap();

        // Verify C2 is scheduled
        let shard_snap = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        assert_eq!(
            shard_snap.read().await.progress.compaction_state.len(),
            1,
            "C2 should be scheduled"
        );

        // Step 8: Simulate race condition - try to finish a compaction with wrong
        // inputs This simulates: compaction was scheduled with inputs A, but we
        // try to finish with inputs B Use original fragment at index 1 which
        // was consumed by C1 and no longer exists
        let phantom_input = vec![all_frags[1].clone()]; // This fragment no longer in manifest
        let phantom_output = vec![make_frag_meta(
            &index_key,
            &mock_frag_path("phantom_output"),
        )];

        // Create a fake state with stale fragments but same key as C2
        let mut phantom_state = state_c2.clone();
        phantom_state.input_fragments = phantom_input.clone();

        // Try to finish C2 but with wrong inputs - should be rejected by splice
        test_suite
            .do_compact_finished_fail(shard_id, phantom_state, phantom_output)
            .await
            .unwrap();

        // Step 9: Verify C2 is still scheduled (phantom finish failed, didn't affect
        // state)
        let shard_snap = test_suite
            .get_shard_from_index(&index_key, shard_id)
            .await
            .unwrap();
        assert_eq!(
            shard_snap.read().await.progress.compaction_state.len(),
            1,
            "C2 should still be in progress after failed finish"
        );

        // Verify manifest still has the correct 3 fragments (phantom didn't modify it)
        let final_frags = shard_snap.read().await.manifest.get_fragments().to_vec();
        assert_eq!(
            final_frags.len(),
            3,
            "Manifest should still have 3 fragments"
        );

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

/// When manifest has more than max_fragments, append fails and schedules
/// MaybeSplitOrCompact(StallTooManyFrags).
#[test]
fn test_append_fail_too_many_frags_schedules_stall_compaction() -> 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 config = make_default_test_meta_config(true);
    config.shard.max_fragments = 5;
    config.shard.auto_compaction = false;
    config.shard.enable_split = false;

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

        test_suite
            .get_worker_node_handle("127.0.0.1:8510")
            .shard_write_manager
            .as_mock()
            .check_to_add_shard
            .store(true, std::sync::atomic::Ordering::Relaxed);

        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 = test_suite.get_first_shard().await;
        let addr = test_suite.get_one_worker_node_addr();

        // Append 6 frags (succeed) so manifest has 6 > max_fragments(5)
        for i in 0..6 {
            let shard_snap = test_suite
                .get_shard_from_index(&index_key, shard.shard_id)
                .await
                .unwrap()
                .read()
                .await
                .clone();
            let frag = make_frag_meta_by_range(
                &mock_frag_path(format!("toomany_{}", i)),
                &shard_snap.range.debug_derive_frag_range(),
            );
            test_suite
                .do_append_delta_frag(
                    addr.as_str(),
                    shard.shard_id,
                    &index_key,
                    create_cdc_file_name(i).to_string(),
                    vec![frag],
                )
                .await
                .unwrap();
        }

        // 7th append should fail (manifest has 6 frags, 6 > 5)
        let shard_snap = test_suite
            .get_shard_from_index(&index_key, shard.shard_id)
            .await
            .unwrap()
            .read()
            .await
            .clone();
        let frag = make_frag_meta_by_range(
            &mock_frag_path("toomany_6"),
            &shard_snap.range.debug_derive_frag_range(),
        );
        let (_, resp) = test_suite
            .do_append_delta_frag(
                addr.as_str(),
                shard.shard_id,
                &index_key,
                create_cdc_file_name(7).to_string(),
                vec![frag],
            )
            .await
            .unwrap();
        assert_eq!(
            resp.status,
            ErrorCode::Throttled.to_i32(),
            "append should return throttled when manifest has too many frags"
        );

        // Wait for MaybeSplitOrCompact(StallTooManyFrags) to be processed
        must_finally(async || {
            SHARD_EVENT_COUNTER_VEC
                .get_metric_with_label_values(&["append_fail_toomanyfrags"])
                .iter()
                .next()
                .map_or(0, |m| m.get())
                >= 1
        })
        .await;

        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        test_suite.check_meta_subscriber().await;
        test_suite.stop().await.unwrap();
    });
    Ok(())
}