// Copyright 2026 PingCAP, Inc.

use std::collections::{BTreeMap, BTreeSet};

use futures::executor::block_on;
use kvproto::metapb;
use pd_client::PdClient;
use raftstore::coprocessor::RegionInfoProvider;
use test_pd_client::{PdClientExt, TestPdClient};
use tikv_util::store::{new_learner_peer, new_peer};

use crate::{ServerCluster, try_wait_result};

/// Moves a region to the target voter stores and transfers leadership to the
/// specified store.
///
/// Each new target store is added as a learner first, then promoted to a voter
/// with the same peer id. The helper waits for each intermediate state to
/// apply, for pending peers to clear, for the leader transfer to finish, and
/// for removed peers to disappear from local kvengine and `RegionInfoAccessor`.
pub fn move_region_to_voter_stores(
    cluster: &ServerCluster,
    pd_client: &TestPdClient,
    region_id: u64,
    target_voter_store_ids: &[u64],
    leader_store_id: u64,
) -> metapb::Region {
    let target_voters: BTreeSet<_> = target_voter_store_ids.iter().copied().collect();
    assert_eq!(
        target_voters.len(),
        target_voter_store_ids.len(),
        "duplicate target voters: {:?}",
        target_voter_store_ids
    );
    assert!(
        target_voters.contains(&leader_store_id),
        "leader store {} is not in target voters {:?}",
        leader_store_id,
        target_voter_store_ids
    );

    let mut region = block_on(pd_client.get_region_by_id(region_id))
        .unwrap()
        .unwrap_or_else(|| panic!("region {} not found", region_id));
    let mut expected_voters: BTreeSet<u64> = region
        .peers
        .iter()
        .filter(|peer| peer.get_role() != metapb::PeerRole::Learner)
        .map(|peer| peer.get_store_id())
        .collect();
    let mut expected_learners: BTreeSet<u64> = region
        .peers
        .iter()
        .filter(|peer| peer.get_role() == metapb::PeerRole::Learner)
        .map(|peer| peer.get_store_id())
        .collect();
    let mut absent_store_ids = Vec::new();

    for &store_id in &target_voters {
        let learner_peer = match region
            .peers
            .iter()
            .find(|peer| peer.get_store_id() == store_id)
        {
            Some(peer) => match peer.get_role() {
                metapb::PeerRole::Voter => None,
                metapb::PeerRole::Learner => {
                    let expected_voters = expected_voters.iter().copied().collect::<Vec<_>>();
                    let expected_learners = expected_learners.iter().copied().collect::<Vec<_>>();
                    region = wait_region_state(
                        cluster,
                        pd_client,
                        region_id,
                        &expected_voters,
                        &expected_learners,
                        None,
                        &absent_store_ids,
                    );
                    Some(
                        region
                            .peers
                            .iter()
                            .find(|peer| peer.get_store_id() == store_id)
                            .unwrap_or_else(|| {
                                panic!(
                                    "target learner {} missing from region {:?}",
                                    store_id, region
                                )
                            })
                            .clone(),
                    )
                }
                role => panic!(
                    "store {} has unexpected role {:?} in region {:?}",
                    store_id, role, region
                ),
            },
            None => {
                let peer = new_learner_peer(store_id, pd_client.alloc_id().unwrap());
                pd_client.add_peer(region_id, peer);
                expected_learners.insert(store_id);
                let expected_voters = expected_voters.iter().copied().collect::<Vec<_>>();
                let expected_learners = expected_learners.iter().copied().collect::<Vec<_>>();
                region = wait_region_state(
                    cluster,
                    pd_client,
                    region_id,
                    &expected_voters,
                    &expected_learners,
                    None,
                    &absent_store_ids,
                );
                Some(
                    region
                        .peers
                        .iter()
                        .find(|peer| peer.get_store_id() == store_id)
                        .unwrap_or_else(|| {
                            panic!(
                                "target learner {} missing from region {:?}",
                                store_id, region
                            )
                        })
                        .clone(),
                )
            }
        };
        if let Some(learner_peer) = learner_peer {
            pd_client.must_none_pending_peer(learner_peer.clone());
            pd_client.add_peer(region_id, new_peer(store_id, learner_peer.get_id()));
            expected_learners.remove(&store_id);
            expected_voters.insert(store_id);
            let expected_voters = expected_voters.iter().copied().collect::<Vec<_>>();
            let expected_learners = expected_learners.iter().copied().collect::<Vec<_>>();
            region = wait_region_state(
                cluster,
                pd_client,
                region_id,
                &expected_voters,
                &expected_learners,
                None,
                &absent_store_ids,
            );
        }
    }

    let expected_voters = expected_voters.iter().copied().collect::<Vec<_>>();
    let expected_learners = expected_learners.iter().copied().collect::<Vec<_>>();

    let leader_peer = region
        .peers
        .iter()
        .find(|peer| peer.get_store_id() == leader_store_id)
        .unwrap_or_else(|| {
            panic!(
                "leader store {} missing from region {:?}",
                leader_store_id, region
            )
        })
        .clone();
    pd_client.must_transfer_leader(region_id, leader_peer);
    region = wait_region_state(
        cluster,
        pd_client,
        region_id,
        &expected_voters,
        &expected_learners,
        Some(leader_store_id),
        &absent_store_ids,
    );

    let mut current_voters: BTreeSet<_> = expected_voters.iter().copied().collect();
    let mut current_learners: BTreeSet<_> = expected_learners.iter().copied().collect();
    let removable_peers = region
        .peers
        .iter()
        .filter(|peer| !target_voters.contains(&peer.get_store_id()))
        .map(|peer| peer.get_store_id())
        .collect::<Vec<_>>();
    for old_store_id in removable_peers {
        let peer = region
            .peers
            .iter()
            .find(|peer| peer.get_store_id() == old_store_id)
            .unwrap_or_else(|| {
                panic!(
                    "old voter {} missing from region {:?}",
                    old_store_id, region
                )
            })
            .clone();
        pd_client.remove_peer(region_id, peer);
        current_voters.remove(&old_store_id);
        current_learners.remove(&old_store_id);
        absent_store_ids.push(old_store_id);
        let current_voters = current_voters.iter().copied().collect::<Vec<_>>();
        let current_learners = current_learners.iter().copied().collect::<Vec<_>>();
        region = wait_region_state(
            cluster,
            pd_client,
            region_id,
            &current_voters,
            &current_learners,
            Some(leader_store_id),
            &absent_store_ids,
        );
    }

    region
}

