// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
    thread,
    time::{Duration, Instant},
};

use raftstore::coprocessor::RegionInfoProvider;
use rfstore::{
    RaftRouter,
    store::{ColdRegionPolicy, PeerMsg, metrics::COLD_CAMPAIGN_EVENT_COUNTER_VEC},
};
use test_cloud_server::{
    ServerCluster, ServerClusterBuilder, alloc_node_id_vec,
    keyspace::{CreateKeyspaceOptions, make_row_key},
};
use tikv::config::TikvConfig;
use tikv_util::config::ReadableDuration;

const NODE_COUNT: usize = 3;
const BASE_TICK: Duration = Duration::from_millis(50);
const ELECTION_TIMEOUT: Duration = Duration::from_millis(500);
const WAIT_TIMEOUT: Duration = Duration::from_secs(20);
const POLL_INTERVAL: Duration = Duration::from_millis(20);

fn configure_cold_region(_: u16, cfg: &mut TikvConfig) {
    cfg.rocksdb.max_open_files = 1_024;
    cfg.raftdb.max_open_files = 1_024;
    cfg.raft_store.cold_region.policy = ColdRegionPolicy::Enabled;
    cfg.raft_store.raft_base_tick_interval = ReadableDuration(BASE_TICK);
    cfg.raft_store.raft_heartbeat_ticks = 2;
    cfg.raft_store.raft_election_timeout_ticks =
        (ELECTION_TIMEOUT.as_millis() / BASE_TICK.as_millis()) as usize;
    cfg.raft_store.raft_store_max_leader_lease = ReadableDuration::millis(100);
    cfg.raft_store.split_region_check_tick_interval = ReadableDuration::millis(100);
    cfg.raft_store.pd_heartbeat_tick_interval = ReadableDuration::millis(100);
    cfg.raft_store.pd_store_heartbeat_tick_interval = ReadableDuration::millis(100);
    cfg.raft_store.cold_region.max_ts_threshold = ReadableDuration::millis(100);
    cfg.raft_store.cold_region.store_down_grace_period = ReadableDuration::secs(1);
    cfg.raft_store.cold_region.region_heartbeat_interval = ReadableDuration::millis(100);
}

fn wait_until(mut condition: impl FnMut() -> bool, description: &str) {
    let started_at = Instant::now();
    while started_at.elapsed() < WAIT_TIMEOUT {
        if condition() {
            return;
        }
        thread::sleep(POLL_INTERVAL);
    }
    panic!(
        "{} did not become true within {:?}",
        description, WAIT_TIMEOUT
    );
}

fn peer_is_cold(cluster: &ServerCluster, node_id: u16, region_id: u64) -> bool {
    cluster
        .get_region_info_accessor(node_id)
        .find_region_by_id(region_id)
        .is_some_and(|region| region.is_cold_peer)
}

fn wait_all_peers_cold(
    cluster: &ServerCluster,
    nodes: &[u16],
    routers: &Mutex<HashMap<u16, RaftRouter>>,
    region_id: u64,
) {
    wait_until(
        || {
            for &node_id in nodes {
                if !peer_is_cold(cluster, node_id, region_id) {
                    routers.lock().unwrap()[&node_id]
                        .send_peer_msg_for_test(region_id, PeerMsg::ColdRegionTick(1));
                }
            }
            nodes
                .iter()
                .all(|&node_id| peer_is_cold(cluster, node_id, region_id))
        },
        "all peers entering Cold",
    );
}

#[test]
fn test_cold_region_recovers_after_leader_store_down() {
    let nodes = alloc_node_id_vec(NODE_COUNT);
    let routers = Arc::new(Mutex::new(HashMap::new()));
    let hook_routers = routers.clone();
    let mut cluster = ServerClusterBuilder::new(nodes.clone(), configure_cold_region)
        .memory_capacity_ratio(1.0 / NODE_COUNT as f64)
        .before_run_server(move |node_id, server| {
            hook_routers
                .lock()
                .unwrap()
                .insert(node_id, server.get_raft_router());
        })
        .build();
    cluster.wait_region_replicated(&[], NODE_COUNT);

    // Create an empty whole-keyspace Region and make every replica Cold.
    let mut client = cluster.new_client();
    let keyspace_id = cluster.keyspace_manager().new_keyspace_id(1);
    cluster.keyspace_manager().create_keyspaces(
        &[keyspace_id],
        vec!["cold-down-store".to_owned()],
        &CreateKeyspaceOptions::default(),
        None,
    );
    client.split_keyspaces([keyspace_id]);
    client.clear_region_cache();
    let key = make_row_key(keyspace_id, 1, b"key");
    let region_id = client.get_region_id(&key);
    cluster.wait_region_replicated(&key, NODE_COUNT);
    wait_all_peers_cold(&cluster, &nodes, &routers, region_id);

    // Stop the cached leader without sending any foreground request.
    let old_leader = cluster.leader_of_region(region_id).unwrap();
    let old_leader_node = cluster
        .get_server_node_id(old_leader.get_store_id())
        .unwrap();
    let live_nodes = nodes
        .iter()
        .copied()
        .filter(|node_id| *node_id != old_leader_node)
        .collect::<Vec<_>>();
    let campaign_count = COLD_CAMPAIGN_EVENT_COUNTER_VEC.store_down.get();
    cluster.stop_node_force(old_leader_node, true);

    // PD down-store detection plus the cold tick must recover a leader.
    wait_until(
        || {
            for &node_id in &live_nodes {
                let router = &routers.lock().unwrap()[&node_id];
                router.send_peer_msg_for_test(region_id, PeerMsg::ColdRegionTick(1));
                router.send_peer_msg_for_test(region_id, PeerMsg::ColdLeaderProbeTimeout);
            }
            COLD_CAMPAIGN_EVENT_COUNTER_VEC.store_down.get() > campaign_count
                && cluster
                    .leader_of_region(region_id)
                    .is_some_and(|leader| leader.get_store_id() != old_leader.get_store_id())
        },
        "cold Region recovering a leader after its leader store is down",
    );
}