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

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

use futures::executor::block_on;
use kvproto::{
    kvrpcpb::Op,
    metapb::{Peer, Region},
};
use pd_client::PdClient;
use raftstore::coprocessor::RegionInfoProvider;
use rfstore::{
    RaftRouter, RaftStoreRouter,
    store::{
        Callback, ColdPeerStateKind, ColdRegionPolicy, PeerMsg, ReadResponse, StoreMsg,
        metrics::COLD_CAMPAIGN_EVENT_COUNTER_VEC,
    },
};
use test_cloud_server::{
    ServerCluster, ServerClusterBuilder, alloc_node_id_vec,
    client::{ClusterClient, TxnMutations},
    keyspace::{CreateKeyspaceOptions, make_row_key},
    util::Mutation,
};
use test_raftstore::{new_read_index_cmd, new_request};
use tikv::config::TikvConfig;
use tikv_util::config::ReadableDuration;

use crate::cases::{FailpointGuard, cfg_notify_once, i_to_key};

const NODE_COUNT: usize = 3;
const BASE_TICK: Duration = Duration::from_millis(50);
const ELECTION_TIMEOUT: Duration = Duration::from_millis(500);
const STATE_WAIT_TIMEOUT: Duration = Duration::from_secs(20);
const STATE_QUERY_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.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 new_cluster() -> (
    ServerCluster,
    Vec<u16>,
    Arc<Mutex<HashMap<u16, RaftRouter>>>,
) {
    let nodes = alloc_node_id_vec(NODE_COUNT);
    let routers = Arc::new(Mutex::new(HashMap::new()));
    let hook_routers = routers.clone();
    let 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);
    (cluster, nodes, routers)
}

