// Copyright 2026 PingCAP, Inc.

use rfstore::store::{load_raft_truncated_state, metrics::RAFT_ENTRY_FETCHES_VEC};
use test_cloud_server::{ServerCluster, must_wait};
use test_pd_client::PdClientExt;
use tikv::config::TikvConfig;
use tikv_util::config::ReadableSize;

use crate::cases::i_to_key;

#[test]
fn test_persist_delay_block_log_compaction() {
    let nodes = crate::cases::alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(nodes.clone(), |_, conf: &mut TikvConfig| {
        conf.raft_store.raft_log_gc_count_limit = Some(100000);
        conf.raft_store.raft_log_gc_threshold = 200;
        conf.raft_store.raft_log_gc_size_limit = Some(ReadableSize::mb(20));
        conf.rfengine.target_file_size = ReadableSize::kb(1);
        conf.rfengine.rlog_soft_memory_limit = ReadableSize::kb(4);
        // The generic test fixture randomizes these values. Use the minimum
        // valid epoch window so the writes below deterministically offload.
        conf.rfengine.epoch_rotate_len = 4;
        conf.rfengine.delay_compaction_epoches = 0;
    });
    cluster.wait_region_replicated(&[], nodes.len());
    cluster.get_pd_client().disable_default_operator();

    let node_store_ids: Vec<(u16, u64)> = nodes
        .iter()
        .map(|&node_id| (node_id, cluster.get_store_id(node_id)))
        .collect();
    let mut client = cluster.new_client();

    client.put_kv(1..2, i_to_key, |i| format!("v{i}").into_bytes());
    let region = client.get_region_by_key(&[]);
    let leader_store_id = region.get_leader().get_store_id();
    let leader_node_id = node_store_ids
        .iter()
        .find_map(|&(node_id, store_id)| (store_id == leader_store_id).then_some(node_id))
        .unwrap();
    let leader_peer_id = client.get_peer_id(&[], leader_store_id);
    let before_truncated_index =
        load_raft_truncated_state(&cluster.get_rfengine(leader_node_id), leader_peer_id)
            .map(|state| state.get_index())
            .unwrap_or_default();
    let stopped_node_id = node_store_ids
        .iter()
        .find_map(|&(node_id, store_id)| (store_id != leader_store_id).then_some(node_id))
        .unwrap();

    cluster.stop_node(stopped_node_id);

    let large_value = vec![b'x'; 1024];
    client.put_kv(2..1200, i_to_key, |_| large_value.clone());

    must_wait(
        || {
            cluster
                .get_rfengine(leader_node_id)
                .get_peer_stats(leader_peer_id)
                .num_logs_offloaded
                > 0
        },
        30,
        || {
            format!(
                "wait for leader offload timeout, leader_node_id={}, leader_peer_id={}",
                leader_node_id, leader_peer_id
            )
        },
    );
    must_wait(
        || {
            load_raft_truncated_state(&cluster.get_rfengine(leader_node_id), leader_peer_id)
                .is_some_and(|state| state.get_index() > before_truncated_index)
        },
        30,
        || {
            format!(
                "wait for leader truncated index increase timeout, leader_node_id={}, leader_peer_id={}",
                leader_node_id, leader_peer_id
            )
        },
    );

    let async_fetch_count = RAFT_ENTRY_FETCHES_VEC
        .with_label_values(&["async_fetch"])
        .get();
    cluster.start_node(stopped_node_id, |_, _| {});
    must_wait(
        || {
            RAFT_ENTRY_FETCHES_VEC
                .with_label_values(&["async_fetch"])
                .get()
                > async_fetch_count
        },
        10,
        || "wait for async fetch timeout".to_string(),
    );

    client.verify_data_with_ref_store();
    cluster.stop();
}

