// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2023 PingCAP, Inc.

use std::{thread, time::Duration};

use futures::executor::block_on;
use pd_client::PdClient;
use test_cloud_server::{ServerCluster, must_wait};
use tikv_util::{
    config::{ReadableDuration, ReadableSize},
    defer,
};

use super::{i_to_key, i_to_keyspace_key, i_to_val};
use crate::cases::alloc_node_id;

#[test]
fn test_retry_failed_flush() {
    for flush_fp in [
        "kvengine_flush_normal",
        "rfstore::util::alloc_id_async_timeout",
    ] {
        let node_id = alloc_node_id();
        let mut cluster = ServerCluster::new(vec![node_id], |_, cfg| {
            cfg.rocksdb.writecf.write_buffer_size = ReadableSize::kb(16);
        });

        let mut client = cluster.new_client();
        let region_id = client.get_region_id(&[]);
        fail::cfg(flush_fp, "return").unwrap();
        client.put_kv(0..100, i_to_key, i_to_val);
        client.put_kv(100..200, i_to_key, i_to_val);
        thread::sleep(Duration::from_secs(1));
        assert_eq!(
            cluster
                .get_kvengine(node_id)
                .get_shard_stat(region_id)
                .l0_table_count,
            0
        );
        fail::remove(flush_fp);
        must_wait(
            || {
                cluster
                    .get_kvengine(node_id)
                    .get_shard_stat(region_id)
                    .l0_table_count
                    > 0
            },
            3,
            || "failed to flush memtables in time".to_string(),
        );
        cluster.stop();
    }
}

#[test]
fn test_resume_failed_compaction() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, cfg| {
        cfg.rocksdb.writecf.write_buffer_size = ReadableSize::kb(4);
        cfg.raft_store.peer_long_check_interval = ReadableDuration::millis(500);
    });
    let compact_fp = "kvengine_compact";
    fail::cfg(compact_fp, "return").unwrap();
    let mut client = cluster.new_client();
    let region_id = client.get_region_id(&[]);
    for i in 0..10 {
        client.put_kv(i * 100..(i + 1) * 100, i_to_key, i_to_val);
    }
    let kv = cluster.get_kvengine(node_id);
    let stats = kv.get_shard_stat(region_id);
    assert!(stats.compaction_score > 1.0);
    fail::remove(compact_fp);
    must_wait(
        || kv.get_shard_stat(region_id).compaction_score == 0.0,
        10,
        || "failed to resume compaction in time".to_string(),
    );
    cluster.stop();
}

// A failed file-GC compaction must be atomic: it drops nothing on failure and
// the leader tick keeps retrying until it succeeds. `kvengine_compact` fails
// every compaction (it returns before the priority match, so file-GC is
// included too), so the tombstone-heavy file stays intact while the failpoint
// is set, then drops once it is removed.
#[test]
fn test_resume_failed_file_gc_compaction() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, cfg| {
        cfg.raft_store.file_gc_compaction_check_tick_interval = ReadableDuration::millis(100);
        cfg.raft_store.peer_long_check_interval = ReadableDuration::millis(500);
        cfg.kvengine.file_gc_compaction_min_interval = ReadableDuration::secs(0);
    });
    // Use a registered keyspace so the GC safe point can be set through PD
    // below, rather than injected into the engine cache directly.
    let keyspace_id =
        block_on(cluster.create_keyspace(&Default::default(), Duration::from_secs(10)));
    let pd_client = cluster.get_pd_client();

    let compact_fp = "kvengine_compact";
    fail::cfg(compact_fp, "return").unwrap();

    let mut client = cluster.new_client();
    client.split_keyspace(keyspace_id);
    let gen_key = i_to_keyspace_key(keyspace_id);
    // Put then delete within one mem-table generation, flush once: a single
    // tombstone-heavy L0 that only the file-GC tick can compact.
    client.put_kv(0..100, &gen_key, i_to_val);
    client.del_kv(0..100, &gen_key);
    let region_id = client.get_region_id(&gen_key(0));
    cluster.flush_memtable(region_id).unwrap();
    assert!(cluster.wait_for_memtable_flushed(region_id, Duration::from_secs(10)));
    let kv = cluster.get_kvengine(node_id);
    must_wait(
        || kv.get_shard_stat(region_id).l0_table_count > 0,
        10,
        || "flushed L0 table not visible".to_string(),
    );
    let tombs_before = kv.get_shard_stat(region_id).tombs;
    assert!(tombs_before > 0);

    // Advance the GC safe point past the deletes so the tick wants to compact.
    // Set it through PD instead of injecting the engine cache directly: the
    // store's periodic gc-states refresh re-fetches from PD, so a PD-backed
    // safe point survives the whole test instead of being wiped mid-run.
    let safe_point = client.get_ts();
    block_on(pd_client.advance_txn_safe_point(keyspace_id, safe_point)).unwrap();
    block_on(pd_client.advance_gc_safe_point(keyspace_id, safe_point)).unwrap();
    block_on(cluster.update_gc_states_immediately()).unwrap();

    // Every triggered compaction fails, so no data is dropped.
    thread::sleep(Duration::from_secs(2));
    assert_eq!(
        kv.get_shard_stat(region_id).tombs,
        tombs_before,
        "a failed file-GC compaction must not drop data"
    );

    // Once compaction can succeed, the tick resumes and the tombstones drop.
    fail::remove(compact_fp);
    must_wait(
        || kv.get_shard_stat(region_id).tombs == 0,
        60,
        || {
            format!(
                "file-GC did not resume after failure: tombs={}",
                kv.get_shard_stat(region_id).tombs
            )
        },
    );
    cluster.stop();
}

