// Copyright 2023 PingCAP, Inc.

use std::{
    sync::atomic::{AtomicUsize, Ordering},
    thread::JoinHandle,
    time::Duration,
};

use kvproto::raft_serverpb::PeerState;
use rfstore::store::load_last_peer_state;
use test_cloud_server::{ServerCluster, client::ClusterClient};
use tikv_util::{info, time::Instant};

use crate::cases::{alloc_node_id_vec, i_to_key, i_to_val};

// Test for the scenario that prepare merge & rollback merge have write conflict
// on pending merge state.
// See https://github.com/tidbcloud/cloud-storage-engine/issues/664
#[test]
fn test_merge_pending_state_conflict() {
    const TEST_DURATION: Duration = Duration::from_secs(3);
    const APPLY_DELAY: &str = "sleep(100)";

    let schedule_merge_error_fp = "on_schedule_merge_error";
    let on_follower_exec_rollback_merge_fp = "on_follower_exec_rollback_merge";

    let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, _| {});
    cluster.wait_region_replicated(&[], 3);
    cluster.get_pd_client().disable_default_operator();

    let mut client = cluster.new_client();
    let prev_keyspace = i_to_key(0);
    client.split(&prev_keyspace);

    let split_key = i_to_key(5);
    client.split(&split_key);
    cluster.wait_pd_region_count(3);

    client.put_kv(0..10, i_to_key, i_to_val);

    // `schedule_merge` returns error to make merges rollback.
    fail::cfg(schedule_merge_error_fp, "return").unwrap();

    let merge_thread = spawn_merge(
        cluster.new_client(),
        i_to_key(0),
        i_to_key(10),
        TEST_DURATION,
    );

    // Delay the apply of rollback merge.
    // The delay is applied to followers only, as this issue would not happen on
    // leader (when there is a pending merge, leader will reject other merge
    // proposals).
    fail::cfg(on_follower_exec_rollback_merge_fp, APPLY_DELAY).unwrap();

    merge_thread.join().unwrap();

    fail::remove(schedule_merge_error_fp);
    fail::remove(on_follower_exec_rollback_merge_fp);

    // `try_merge` may be rollback by target region changed.
    for _ in 0..10 {
        client.try_merge(&i_to_key(0), &i_to_key(10));
        if cluster.get_pd_client_ext().get_regions_number() == 2 {
            break;
        }
        std::thread::sleep(Duration::from_secs(1));
    }

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

fn spawn_merge(
    mut client: ClusterClient,
    mut source_key: Vec<u8>,
    mut target_key: Vec<u8>,
    duration: Duration,
) -> JoinHandle<()> {
    let merge_counter = AtomicUsize::new(0);
    std::thread::spawn(move || {
        let start_time = Instant::now();
        while start_time.saturating_elapsed() < duration {
            let req_sent = client.try_merge(&source_key, &target_key);
            assert!(req_sent);
            merge_counter.fetch_add(1, Ordering::SeqCst);

            // Swap keys to propose more merges, and make issue reproduction more easily.
            std::mem::swap(&mut source_key, &mut target_key);
        }
        info!(
            "merge thread exit, merge counter {}",
            merge_counter.load(Ordering::Relaxed)
        );
    })
}

#[test]
fn test_async_io_apply_before_leader_persist_merge() {
    let nodes = alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(nodes.clone(), |_, _| {});
    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();

    let prev_keyspace = i_to_key(0);
    client.split(&prev_keyspace);
    let split_key = i_to_key(5);
    client.split(&split_key);
    cluster.wait_pd_region_count(3);
    client.put_kv(0..10, i_to_key, i_to_val);

    let source_region = client.get_region_by_key(&i_to_key(0));
    let target_region = client.get_region_by_key(&i_to_key(10));
    assert_ne!(source_region.id, target_region.id);
    let source_peers = source_region
        .peers
        .iter()
        .map(|peer| (peer.get_store_id(), peer.get_id()))
        .collect::<Vec<_>>();

    fail::cfg("on_schedule_merge_error", "return").unwrap();
    for _ in 0..5 {
        client.try_merge(&i_to_key(0), &i_to_key(10));
        std::thread::sleep(Duration::from_millis(100));
    }
    assert_eq!(cluster.get_pd_client_ext().get_regions_number(), 3);

    let source_leader_store_id = source_region.get_leader().get_store_id();
    let restart_node_id = node_store_ids
        .iter()
        .find_map(|&(node_id, store_id)| (store_id != source_leader_store_id).then_some(node_id))
        .unwrap();
    cluster.stop_node(restart_node_id);
    std::thread::sleep(Duration::from_millis(100));
    cluster.start_node(restart_node_id, |_, _| {});
    cluster.wait_region_replicated(&i_to_key(0), nodes.len());

    fail::remove("on_schedule_merge_error");

    assert!(client.try_merge_and_wait(&i_to_key(0), &i_to_key(10), 10));
    cluster.wait_pd_region_count(2);
    client.verify_data_with_ref_store();

    for (store_id, peer_id) in source_peers {
        let node_id = node_store_ids
            .iter()
            .find_map(|&(node_id, id)| (id == store_id).then_some(node_id))
            .unwrap();
        let state = load_last_peer_state(&cluster.get_rfengine(node_id), peer_id).unwrap();
        assert_eq!(state.get_state(), PeerState::Tombstone, "{:?}", state);
    }

    cluster.stop();
}