#[test]
fn test_rfstore_async_io_apply_without_leader_persist() {
    let nodes = crate::cases::alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(nodes.clone(), |_, conf: &mut TikvConfig| {
        conf.raft_store.raft_log_gc_count_limit = Some(100000);
        conf.raft_store.raft_log_gc_threshold = 200;
        conf.raft_store.raft_log_gc_size_limit = Some(ReadableSize::mb(20));
        conf.rfengine.target_file_size = ReadableSize::kb(1);
        conf.rfengine.rlog_soft_memory_limit = ReadableSize::kb(4);
        // Avoid the fixture's randomized epoch window so async-fetch
        // readiness depends only on the writes made by this test.
        conf.rfengine.epoch_rotate_len = 4;
        conf.rfengine.delay_compaction_epoches = 0;
    });
    cluster.wait_region_replicated(&[], nodes.len());
    cluster.get_pd_client().disable_default_operator();

    let node_store_ids: Vec<(u16, u64)> = nodes
        .iter()
        .map(|&node_id| (node_id, cluster.get_store_id(node_id)))
        .collect();
    let mut client = cluster.new_client();

    client.put_kv(1..2, i_to_key, |i| format!("v{i}").into_bytes());
    let mut region = client.get_region_by_key(&[]);
    let leader_store_id = region.get_leader().get_store_id();
    let (stopped_node_id, stopped_store_id) = node_store_ids
        .iter()
        .find_map(|&(node_id, store_id)| {
            (store_id != leader_store_id).then_some((node_id, store_id))
        })
        .unwrap();
    let target_store_id = region
        .peers
        .iter()
        .find_map(|peer| {
            let store_id = peer.get_store_id();
            (store_id != leader_store_id && store_id != stopped_store_id).then_some(store_id)
        })
        .unwrap();
    let target_peer = region
        .peers
        .iter()
        .find(|peer| peer.get_store_id() == target_store_id)
        .cloned()
        .unwrap();
    let target_peer_id = target_peer.get_id();
    let target_node_id = node_store_ids
        .iter()
        .find_map(|&(node_id, store_id)| (store_id == target_store_id).then_some(node_id))
        .unwrap();
    // Measure the lagging peer directly before stopping it. The restarted peer
    // will resume from `last_idx + 1`, so the new leader must have already
    // offloaded entries below that boundary to force async fetch.
    let stopped_peer_id = client.get_peer_id(&[], stopped_store_id);
    let stopped_last_index = cluster
        .get_rfengine(stopped_node_id)
        .get_peer_stats(stopped_peer_id)
        .last_idx;

    cluster.stop_node(stopped_node_id);
    let large_value = vec![b'y'; 1024];
    client.put_kv(2..1000, i_to_key, |_| large_value.clone());

    region = client.get_region_by_key(&[]);
    cluster
        .get_pd_client()
        .must_transfer_leader(region.id, target_peer.clone());

    // Drive the transferred leader forward until it really offloads entries
    // below the restarted peer's resume point. After leadership changes, the
    // new leader may need additional writes before offloading advances.
    let leader_ready_for_async_fetch = || {
        let stats = cluster
            .get_rfengine(target_node_id)
            .get_peer_stats(target_peer_id);
        stats.num_logs_offloaded > 0 && stats.first_idx > stopped_last_index + 1
    };
    let mut next_key = 2000;
    for _ in 0..10 {
        if leader_ready_for_async_fetch() {
            break;
        }
        client.put_kv(next_key..next_key + 200, i_to_key, |_| large_value.clone());
        next_key += 200;
    }
    must_wait(
        || leader_ready_for_async_fetch(),
        30,
        || {
            let stats = cluster
                .get_rfengine(target_node_id)
                .get_peer_stats(target_peer_id);
            format!(
                "wait for leader async-fetch readiness timeout, leader_node_id={}, leader_peer_id={}, stopped_node_id={}, stopped_peer_id={}, stopped_last_index={}, leader_first_idx={}, leader_last_idx={}, leader_truncated_idx={}, leader_num_logs_offloaded={}",
                target_node_id,
                target_peer_id,
                stopped_node_id,
                stopped_peer_id,
                stopped_last_index,
                stats.first_idx,
                stats.last_idx,
                stats.truncated_idx,
                stats.num_logs_offloaded
            )
        },
    );

    let async_fetch_count = RAFT_ENTRY_FETCHES_VEC
        .with_label_values(&["async_fetch"])
        .get();
    cluster.start_node(stopped_node_id, |_, _| {});
    must_wait(
        || {
            RAFT_ENTRY_FETCHES_VEC
                .with_label_values(&["async_fetch"])
                .get()
                > async_fetch_count
        },
        10,
        || "wait for async fetch timeout".to_string(),
    );

    client.put_kv(1000..1100, i_to_key, |i| format!("v{i}").into_bytes());
    client.verify_data_with_ref_store();
    cluster.stop();
}