/// Waits until PD, local kvengine state, and local `RegionInfoAccessor`
/// converge to the expected peer layout.
pub fn wait_region_state(
    cluster: &ServerCluster,
    pd_client: &TestPdClient,
    region_id: u64,
    expected_voter_store_ids: &[u64],
    expected_learner_store_ids: &[u64],
    expected_leader_store_id: Option<u64>,
    absent_store_ids: &[u64],
) -> metapb::Region {
    let expected_voters: BTreeSet<_> = expected_voter_store_ids.iter().copied().collect();
    let expected_learners: BTreeSet<_> = expected_learner_store_ids.iter().copied().collect();
    assert!(
        expected_voters.is_disjoint(&expected_learners),
        "expected voters and learners overlap: voters={:?}, learners={:?}",
        expected_voter_store_ids,
        expected_learner_store_ids
    );

    let mut expected_peer_roles = BTreeMap::new();
    for &store_id in &expected_voters {
        expected_peer_roles.insert(store_id, metapb::PeerRole::Voter);
    }
    for &store_id in &expected_learners {
        expected_peer_roles.insert(store_id, metapb::PeerRole::Learner);
    }
    let expected_present_store_ids = expected_peer_roles.keys().copied().collect::<Vec<_>>();

    let fail_msg = || {
        format!(
            "region {} did not reach expected state: voters={:?}, learners={:?}, leader={:?}, absent={:?}",
            region_id,
            expected_voter_store_ids,
            expected_learner_store_ids,
            expected_leader_store_id,
            absent_store_ids
        )
    };
    try_wait_result(
        || {
            let (region, leader) = block_on(pd_client.get_region_leader_by_id(region_id))
                .map_err(|err| format!("get_region_leader_by_id({region_id}) failed: {err:?}"))?
                .ok_or_else(|| format!("region {} not found", region_id))?;
            let actual_peer_roles = region_peer_roles(&region);
            if actual_peer_roles != expected_peer_roles {
                return Err(format!(
                    "region {} peer set mismatch: expected {:?}, got {:?}",
                    region_id, expected_peer_roles, actual_peer_roles
                ));
            }
            if let Some(store_id) = expected_leader_store_id {
                if leader.get_store_id() != store_id {
                    return Err(format!(
                        "region {} leader mismatch: expected {}, got {:?}",
                        region_id, store_id, leader
                    ));
                }
            }

            let region_ver = region.get_region_epoch().get_version();
            let region_conf_ver = region.get_region_epoch().get_conf_ver();
            for &store_id in &expected_present_store_ids {
                let Some(node_id) = cluster.get_server_node_id(store_id) else {
                    return Err(format!(
                        "store {} is not backed by a running server",
                        store_id
                    ));
                };
                let kv = cluster.get_kvengine(node_id);
                if kv.get_shard_with_ver(region_id, region_ver).is_err() {
                    return Err(format!(
                        "store {} has not applied region {} version {} yet",
                        store_id, region_id, region_ver
                    ));
                }
                let accessor = cluster.get_region_info_accessor(node_id);
                let Some(local_region_info) = accessor.find_region_by_id(region_id) else {
                    return Err(format!(
                        "store {} has not published region {} to RegionInfoAccessor yet",
                        store_id, region_id
                    ));
                };
                let local_peer_roles = region_peer_roles(&local_region_info.region);
                if local_region_info.region.get_region_epoch().get_conf_ver() != region_conf_ver {
                    return Err(format!(
                        "store {} sees conf_ver {} for region {}, expected {}",
                        store_id,
                        local_region_info.region.get_region_epoch().get_conf_ver(),
                        region_id,
                        region_conf_ver
                    ));
                }
                if local_peer_roles != expected_peer_roles {
                    return Err(format!(
                        "store {} sees peer set {:?} for region {}, expected {:?}",
                        store_id, local_peer_roles, region_id, expected_peer_roles
                    ));
                }
                if expected_leader_store_id == Some(store_id)
                    && (!local_region_info.is_leader()
                        || local_region_info
                            .leader_peer()
                            .is_none_or(|leader| leader.get_store_id() != store_id))
                {
                    return Err(format!(
                        "store {} does not see itself as leader for region {}",
                        store_id, region_id
                    ));
                }
            }
            for &store_id in absent_store_ids {
                let Some(node_id) = cluster.get_server_node_id(store_id) else {
                    return Err(format!(
                        "store {} is not backed by a running server",
                        store_id
                    ));
                };
                let kv = cluster.get_kvengine(node_id);
                if kv.get_shard(region_id).is_some() {
                    return Err(format!(
                        "store {} still has region {} locally",
                        store_id, region_id
                    ));
                }
                if cluster
                    .get_region_info_accessor(node_id)
                    .find_region_by_id(region_id)
                    .is_some()
                {
                    return Err(format!(
                        "store {} still reports region {} in RegionInfoAccessor",
                        store_id, region_id
                    ));
                }
            }
            Ok(region)
        },
        45,
    )
    .unwrap_or_else(|last_err| panic!("{}: {}", fail_msg(), last_err))
}

fn region_peer_roles(region: &metapb::Region) -> BTreeMap<u64, metapb::PeerRole> {
    region
        .peers
        .iter()
        .map(|peer| (peer.get_store_id(), peer.get_role()))
        .collect()
}