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

use std::{
    collections::{HashMap, HashSet},
    env,
    sync::{
        Arc, Mutex, RwLock,
        atomic::{AtomicBool, AtomicU16, Ordering},
        mpsc,
    },
    thread,
    time::{Duration, Instant},
};

use futures::executor::block_on;
use hyper::{Body, Method};
use pd_client::PdClient;
use rand::{Rng, SeedableRng, prelude::SliceRandom, rngs::StdRng};
use rfstore::{
    RaftRouter, RaftStoreRouter,
    store::{Callback, ColdPeerStateSnapshot, ColdRegionPolicy, StoreMsg},
};
use test_cloud_server::{
    ServerCluster, ServerClusterBuilder, alloc_node_id_vec,
    client::{ClusterClient, MutateOptions},
    keyspace::{CreateKeyspaceOptions, make_row_key},
    move_region_to_voter_stores,
};
use test_pd_client::PdClientExt;
use test_raftstore::{new_get_cmd, new_request, new_snap_cmd};
use tikv::config::TikvConfig;
use tikv_util::{config::ReadableDuration, info};

use crate::init_logger;

const NODE_COUNT: usize = 4;
const KEYSPACE_COUNT: usize = 1_000;
const PERSISTED_COLD_COUNT: usize = 18;
const COLD_REGION_COUNT: usize = KEYSPACE_COUNT - 2;
const TRACKED_REGION_COUNT: usize = KEYSPACE_COUNT + 1;
const TRACKED_PEER_COUNT: usize = TRACKED_REGION_COUNT * 3;
const COLD_PEER_COUNT: usize = COLD_REGION_COUNT * 3;
const STEADY_REGULAR_COUNT: usize = 9;
const DEPLOYMENT_REGION_COUNT: usize = TRACKED_REGION_COUNT + 2;
const DEPLOYMENT_PEER_COUNT: usize = DEPLOYMENT_REGION_COUNT * 3;
// Of the two untracked boundary regions, the bootstrap region stays active.
const DEPLOYMENT_COLD_REGION_COUNT: usize = COLD_REGION_COUNT + 1;
const DEPLOYMENT_COLD_PEER_COUNT: usize = DEPLOYMENT_COLD_REGION_COUNT * 3;
const DEPLOYMENT_REGULAR_COUNT: usize = DEPLOYMENT_PEER_COUNT - DEPLOYMENT_COLD_PEER_COUNT;
const FOLLOWER_ONLY_TEMP_REGULAR_COUNT: usize = DEPLOYMENT_COLD_REGION_COUNT * 2;

const BASE_TICK: Duration = Duration::from_millis(50);
const ELECTION_TIMEOUT: Duration = Duration::from_millis(500);
const MAX_LEADER_LEASE: Duration = Duration::from_millis(100);
const SPLIT_CHECK_INTERVAL: Duration = Duration::from_millis(100);
const MAX_TS_THRESHOLD: Duration = Duration::from_millis(300);
const STORE_DOWN_GRACE_PERIOD: Duration = Duration::from_secs(5);
// Keep the production 5m:30m:2h ratio at one tenth of the duration.
const COLD_REGION_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
const COLD_LEADER_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(3 * 60);
const MAX_LEADER_MISSING_DURATION: Duration = Duration::from_secs(12 * 60);

const STATE_QUERY_INTERVAL: Duration = Duration::from_millis(100);
const STATE_QUERY_TIMEOUT: Duration = Duration::from_secs(2);
const STATE_WAIT_TIMEOUT: Duration = Duration::from_secs(180);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const FINAL_CHECK_RESERVE: Duration = Duration::from_secs(45);
const SHORT_RESTART_BUDGET: Duration = Duration::from_secs(40);
const LONG_RESTART_BUDGET: Duration = Duration::from_secs(100);
const SHORT_OUTAGE: Duration = Duration::from_millis(100);
const LONG_OUTAGE: Duration = Duration::from_secs(67);
const STALE_LEADER_RECOVERY_TIMEOUT: Duration = Duration::from_secs(2);
const DEFAULT_TEST_SECS: u64 = 600;

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
enum Action {
    Read,
    Write,
    TransferLeader,
    MovePeer,
    StaleLeaderRead,
    ShortRestart,
    LongRestart,
}

#[derive(Clone)]
struct ColdRegion {
    region_id: u64,
    key: Vec<u8>,
}

#[derive(Clone, Copy, Debug)]
struct ColdMetrics {
    cold_regions: usize,
    campaign_events: u64,
}

struct StatusClient {
    runtime: tokio::runtime::Runtime,
    client: hyper::Client<hyper::client::HttpConnector>,
}

impl StatusClient {
    fn new() -> Self {
        Self {
            runtime: tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap(),
            client: hyper::Client::new(),
        }
    }

    fn update_cold_policy(&self, cluster: &ServerCluster, nodes: &[u16], policy: ColdRegionPolicy) {
        let body = serde_json::to_vec(&HashMap::from([(
            "raftstore.cold-region.policy".to_owned(),
            policy.to_string(),
        )]))
        .unwrap();
        for &node_id in nodes {
            let (status, response) = self.request(
                Method::POST,
                format!("http://{}/config", cluster.status_addr(node_id)),
                body.clone(),
            );
            assert!(
                status.is_success(),
                "update cold-region policy on node {} failed: policy={}, status={}, response={}",
                node_id,
                policy,
                status,
                response,
            );
        }
    }