// While disk space is critically low the file-GC tick must bail at the
// low-space gate (a reclaim writes new SSTs before deleting inputs, worsening
// pressure), and it must resume once space recovers -- so the gate never wedges
// file-GC off permanently.
#[test]
fn test_file_gc_skipped_on_critical_low_space() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, cfg| {
        cfg.raft_store.file_gc_compaction_check_tick_interval = ReadableDuration::millis(100);
        cfg.kvengine.file_gc_compaction_min_interval = ReadableDuration::secs(0);
    });
    let keyspace_id =
        block_on(cluster.create_keyspace(&Default::default(), Duration::from_secs(10)));
    let pd_client = cluster.get_pd_client();

    let mut client = cluster.new_client();
    client.split_keyspace(keyspace_id);
    let gen_key = i_to_keyspace_key(keyspace_id);
    client.put_kv(0..100, &gen_key, i_to_val);
    client.del_kv(0..100, &gen_key);
    let region_id = client.get_region_id(&gen_key(0));
    cluster.flush_memtable(region_id).unwrap();
    assert!(cluster.wait_for_memtable_flushed(region_id, Duration::from_secs(10)));
    let kv = cluster.get_kvengine(node_id);
    must_wait(
        || kv.get_shard_stat(region_id).l0_table_count > 0,
        10,
        || "flushed L0 table not visible".to_string(),
    );
    let tombs_before = kv.get_shard_stat(region_id).tombs;
    assert!(tombs_before > 0);

    // Force the disk to read as critically low so the tick bails at the gate.
    let low_space_fp = "engine_is_critical_low_space";
    fail::cfg(low_space_fp, "return").unwrap();
    defer!(fail::remove(low_space_fp));
    let skip_count = || {
        kvengine::metrics::ENGINE_FILE_GC_COMPACTION_TRIGGER_SKIP_COUNTER
            .with_label_values(&["critical_low_space"])
            .get()
    };
    let skips_before = skip_count();

    // Advance the safe point so the tick would otherwise compact.
    let safe_point = client.get_ts();
    block_on(pd_client.advance_txn_safe_point(keyspace_id, safe_point)).unwrap();
    block_on(pd_client.advance_gc_safe_point(keyspace_id, safe_point)).unwrap();
    block_on(cluster.update_gc_states_immediately()).unwrap();

    // The gate fires (the skip counter climbs) and no data is dropped.
    must_wait(
        || skip_count() > skips_before,
        30,
        || "low-space gate never recorded a skip".to_string(),
    );
    thread::sleep(Duration::from_secs(2));
    assert_eq!(
        kv.get_shard_stat(region_id).tombs,
        tombs_before,
        "file-GC must not compact while disk space is critically low"
    );

    // Once space recovers, the tick resumes and reclaims the tombstones.
    fail::remove(low_space_fp);
    must_wait(
        || kv.get_shard_stat(region_id).tombs == 0,
        60,
        || {
            format!(
                "file-GC did not resume after low space cleared: tombs={}",
                kv.get_shard_stat(region_id).tombs
            )
        },
    );
    cluster.stop();
}