fn create_cold_region(cluster: &ServerCluster, client: &mut ClusterClient) -> (Region, Vec<u8>) {
    let keyspace_id = cluster.keyspace_manager().new_keyspace_id(1);
    cluster.keyspace_manager().create_keyspaces(
        &[keyspace_id],
        vec!["cold-failpoint".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);
    let region = block_on(cluster.get_pd_client().get_region_by_id(region_id))
        .unwrap()
        .unwrap();
    (region, key)
}

fn wait_until(mut condition: impl FnMut() -> bool, description: &str) {
    let started_at = Instant::now();
    while started_at.elapsed() < STATE_WAIT_TIMEOUT {
        if condition() {
            return;
        }
        thread::sleep(STATE_QUERY_INTERVAL);
    }
    panic!(
        "{} did not become true within {:?}",
        description, STATE_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 cold_peer_state(router: &RaftRouter, region_id: u64) -> Option<ColdPeerStateKind> {
    let (tx, rx) = mpsc::channel();
    router.send_store_msg(StoreMsg::QueryColdPeerStates {
        region_ids: Arc::new(HashSet::from([region_id])),
        callback: Box::new(move |snapshot| tx.send(snapshot).unwrap()),
    });
    let snapshot = rx.recv_timeout(STATE_WAIT_TIMEOUT).unwrap();
    assert_eq!(snapshot.peer_count, 1, "{:?}", snapshot);
    if snapshot.cold_count == 1 {
        return None;
    }
    Some(snapshot.non_cold_peers[0].state)
}

fn wait_all_peers_cold(
    cluster: &ServerCluster,
    nodes: &[u16],
    routers: &Mutex<HashMap<u16, RaftRouter>>,
    region_id: u64,
) {
    let leader = cluster.leader_of_region(region_id).unwrap();
    let leader_node = cluster.get_server_node_id(leader.get_store_id()).unwrap();
    wait_until(
        || {
            routers.lock().unwrap()[&leader_node]
                .send_peer_msg_for_test(region_id, PeerMsg::ColdRegionTick(1));
            peer_is_cold(cluster, leader_node, region_id)
        },
        "leader peer cold state",
    );

    wait_until(
        || {
            for &node_id in nodes {
                if node_id != leader_node && !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 peer cold states",
    );
}

fn prewrite(client: &mut ClusterClient, key: Vec<u8>) -> kvproto::kvrpcpb::PrewriteResponse {
    let mut mutation = Mutation::default();
    mutation.set_op(Op::Put);
    mutation.set_key(key.clone());
    mutation.set_value(b"value".to_vec());
    let start_ts = client.get_ts();
    client
        .kv_prewrite_single_region_without_retry(
            key.into(),
            None,
            TxnMutations::from_normal(vec![mutation]),
            start_ts,
        )
        .unwrap()
}

fn send_read_index(
    router: &RaftRouter,
    region: &Region,
    peer: Peer,
) -> mpsc::Receiver<ReadResponse> {
    let mut request = new_request(
        region.get_id(),
        region.get_region_epoch().clone(),
        vec![new_read_index_cmd()],
        true,
    );
    request.mut_header().set_peer(peer);
    let (tx, rx) = mpsc::channel();
    router.send_command(
        request,
        Default::default(),
        Callback::read(Box::new(move |response| tx.send(response).unwrap())),
    );
    rx
}

fn wait_failpoint(
    receiver: tokio::sync::oneshot::Receiver<tokio::sync::oneshot::Sender<()>>,
) -> tokio::sync::oneshot::Sender<()> {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_time()
        .build()
        .unwrap();
    runtime
        .block_on(async { tokio::time::timeout(STATE_WAIT_TIMEOUT, receiver).await })
        .expect("failpoint was not reached")
        .unwrap()
}

#[test]
fn test_scheduler_accepts_cold_leader() {
    let (cluster, nodes, routers) = new_cluster();
    let mut client = cluster.new_client();
    let _guard = FailpointGuard::new().cfg("check_leader_peer_is_leader", "return(false)");

    let response = prewrite(&mut client, i_to_key(1));
    let error = response.get_region_error();
    assert!(error.has_not_leader(), "{:?}", error);
    assert!(error.get_not_leader().has_leader(), "{:?}", error);

    let (region, key) = create_cold_region(&cluster, &mut client);
    wait_all_peers_cold(&cluster, &nodes, &routers, region.get_id());
    client.clear_region_cache();

    let response = prewrite(&mut client, key);
    assert!(!response.has_region_error(), "{:?}", response);
    assert!(response.get_errors().is_empty(), "{:?}", response);
}

#[test]
fn test_cold_read_index_timeout_lifecycle() {
    let (cluster, nodes, routers) = new_cluster();
    let mut client = cluster.new_client();
    let (mut region, _) = create_cold_region(&cluster, &mut client);
    let region_id = region.get_id();
    wait_all_peers_cold(&cluster, &nodes, &routers, region_id);

    let leader = cluster.leader_of_region(region_id).unwrap();
    let follower = region
        .get_peers()
        .iter()
        .find(|peer| peer.get_id() != leader.get_id())
        .unwrap()
        .clone();
    let follower_node = cluster.get_server_node_id(follower.get_store_id()).unwrap();
    let follower_router = routers.lock().unwrap()[&follower_node].clone();

    let campaign_count = COLD_CAMPAIGN_EVENT_COUNTER_VEC.read_index.get();
    let (pause_receiver, pause_guard) = cfg_notify_once("on_step_read_index_msg");
    let response_receiver = send_read_index(&follower_router, &region, follower.clone());
    wait_failpoint(pause_receiver).send(()).unwrap();
    let response = response_receiver.recv_timeout(STATE_WAIT_TIMEOUT).unwrap();
    assert!(
        !response.response.get_header().has_error(),
        "{:?}",
        response
    );
    drop(pause_guard);

    let observation_deadline = Instant::now() + ELECTION_TIMEOUT * 2;
    while Instant::now() < observation_deadline {
        assert_ne!(
            cold_peer_state(&follower_router, region_id),
            Some(ColdPeerStateKind::Regular),
            "successful ReadIndex triggered a stale campaign timeout"
        );
        thread::sleep(STATE_QUERY_INTERVAL);
    }
    assert_eq!(
        COLD_CAMPAIGN_EVENT_COUNTER_VEC.read_index.get(),
        campaign_count,
        "successful ReadIndex triggered a stale campaign timeout"
    );
    wait_all_peers_cold(&cluster, &nodes, &routers, region_id);

    let campaign_count = COLD_CAMPAIGN_EVENT_COUNTER_VEC.read_index.get();
    region = block_on(cluster.get_pd_client().get_region_by_id(region_id))
        .unwrap()
        .unwrap();
    let leader = cluster.leader_of_region(region_id).unwrap();
    let follower = region
        .get_peers()
        .iter()
        .find(|peer| peer.get_id() != leader.get_id())
        .unwrap()
        .clone();
    let follower_node = cluster.get_server_node_id(follower.get_store_id()).unwrap();
    let follower_router = routers.lock().unwrap()[&follower_node].clone();

    let (pause_receiver, pause_guard) = cfg_notify_once("on_step_read_index_msg");
    let response_receiver = send_read_index(&follower_router, &region, follower);
    let release = wait_failpoint(pause_receiver);
    let response = response_receiver.recv_timeout(STATE_WAIT_TIMEOUT).unwrap();
    assert!(
        response
            .response
            .get_header()
            .get_error()
            .has_stale_command(),
        "{:?}",
        response
    );
    wait_until(
        || COLD_CAMPAIGN_EVENT_COUNTER_VEC.read_index.get() > campaign_count,
        "timed-out follower campaign",
    );
    release.send(()).unwrap();
    drop(pause_guard);

    wait_all_peers_cold(&cluster, &nodes, &routers, region_id);
}