    fn cold_metrics(&self, cluster: &ServerCluster, node_id: u16) -> ColdMetrics {
        let (status, body) = self.request(
            Method::GET,
            format!("http://{}/metrics", cluster.status_addr(node_id)),
            Vec::new(),
        );
        assert!(
            status.is_success(),
            "query metrics on node {} failed: status={}, response={}",
            node_id,
            status,
            body,
        );
        ColdMetrics {
            cold_regions: prometheus_metric_sum(&body, "rfstore_cold_regions_count") as usize,
            campaign_events: prometheus_metric_sum(&body, "rfstore_cold_campaign_events_total"),
        }
    }

    fn request(&self, method: Method, uri: String, body: Vec<u8>) -> (hyper::StatusCode, String) {
        self.runtime.block_on(async {
            let request = hyper::Request::builder()
                .method(method)
                .uri(uri)
                .body(Body::from(body))
                .unwrap();
            let response = self.client.request(request).await.unwrap();
            let status = response.status();
            let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
            (status, String::from_utf8(body.to_vec()).unwrap())
        })
    }
}

struct StateMonitor {
    running: Arc<AtomicBool>,
    failure: Arc<Mutex<Option<String>>>,
    handle: Option<thread::JoinHandle<()>>,
}

impl StateMonitor {
    fn start(
        routers: Arc<Mutex<HashMap<u16, RaftRouter>>>,
        region_ids: Arc<HashSet<u64>>,
        offline_node: Arc<AtomicU16>,
        current_action: Arc<RwLock<String>>,
    ) -> Self {
        let running = Arc::new(AtomicBool::new(true));
        let failure = Arc::new(Mutex::new(None));
        let thread_running = running.clone();
        let thread_failure = failure.clone();
        let handle = thread::spawn(move || {
            while thread_running.load(Ordering::Acquire) {
                let started_at = Instant::now();
                let offline_node = offline_node.load(Ordering::Acquire);
                match query_cold_states(&routers, region_ids.clone(), offline_node) {
                    Ok(snapshot) => {
                        let non_cold = snapshot.regular_count + snapshot.temp_regular_count;
                        let invalid_regular = snapshot.regular_count * 100 > snapshot.peer_count;
                        let invalid_non_cold = non_cold * 100 > snapshot.peer_count * 5;
                        if invalid_regular || invalid_non_cold {
                            let action = current_action.read().unwrap().clone();
                            *thread_failure.lock().unwrap() = Some(format!(
                                "cold-state limit exceeded: action={action}, offline_node={offline_node}, \
                                 peers={}, cold={}, regular={}, temp_regular={}, non_cold={:?}",
                                snapshot.peer_count,
                                snapshot.cold_count,
                                snapshot.regular_count,
                                snapshot.temp_regular_count,
                                snapshot.non_cold_peers,
                            ));
                            return;
                        }
                    }
                    Err(err) => {
                        let action = current_action.read().unwrap().clone();
                        *thread_failure.lock().unwrap() =
                            Some(format!("cold-state query failed: action={action}, {err}"));
                        return;
                    }
                }
                thread::sleep(STATE_QUERY_INTERVAL.saturating_sub(started_at.elapsed()));
            }
        });
        Self {
            running,
            failure,
            handle: Some(handle),
        }
    }

    fn check(&self) {
        if let Some(err) = self.failure.lock().unwrap().clone() {
            panic!("{}", err);
        }
    }

    fn stop(&mut self) {
        self.running.store(false, Ordering::Release);
        if let Some(handle) = self.handle.take() {
            handle.join().unwrap();
        }
        self.check();
    }
}

impl Drop for StateMonitor {
    fn drop(&mut self) {
        self.running.store(false, Ordering::Release);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

struct ColdRegionTest {
    cluster: ServerCluster,
    client: ClusterClient,
    routers: Arc<Mutex<HashMap<u16, RaftRouter>>>,
    offline_node: Arc<AtomicU16>,
    store_nodes: HashMap<u64, u16>,
    tracked_region_ids: Arc<HashSet<u64>>,
    cold_region_ids: Arc<HashSet<u64>>,
    cold_regions: Vec<ColdRegion>,
    expected_values: HashMap<Vec<u8>, Vec<u8>>,
    hot_key: Vec<u8>,
    value_seq: u64,
}

impl ColdRegionTest {
    fn query(&self, region_ids: Arc<HashSet<u64>>) -> ColdPeerStateSnapshot {
        query_cold_states(
            &self.routers,
            region_ids,
            self.offline_node.load(Ordering::Acquire),
        )
        .unwrap()
    }

    fn wait_region_cold(&self, region_id: u64) {
        let region_ids = Arc::new(HashSet::from([region_id]));
        self.wait_until(region_ids, |snapshot| {
            snapshot.peer_count == 3 && snapshot.cold_count == 3
        });
    }

    fn wait_all_cold(&self) {
        self.wait_until(self.cold_region_ids.clone(), |snapshot| {
            snapshot.peer_count == COLD_PEER_COUNT && snapshot.cold_count == COLD_PEER_COUNT
        });
    }

    fn wait_exact_steady_state(&self) {
        self.wait_until(self.tracked_region_ids.clone(), |snapshot| {
            snapshot.peer_count == TRACKED_PEER_COUNT
                && snapshot.cold_count == COLD_PEER_COUNT
                && snapshot.regular_count == STEADY_REGULAR_COUNT
                && snapshot.temp_regular_count == 0
        });
    }

    fn wait_all_regular(&self, region_ids: Arc<HashSet<u64>>) {
        self.wait_until(region_ids, |snapshot| {
            snapshot.peer_count == DEPLOYMENT_PEER_COUNT
                && snapshot.regular_count == DEPLOYMENT_PEER_COUNT
                && snapshot.temp_regular_count == 0
                && snapshot.cold_count == 0
        });
    }

    fn wait_follower_only(&self, region_ids: Arc<HashSet<u64>>) {
        self.wait_until(region_ids, |snapshot| {
            snapshot.peer_count == DEPLOYMENT_PEER_COUNT
                && snapshot.cold_count == 0
                && snapshot.temp_regular_count == FOLLOWER_ONLY_TEMP_REGULAR_COUNT
                && snapshot.regular_count + snapshot.temp_regular_count == DEPLOYMENT_PEER_COUNT
        });
    }

    fn wait_deployment_enabled(&self, region_ids: Arc<HashSet<u64>>) {
        self.wait_until(region_ids, |snapshot| {
            snapshot.peer_count == DEPLOYMENT_PEER_COUNT
                && snapshot.regular_count == DEPLOYMENT_REGULAR_COUNT
                && snapshot.temp_regular_count == 0
                && snapshot.cold_count == DEPLOYMENT_COLD_PEER_COUNT
        });
    }

    fn wait_cold_metrics(
        &self,
        status_client: &StatusClient,
        node_id: u16,
        expected_cold_regions: usize,
        expected_campaign_events: Option<u64>,
    ) -> ColdMetrics {
        let started_at = Instant::now();
        let mut last = None;
        while started_at.elapsed() < STATE_WAIT_TIMEOUT {
            let metrics = status_client.cold_metrics(&self.cluster, node_id);
            if let Some(expected) = expected_campaign_events {
                assert_eq!(
                    metrics.campaign_events, expected,
                    "cold campaign occurred while changing cold-region policy"
                );
            }
            if metrics.cold_regions == expected_cold_regions {
                return metrics;
            }
            last = Some(metrics);
            thread::sleep(STATE_QUERY_INTERVAL);
        }
        panic!(
            "cold-region metrics did not converge within {:?}: expected_cold_regions={}, last={:?}",
            STATE_WAIT_TIMEOUT, expected_cold_regions, last
        );
    }

    fn set_restart_policy(&mut self, nodes: &[u16], policy: ColdRegionPolicy) {
        // ServerCluster restarts from its stored config rather than the config
        // persisted by the online endpoint.
        for &node_id in nodes {
            let mut config = self.cluster.get_node_config(node_id).clone();
            config.raft_store.cold_region.policy = policy;
            self.cluster.update_node_config(node_id, config);
        }
    }

    fn wait_until(
        &self,
        region_ids: Arc<HashSet<u64>>,
        mut check: impl FnMut(&ColdPeerStateSnapshot) -> bool,
    ) -> ColdPeerStateSnapshot {
        let started_at = Instant::now();
        let mut last = None;
        while started_at.elapsed() < STATE_WAIT_TIMEOUT {
            let snapshot = self.query(region_ids.clone());
            if check(&snapshot) {
                return snapshot;
            }
            last = Some(snapshot);
            thread::sleep(STATE_QUERY_INTERVAL);
        }
        panic!(
            "cold state did not converge within {:?}: {:?}",
            STATE_WAIT_TIMEOUT, last
        );
    }

    fn read(&mut self, region_idx: usize, direct_snap: bool) {
        let region = self.cold_regions[region_idx].clone();
        if direct_snap {
            self.direct_snap(region.region_id, None);
        }
        self.assert_read(&region.key);
        self.wait_region_cold(region.region_id);
    }

    fn write(&mut self, region_idx: usize) {
        let region = self.cold_regions[region_idx].clone();
        let value = self.next_value();
        put(&mut self.client, &region.key, &value);
        self.expected_values.insert(region.key.clone(), value);
        self.assert_read(&region.key);
        self.flush_and_wait(region.region_id);
    }

    fn transfer_leader(&mut self, region_idx: usize, rng: &mut StdRng) {
        let region_id = self.cold_regions[region_idx].region_id;
        let region = block_on(self.cluster.get_pd_client().get_region_by_id(region_id))
            .unwrap()
            .unwrap();
        let leader = self.cluster.leader_of_region(region_id).unwrap();
        let candidates = region
            .peers
            .iter()
            .filter(|peer| peer.get_id() != leader.get_id())
            .cloned()
            .collect::<Vec<_>>();
        let target = candidates[rng.gen_range(0..candidates.len())].clone();
        info!("transfer cold-region leader";
            "region_id" => region_id,
            "target_peer" => target.get_id(),
            "target_store" => target.get_store_id());
        self.cluster.transfer_leader(region_id, target);
        self.read(region_idx, false);
    }

    fn move_peer(&mut self, region_idx: usize, rng: &mut StdRng) {
        let region_id = self.cold_regions[region_idx].region_id;
        let pd_client = self.cluster.get_pd_client();
        let region = block_on(pd_client.get_region_by_id(region_id))
            .unwrap()
            .unwrap();
        let all_stores = self.store_nodes.keys().copied().collect::<Vec<_>>();
        let target_store = *all_stores
            .iter()
            .find(|store_id| {
                region
                    .peers
                    .iter()
                    .all(|peer| peer.get_store_id() != **store_id)
            })
            .unwrap();
        let remove_idx = rng.gen_range(0..region.peers.len());
        let mut target_voters = region
            .peers
            .iter()
            .enumerate()
            .filter_map(|(idx, peer)| (idx != remove_idx).then_some(peer.get_store_id()))
            .collect::<Vec<_>>();
        target_voters.push(target_store);
        let leader_store = target_voters[rng.gen_range(0..target_voters.len())];
        info!("move cold-region peer";
            "region_id" => region_id,
            "removed_store" => region.peers[remove_idx].get_store_id(),
            "target_store" => target_store,
            "leader_store" => leader_store);
        let region = move_region_to_voter_stores(
            &self.cluster,
            &pd_client,
            region_id,
            &target_voters,
            leader_store,
        );
        assert_eq!(region.peers.len(), 3);
        assert!(
            region
                .peers
                .iter()
                .all(|peer| peer.get_role() == kvproto::metapb::PeerRole::Voter)
        );
        self.read(region_idx, false);
    }

    fn stale_leader_read(&mut self, region_idx: usize, rng: &mut StdRng) {
        let region = self.cold_regions[region_idx].clone();
        let pd_region = block_on(
            self.cluster
                .get_pd_client()
                .get_region_by_id(region.region_id),
        )
        .unwrap()
        .unwrap();
        let leader = self.cluster.leader_of_region(region.region_id).unwrap();
        let followers = pd_region
            .peers
            .iter()
            .filter(|peer| peer.get_id() != leader.get_id())
            .cloned()
            .collect::<Vec<_>>();
        let mut follower_idx = rng.gen_range(0..followers.len());
        let follower = followers[follower_idx].clone();
        let stopped_node = self.store_nodes[&leader.get_store_id()];
        let get = rng.gen_bool(0.5);
        let force = rng.gen_bool(0.5);
        info!("read cold follower with stale leader hint";
            "region_id" => region.region_id,
            "leader_peer" => leader.get_id(),
            "leader_store" => leader.get_store_id(),
            "follower_peer" => follower.get_id(),
            "force" => force,
            "get" => get);

        self.offline_node.store(stopped_node, Ordering::Release);
        self.cluster.stop_node_force(stopped_node, force);
        let started_at = Instant::now();
        let mut response = self.direct_read(
            &pd_region,
            follower.clone(),
            get.then_some(region.key.as_slice()),
        );
        let error = response.get_header().get_error();
        assert!(error.has_not_leader());
        if error.get_not_leader().has_leader() {
            assert_eq!(
                error.get_not_leader().get_leader().get_id(),
                leader.get_id()
            );
        }

        while response.get_header().has_error() {
            assert!(
                started_at.elapsed() < STALE_LEADER_RECOVERY_TIMEOUT,
                "stale leader-hint {:?} did not recover within {:?}: {:?}",
                if get { "Get" } else { "Snap" },
                STALE_LEADER_RECOVERY_TIMEOUT,
                response.get_header().get_error(),
            );
            thread::sleep(BASE_TICK);
            follower_idx = (follower_idx + 1) % followers.len();
            response = self.direct_read(
                &pd_region,
                followers[follower_idx].clone(),
                get.then_some(region.key.as_slice()),
            );
        }
        assert!(started_at.elapsed() < STALE_LEADER_RECOVERY_TIMEOUT);
        info!("stale leader-hint read recovered";
            "region_id" => region.region_id,
            "get" => get,
            "elapsed" => ?started_at.elapsed());

        self.cluster.start_node(stopped_node, |_, _| {});
        self.offline_node.store(0, Ordering::Release);
        self.assert_read(&region.key);
        self.wait_all_cold();
    }

    fn restart(&mut self, region_idx: usize, long: bool, rng: &mut StdRng) {
        let region = self.cold_regions[region_idx].clone();
        let pd_region = block_on(
            self.cluster
                .get_pd_client()
                .get_region_by_id(region.region_id),
        )
        .unwrap()
        .unwrap();
        let leader = self.cluster.leader_of_region(region.region_id).unwrap();
        let stop_leader = rng.gen_bool(0.5);
        let stopped_peer = if stop_leader {
            leader
        } else {
            pd_region
                .peers
                .iter()
                .find(|peer| peer.get_id() != leader.get_id())
                .unwrap()
                .clone()
        };
        let stopped_node = self.store_nodes[&stopped_peer.get_store_id()];
        let force = rng.gen_bool(0.5);
        let write = rng.gen_bool(0.5);
        let value = write.then(|| self.next_value());
        let expected = self.expected_values.get(&region.key).cloned();
        let mut request_client = self.client.clone();
        let request_key = region.key.clone();
        let request_value = value.clone();
        let outage = if long { LONG_OUTAGE } else { SHORT_OUTAGE };
        let request_timeout = if long {
            outage + REQUEST_TIMEOUT
        } else {
            REQUEST_TIMEOUT
        };
        info!("restart store for cold region";
            "region_id" => region.region_id,
            "node_id" => stopped_node,
            "store_id" => stopped_peer.get_store_id(),
            "stop_leader" => stop_leader,
            "force" => force,
            "long" => long,
            "write" => write);

        self.offline_node.store(stopped_node, Ordering::Release);
        self.cluster.stop_node_force(stopped_node, force);
        let request_started_at = Instant::now();
        let (request_tx, request_rx) = mpsc::sync_channel(1);
        let request = thread::spawn(move || {
            if let Some(value) = request_value {
                put_with_retry(&mut request_client, &request_key, &value, request_timeout);
            } else {
                let actual = get_with_retry(&mut request_client, &request_key, request_timeout);
                assert_eq!(actual, expected);
            }
            let _ = request_tx.send(request_started_at.elapsed());
        });

        thread::sleep(outage);
        self.cluster.start_node(stopped_node, |_, _| {});
        self.offline_node.store(0, Ordering::Release);

        if stop_leader {
            self.direct_snap(region.region_id, Some(stopped_peer));
        }
        let request_elapsed = request_rx.recv_timeout(REQUEST_TIMEOUT).unwrap();
        assert!(request_elapsed <= request_timeout);
        info!("cold-region restart request completed";
            "region_id" => region.region_id,
            "elapsed" => ?request_elapsed);
        request.join().unwrap();
        if let Some(value) = value {
            self.expected_values.insert(region.key.clone(), value);
            self.cluster.flush_memtable(region.region_id).unwrap();
            assert!(
                self.cluster
                    .wait_for_memtable_flushed(region.region_id, Duration::from_secs(30))
            );
        }
        self.assert_read(&region.key);
        self.wait_all_cold();
    }

    fn direct_snap(&self, region_id: u64, target: Option<kvproto::metapb::Peer>) {
        let pd_client = self.cluster.get_pd_client();
        let region = block_on(pd_client.get_region_by_id(region_id))
            .unwrap()
            .unwrap();
        let peer = target.unwrap_or_else(|| self.cluster.leader_of_region(region_id).unwrap());
        let response = self.direct_read(&region, peer, None);
        if response.get_header().has_error() {
            assert!(response.get_header().get_error().has_not_leader());
        }
    }

    fn direct_read(
        &self,
        region: &kvproto::metapb::Region,
        peer: kvproto::metapb::Peer,
        key: Option<&[u8]>,
    ) -> kvproto::raft_cmdpb::RaftCmdResponse {
        let node_id = self.store_nodes[&peer.get_store_id()];
        let router = self.routers.lock().unwrap()[&node_id].clone();
        let mut request = new_request(
            region.get_id(),
            region.get_region_epoch().clone(),
            vec![key.map_or_else(new_snap_cmd, |key| new_get_cmd(key))],
            false,
        );
        request.mut_header().set_peer(peer);
        let (tx, rx) = mpsc::sync_channel(1);
        router.send_command(
            request,
            Default::default(),
            Callback::read(Box::new(move |response| {
                let _ = tx.send(response.response);
            })),
        );
        rx.recv_timeout(Duration::from_secs(3)).unwrap()
    }

    fn flush_and_wait(&self, region_id: u64) {
        self.cluster.flush_memtable(region_id).unwrap();
        assert!(
            self.cluster
                .wait_for_memtable_flushed(region_id, Duration::from_secs(30))
        );
        self.wait_region_cold(region_id);
    }

    fn refresh_hot_region(&mut self) {
        let value = self.next_value();
        put(&mut self.client, &self.hot_key, &value);
        self.expected_values.insert(self.hot_key.clone(), value);
    }

    fn assert_read(&mut self, key: &[u8]) {
        let actual = self.client.get_key_version(key, u64::MAX).unwrap();
        assert_eq!(actual.as_ref(), self.expected_values.get(key));
    }

    fn verify_data(&mut self) {
        for (key, value) in &self.expected_values {
            assert_eq!(
                self.client.get_key_version(key, u64::MAX).unwrap(),
                Some(value.clone()),
                "unexpected value for key {key:?}"
            );
        }
    }

    fn next_value(&mut self) -> Vec<u8> {
        self.value_seq += 1;
        format!("cold-region-random-{}", self.value_seq).into_bytes()
    }
}

fn prometheus_metric_sum(metrics: &str, name: &str) -> u64 {
    let value = metrics
        .lines()
        .filter_map(|line| {
            let mut fields = line.split_whitespace();
            let series = fields.next()?;
            if series != name
                && !series
                    .strip_prefix(name)
                    .is_some_and(|suffix| suffix.starts_with('{'))
            {
                return None;
            }
            Some(fields.next().unwrap().parse::<f64>().unwrap())
        })
        .sum::<f64>();
    assert!(
        value >= 0.0 && value.fract() == 0.0,
        "metric {} is not a non-negative integer: {}",
        name,
        value,
    );
    value as u64
}

#[test]
fn test_random_cold_region() {
    init_logger();
    let test_started_at = Instant::now();
    let test_secs = env::var("COLD_REGION_RANDOM_TEST_SECS")
        .ok()
        .map(|value| value.parse::<u64>().unwrap())
        .unwrap_or(DEFAULT_TEST_SECS);
    assert!(test_secs <= DEFAULT_TEST_SECS);
    let test_deadline = test_started_at + Duration::from_secs(test_secs);
    let random_deadline = test_deadline
        .checked_sub(FINAL_CHECK_RESERVE)
        .unwrap_or(test_started_at);
    let seed = env::var("COLD_REGION_RANDOM_TEST_SEED")
        .ok()
        .map(|value| value.parse::<u64>().unwrap())
        .unwrap_or_else(rand::random);
    info!("start cold-region random test"; "seed" => seed, "test_secs" => test_secs);
    let mut rng = StdRng::seed_from_u64(seed);

    let routers = Arc::new(Mutex::new(HashMap::new()));
    let hook_routers = routers.clone();
    let nodes = alloc_node_id_vec(NODE_COUNT);
    let mut cluster = ServerClusterBuilder::new(vec![nodes[0]], 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();
    let pd_client = cluster.get_pd_client();
    pd_client.disable_default_operator();
    for &node_id in &nodes[1..] {
        cluster.start_node(node_id, configure_cold_region);
    }
    let mut client = cluster.new_client();
    let bootstrap_region_id = client.get_region_id(&[]);
    let bootstrap_voters = nodes[..3]
        .iter()
        .map(|&node_id| cluster.get_store_id(node_id))
        .collect::<Vec<_>>();
    move_region_to_voter_stores(
        &cluster,
        &pd_client,
        bootstrap_region_id,
        &bootstrap_voters,
        bootstrap_voters[0],
    );
    cluster.wait_region_replicated(&[], 3);

    let keyspace_ids = (0..KEYSPACE_COUNT)
        .map(|_| cluster.keyspace_manager().new_keyspace_id(1))
        .collect::<Vec<_>>();
    cluster.keyspace_manager().create_keyspaces(
        &keyspace_ids,
        keyspace_ids.iter().map(|id| format!("cold-{id}")).collect(),
        &CreateKeyspaceOptions::default(),
        None,
    );
    client.split_keyspaces(keyspace_ids.iter().copied());
    cluster.wait_pd_region_min_count(KEYSPACE_COUNT + 1);

    let mut expected_values = HashMap::new();
    let mut cold_keys = keyspace_ids[..COLD_REGION_COUNT]
        .iter()
        .map(|&keyspace_id| make_row_key(keyspace_id, 1, b"cold"))
        .collect::<Vec<_>>();
    for (idx, key) in cold_keys.iter().take(PERSISTED_COLD_COUNT).enumerate() {
        let value = format!("initial-{idx}").into_bytes();
        put(&mut client, key, &value);
        expected_values.insert(key.clone(), value);
    }

    let hot_key = make_row_key(keyspace_ids[KEYSPACE_COUNT - 2], 1, b"hot");
    let hot_value = b"initial-hot".to_vec();
    put(&mut client, &hot_key, &hot_value);
    expected_values.insert(hot_key.clone(), hot_value);

    let partial_keyspace = keyspace_ids[KEYSPACE_COUNT - 1];
    let partial_left = make_row_key(partial_keyspace, 1, b"a");
    let partial_split = make_row_key(partial_keyspace, 1, b"m");
    let partial_empty = make_row_key(partial_keyspace, 1, b"z");
    // Keep the right half empty to cover ranges pre-split for restore or ingest.
    let partial_value = b"partial-left".to_vec();
    put(&mut client, &partial_left, &partial_value);
    expected_values.insert(partial_left.clone(), partial_value);
    client.try_split(&partial_split, 30).unwrap();
    client.clear_region_cache();

    let cold_regions = cold_keys
        .drain(..)
        .map(|key| ColdRegion {
            region_id: client.get_region_id(&key),
            key,
        })
        .collect::<Vec<_>>();
    let cold_region_ids = Arc::new(
        cold_regions
            .iter()
            .map(|region| region.region_id)
            .collect::<HashSet<_>>(),
    );
    assert_eq!(cold_region_ids.len(), COLD_REGION_COUNT);
    let hot_region_id = client.get_region_id(&hot_key);
    let partial_region_ids = HashSet::from([
        client.get_region_id(&partial_left),
        client.get_region_id(&partial_empty),
    ]);
    assert_eq!(partial_region_ids.len(), 2);
    let tracked_region_ids = Arc::new(
        cold_region_ids
            .iter()
            .copied()
            .chain([hot_region_id])
            .chain(partial_region_ids.iter().copied())
            .collect::<HashSet<_>>(),
    );
    assert_eq!(tracked_region_ids.len(), TRACKED_REGION_COUNT);
    let deployment_region_ids = Arc::new(
        pd_client
            .get_all_regions()
            .iter()
            .map(|region| region.get_id())
            .collect::<HashSet<_>>(),
    );
    assert_eq!(deployment_region_ids.len(), DEPLOYMENT_REGION_COUNT);

    for region in cold_regions.iter().take(PERSISTED_COLD_COUNT) {
        cluster.flush_memtable(region.region_id).unwrap();
        assert!(cluster.wait_for_memtable_flushed(region.region_id, Duration::from_secs(30)));
    }
    for region_id in &partial_region_ids {
        cluster.flush_memtable(*region_id).unwrap();
        assert!(cluster.wait_for_memtable_flushed(*region_id, Duration::from_secs(30)));
    }

    let store_nodes = nodes
        .iter()
        .map(|&node_id| (cluster.get_store_id(node_id), node_id))
        .collect();
    let offline_node = Arc::new(AtomicU16::new(0));
    let mut test = ColdRegionTest {
        cluster,
        client,
        routers: routers.clone(),
        offline_node: offline_node.clone(),
        store_nodes,
        tracked_region_ids: tracked_region_ids.clone(),
        cold_region_ids,
        cold_regions,
        expected_values,
        hot_key,
        value_seq: 0,
    };

    // Verify the deployment sequence from disabled through follower-only to
    // fully enabled before starting disruptive random actions.
    let status_client = StatusClient::new();
    test.refresh_hot_region();
    test.wait_all_regular(deployment_region_ids.clone());
    let initial_metrics = status_client.cold_metrics(&test.cluster, nodes[0]);
    assert_eq!(initial_metrics.cold_regions, 0);

    status_client.update_cold_policy(&test.cluster, &nodes, ColdRegionPolicy::EnableFollower);
    test.wait_follower_only(deployment_region_ids.clone());
    let follower_only_metrics = status_client.cold_metrics(&test.cluster, nodes[0]);
    assert_eq!(follower_only_metrics.cold_regions, 0);
    assert_eq!(
        follower_only_metrics.campaign_events,
        initial_metrics.campaign_events
    );

    status_client.update_cold_policy(&test.cluster, &nodes, ColdRegionPolicy::Enabled);
    test.set_restart_policy(&nodes, ColdRegionPolicy::Enabled);
    test.wait_cold_metrics(
        &status_client,
        nodes[0],
        DEPLOYMENT_COLD_PEER_COUNT,
        Some(initial_metrics.campaign_events),
    );
    test.wait_deployment_enabled(deployment_region_ids.clone());

    let current_action = Arc::new(RwLock::new(format!("seed={seed} initial steady state")));
    let mut monitor = StateMonitor::start(
        routers,
        tracked_region_ids,
        offline_node,
        current_action.clone(),
    );
    let mut required = vec![
        Action::Read,
        Action::Write,
        Action::TransferLeader,
        Action::MovePeer,
        Action::StaleLeaderRead,
        Action::ShortRestart,
        Action::LongRestart,
    ];
    required.shuffle(&mut rng);
    let mut action_counts = HashMap::new();

    while Instant::now() < random_deadline || !required.is_empty() {
        monitor.check();
        let action = required.pop().unwrap_or_else(|| {
            random_action(
                &mut rng,
                random_deadline.saturating_duration_since(Instant::now()),
            )
        });
        let region_idx = rng.gen_range(0..test.cold_regions.len());
        *current_action.write().unwrap() = format!(
            "seed={seed} {action:?} region={} elapsed={:?}",
            test.cold_regions[region_idx].region_id,
            test_started_at.elapsed()
        );
        info!("run cold-region random action";
            "seed" => seed,
            "action" => ?action,
            "region_id" => test.cold_regions[region_idx].region_id,
            "elapsed" => ?test_started_at.elapsed());
        match action {
            Action::Read => test.read(region_idx, rng.gen_bool(0.5)),
            Action::Write => test.write(region_idx),
            Action::TransferLeader => test.transfer_leader(region_idx, &mut rng),
            Action::MovePeer => test.move_peer(region_idx, &mut rng),
            Action::StaleLeaderRead => test.stale_leader_read(region_idx, &mut rng),
            Action::ShortRestart => test.restart(region_idx, false, &mut rng),
            Action::LongRestart => test.restart(region_idx, true, &mut rng),
        }
        *action_counts.entry(action).or_insert(0usize) += 1;
    }

    *current_action.write().unwrap() = format!("seed={seed} final convergence");
    for action in [
        Action::Read,
        Action::Write,
        Action::TransferLeader,
        Action::MovePeer,
        Action::StaleLeaderRead,
        Action::ShortRestart,
        Action::LongRestart,
    ] {
        assert!(action_counts.get(&action).copied().unwrap_or(0) > 0);
    }
    test.wait_all_cold();
    test.refresh_hot_region();
    test.wait_exact_steady_state();
    test.verify_data();

    let final_region_idx = rng.gen_range(0..test.cold_regions.len());
    test.transfer_leader(final_region_idx, &mut rng);
    test.move_peer(final_region_idx, &mut rng);
    test.wait_all_cold();
    test.refresh_hot_region();
    test.wait_exact_steady_state();
    monitor.stop();

    // Downgrade online after the random workload and require every Cold peer to
    // drain. Healthy followers remain TempRegular under regular leader heartbeats.
    *current_action.write().unwrap() = format!("seed={seed} online downgrade");
    let before_downgrade =
        test.wait_cold_metrics(&status_client, nodes[0], DEPLOYMENT_COLD_PEER_COUNT, None);
    status_client.update_cold_policy(&test.cluster, &nodes, ColdRegionPolicy::EnableFollower);
    test.wait_follower_only(deployment_region_ids);
    test.wait_cold_metrics(
        &status_client,
        nodes[0],
        0,
        Some(before_downgrade.campaign_events),
    );
    test.set_restart_policy(&nodes, ColdRegionPolicy::EnableFollower);
    test.verify_data();
    info!("cold-region random test passed";
        "seed" => seed,
        "elapsed" => ?test_started_at.elapsed(),
        "actions" => ?action_counts);
}

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::Disabled;
    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(MAX_LEADER_LEASE);
    cfg.raft_store.split_region_check_tick_interval = ReadableDuration(SPLIT_CHECK_INTERVAL);
    cfg.raft_store.cold_region.max_ts_threshold = ReadableDuration(MAX_TS_THRESHOLD);
    cfg.raft_store.cold_region.store_down_grace_period = ReadableDuration(STORE_DOWN_GRACE_PERIOD);
    cfg.raft_store.cold_region.region_heartbeat_interval =
        ReadableDuration(COLD_REGION_HEARTBEAT_INTERVAL);
    cfg.raft_store.cold_region.leader_heartbeat_interval =
        ReadableDuration(COLD_LEADER_HEARTBEAT_INTERVAL);
    cfg.raft_store.max_leader_missing_duration = ReadableDuration(MAX_LEADER_MISSING_DURATION);
}

fn put(client: &mut ClusterClient, key: &[u8], value: &[u8]) {
    client
        .try_put_kv(
            0..1,
            |_| key.to_vec(),
            |_| value.to_vec(),
            MutateOptions::default(),
        )
        .unwrap();
}

fn get_with_retry(client: &mut ClusterClient, key: &[u8], timeout: Duration) -> Option<Vec<u8>> {
    let started_at = Instant::now();
    loop {
        match client.get_key_version(key, u64::MAX) {
            Ok(value) => return value,
            Err(_) if started_at.elapsed() < timeout => {
                info!("retry cold-region read"; "elapsed" => ?started_at.elapsed());
                client.clear_region_cache();
            }
            Err(err) => panic!("cold-region read failed after {:?}: {:?}", timeout, err),
        }
    }
}

fn put_with_retry(client: &mut ClusterClient, key: &[u8], value: &[u8], timeout: Duration) {
    let started_at = Instant::now();
    loop {
        match client.try_put_kv(
            0..1,
            |_| key.to_vec(),
            |_| value.to_vec(),
            MutateOptions::default(),
        ) {
            Ok(_) => return,
            Err(_) if started_at.elapsed() < timeout => {
                info!("retry cold-region write"; "elapsed" => ?started_at.elapsed());
                client.clear_region_cache();
            }
            Err(err) => panic!("cold-region write failed after {:?}: {:?}", timeout, err),
        }
    }
}

fn query_cold_states(
    routers: &Mutex<HashMap<u16, RaftRouter>>,
    region_ids: Arc<HashSet<u64>>,
    offline_node: u16,
) -> Result<ColdPeerStateSnapshot, String> {
    let routers = routers
        .lock()
        .unwrap()
        .iter()
        .filter(|(node_id, _)| **node_id != offline_node)
        .map(|(&node_id, router)| (node_id, router.clone()))
        .collect::<Vec<_>>();
    let expected = routers.len();
    let (tx, rx) = mpsc::channel();
    for (node_id, router) in routers {
        let tx = tx.clone();
        router.send_store_msg(StoreMsg::QueryColdPeerStates {
            region_ids: region_ids.clone(),
            callback: Box::new(move |snapshot| {
                let _ = tx.send((node_id, snapshot));
            }),
        });
    }
    drop(tx);

    let started_at = Instant::now();
    let mut result = ColdPeerStateSnapshot::default();
    for _ in 0..expected {
        let remaining = STATE_QUERY_TIMEOUT.saturating_sub(started_at.elapsed());
        let (node_id, snapshot) = rx
            .recv_timeout(remaining)
            .map_err(|err| format!("timed out waiting for state from online stores: {err}"))?;
        if snapshot.peer_count
            != snapshot.cold_count + snapshot.regular_count + snapshot.temp_regular_count
        {
            return Err(format!(
                "invalid state counts from node {node_id}: {snapshot:?}"
            ));
        }
        result.peer_count += snapshot.peer_count;
        result.cold_count += snapshot.cold_count;
        result.regular_count += snapshot.regular_count;
        result.temp_regular_count += snapshot.temp_regular_count;
        result.non_cold_peers.extend(snapshot.non_cold_peers);
    }
    Ok(result)
}

fn random_action(rng: &mut StdRng, remaining: Duration) -> Action {
    let action = match rng.gen_range(0..100) {
        0..=24 => Action::Read,
        25..=44 => Action::Write,
        45..=64 => Action::TransferLeader,
        65..=74 => Action::MovePeer,
        75..=84 => Action::StaleLeaderRead,
        85..=92 => Action::ShortRestart,
        _ => Action::LongRestart,
    };
    match action {
        Action::LongRestart if remaining < LONG_RESTART_BUDGET => Action::Read,
        Action::ShortRestart if remaining < SHORT_RESTART_BUDGET => Action::Read,
        _ => action,
    }
}