// Copyright 2026 PingCAP, Inc.

use std::{
    collections::{BTreeMap, BTreeSet},
    sync::Arc,
    time::Duration,
};

use api_version::ApiV2;
use futures::executor::block_on;
use kvengine::{Engine, ShardStats};
use merged_engine::{KeyspaceShardMetaScope, KeyspaceShardMetas};
use native_br::{backup, backup_worker};
use pd_client::PdClient;
use replication_worker::{CdcMsg, ReplicationScheduler};
use test_cloud_server::{
    ServerCluster,
    client::ClusterClient,
    keyspace::{CreateKeyspaceOptions, make_row_key},
    move_region_to_voter_stores, must_wait_result,
    ticdc::KeyspaceTestService,
    util::RawRegion,
};
use tikv_util::{config::ReadableDuration, warn};

pub(crate) const WAIT_SECS: usize = 60;
pub(crate) const REGION_REPLICA_COUNT: usize = 3;

#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct StableShardStats {
    pub(crate) id: u64,
    pub(crate) ver: u64,
    pub(crate) kv_size: u64,
    pub(crate) entries: usize,
    pub(crate) max_ts: u64,
}

impl From<ShardStats> for StableShardStats {
    fn from(stats: ShardStats) -> Self {
        Self {
            id: stats.id,
            ver: stats.ver,
            kv_size: stats.kv_size,
            entries: stats.entries,
            max_ts: stats.max_ts,
        }
    }
}

pub(crate) fn create_keyspace_with_one_table(
    runtime: &tokio::runtime::Runtime,
    cluster: &ServerCluster,
    client: &mut ClusterClient,
    keyspace_name: &str,
) -> (u32, i64) {
    let keyspace_id = cluster.keyspace_manager().new_keyspace_id(1);
    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        keyspace_id,
        keyspace_name.to_owned(),
        &CreateKeyspaceOptions {
            table_count: 1,
            schema_enable_ratio: 0.0,
            ..Default::default()
        },
        false,
    ));

    client.split_keyspace(keyspace_id);
    let (start, end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
    if !start.is_empty() {
        cluster.wait_region_replicated(&start, REGION_REPLICA_COUNT);
    }
    if !end.is_empty() {
        cluster.wait_region_replicated(&end, REGION_REPLICA_COUNT);
    }

    let keyspace = cluster
        .keyspace_manager()
        .get_keyspace_meta(keyspace_id)
        .unwrap();
    let tables = keyspace.get_all_available_tables();
    assert_eq!(
        tables.len(),
        1,
        "unexpected tables for keyspace {}",
        keyspace_id
    );
    (keyspace_id, tables[0])
}

pub(crate) fn split_table_rows(
    cluster: &ServerCluster,
    client: &mut ClusterClient,
    keyspace_id: u32,
    table_id: i64,
    split_points: &[u32],
) {
    for &split_idx in split_points {
        let split_key = make_row_key(keyspace_id, table_id, format!("k{split_idx:08}").as_bytes());
        client
            .try_split(&split_key, WAIT_SECS)
            .expect("split_table_rows");
        cluster.wait_region_replicated(&split_key, REGION_REPLICA_COUNT);
    }
}

pub(crate) fn scan_keyspace_regions_upstream(
    client: &mut ClusterClient,
    keyspace_id: u32,
) -> Vec<RawRegion> {
    let (start, end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
    client.clear_region_cache();

    let mut regions = vec![];
    let first_region = client.get_region_by_key(&start);
    if end.is_empty() || first_region.raw_start.as_slice() < end.as_slice() {
        let scan_start = first_region.raw_end.clone();
        regions.push(first_region);
        if end.is_empty() || scan_start.as_slice() < end.as_slice() {
            regions.extend(client.scan_regions(scan_start.as_slice(), &end).unwrap());
        }
    }
    regions
}

pub(crate) fn region_ids(regions: &[RawRegion]) -> Vec<u64> {
    let mut ids = regions.iter().map(|region| region.id).collect::<Vec<_>>();
    ids.sort_unstable();
    ids
}

pub(crate) fn assert_regions_on_store_group(regions: &[RawRegion], expected_store_ids: &[u64]) {
    let expected = expected_store_ids.iter().copied().collect::<BTreeSet<_>>();
    for region in regions {
        let actual = region
            .peers
            .iter()
            .filter(|peer| peer.get_role() != kvproto::metapb::PeerRole::Learner)
            .map(|peer| peer.get_store_id())
            .collect::<BTreeSet<_>>();
        assert_eq!(
            actual, expected,
            "region {} on unexpected stores: {:?}",
            region.id, region
        );
    }
}

pub(crate) fn move_regions_to_group(
    cluster: &ServerCluster,
    pd_client: &test_pd_client::TestPdClient,
    regions: &[RawRegion],
    target_store_ids: &[u64],
    leader_store_id: u64,
) {
    for region in regions {
        move_region_to_voter_stores(
            cluster,
            pd_client,
            region.id,
            target_store_ids,
            leader_store_id,
        );
    }
}

pub(crate) fn move_region_ids_to_group(
    cluster: &ServerCluster,
    pd_client: &test_pd_client::TestPdClient,
    region_ids: &[u64],
    target_store_ids: &[u64],
    leader_store_id: u64,
) {
    for &region_id in region_ids {
        move_region_to_voter_stores(
            cluster,
            pd_client,
            region_id,
            target_store_ids,
            leader_store_id,
        );
    }
}

pub(crate) fn region_versions(regions: &[RawRegion]) -> BTreeMap<u64, u64> {
    regions
        .iter()
        .map(|region| (region.id, region.epoch.version))
        .collect()
}

pub(crate) fn load_keyspace_shard_metas_via_scheduler(
    scheduler: &ReplicationScheduler,
    keyspace_id: u32,
) -> std::result::Result<KeyspaceShardMetas, String> {
    let (cb, rx) = tikv_util::mpsc::paired_callback();
    scheduler.schedule(CdcMsg::LoadKeyspaceShardMetas { keyspace_id, cb });
    rx.recv()
        .map_err(|_| format!("load keyspace {} metas canceled", keyspace_id))?
        .map_err(|err| format!("load keyspace {} metas failed: {:?}", keyspace_id, err))
}

pub(crate) fn wait_keyspace_shard_metas_ready(
    scheduler: &ReplicationScheduler,
    keyspace_id: u32,
    expected_versions: &BTreeMap<u64, u64>,
) {
    must_wait_result(
        || {
            let metas = load_keyspace_shard_metas_via_scheduler(scheduler, keyspace_id)?;
            if metas.scope != KeyspaceShardMetaScope::All {
                return Err(format!(
                    "keyspace {} scope mismatch: expected all, got {:?}",
                    keyspace_id, metas.scope
                ));
            }
            let actual_versions = metas
                .metas
                .iter()
                .map(|(&region_id, meta)| (region_id, meta.ver))
                .collect::<BTreeMap<_, _>>();
            let unexpected = actual_versions
                .keys()
                .filter(|region_id| !expected_versions.contains_key(region_id))
                .copied()
                .collect::<Vec<_>>();
            if !unexpected.is_empty() {
                warn!(
                    "ignore unexpected shard metas returned by replication worker";
                    "keyspace" => keyspace_id,
                    "unexpected_regions" => ?unexpected,
                    "actual" => ?actual_versions,
                );
            }

            let filtered_actual_versions = actual_versions
                .iter()
                .filter_map(|(&region_id, &ver)| {
                    expected_versions
                        .contains_key(&region_id)
                        .then_some((region_id, ver))
                })
                .collect::<BTreeMap<_, _>>();
            if filtered_actual_versions == *expected_versions {
                return Ok(());
            }

            let missing = expected_versions
                .iter()
                .filter_map(|(&region_id, &expected_ver)| {
                    (!filtered_actual_versions.contains_key(&region_id))
                        .then_some((region_id, expected_ver))
                })
                .collect::<Vec<_>>();
            let differing = expected_versions
                .iter()
                .filter_map(|(&region_id, &expected_ver)| {
                    let actual_ver = filtered_actual_versions.get(&region_id).copied()?;
                    (actual_ver != expected_ver).then_some((region_id, expected_ver, actual_ver))
                })
                .collect::<Vec<_>>();
            Err(format!(
                "keyspace {} shard metas mismatch for whole-keyspace regions after filtering unexpected regions: missing {:?}, unexpected {:?}, differing {:?}, actual {:?}, filtered_actual {:?}",
                keyspace_id,
                missing,
                unexpected,
                differing,
                actual_versions,
                filtered_actual_versions
            ))
        },
        WAIT_SECS,
        || format!("wait keyspace {} shard metas ready", keyspace_id),
    );
}

pub(crate) fn add_keyspace_via_scheduler(
    scheduler: &ReplicationScheduler,
    keyspace_id: u32,
    keyspace_svc: &KeyspaceTestService,
) {
    let (cb, rx) = tikv_util::mpsc::paired_callback();
    scheduler.schedule(CdcMsg::AddKeyspace {
        keyspace_id,
        pd_url: keyspace_svc.pd_url().to_owned(),
        cdc_addr: keyspace_svc.cdc_addr().to_owned(),
        cb,
    });
    rx.recv()
        .unwrap_or_else(|_| panic!("add keyspace {} canceled", keyspace_id))
        .unwrap_or_else(|err| panic!("add keyspace {} failed: {:?}", keyspace_id, err));
}

pub(crate) fn wait_worker_store_registered(
    keyspace_svc: &KeyspaceTestService,
    merged_store_id: u64,
) {
    must_wait_result(
        || {
            let stores = keyspace_svc
                .pd_client_ext()
                .get_all_stores(true)
                .map_err(|err| format!("{err:?}"))?;
            if stores.iter().any(|store| store.get_id() == merged_store_id) {
                Ok(())
            } else {
                Err(format!("store {} not registered yet", merged_store_id))
            }
        },
        WAIT_SECS,
        || {
            format!(
                "wait replication worker store {} registered",
                merged_store_id
            )
        },
    );
}

pub(crate) fn collect_upstream_shard_stats(
    cluster: &ServerCluster,
    regions: &[RawRegion],
) -> std::result::Result<BTreeMap<u64, StableShardStats>, String> {
    let pd_client = cluster.get_pd_client();
    let mut stats = BTreeMap::new();
    for region in regions {
        let (_, leader) = block_on(pd_client.get_region_leader_by_id(region.id))
            .map_err(|err| format!("get region {} leader failed: {:?}", region.id, err))?
            .ok_or_else(|| format!("region {} missing leader", region.id))?;
        let leader_store_id = leader.get_store_id();
        let Some(node_id) = cluster.get_server_node_id(leader_store_id) else {
            return Err(format!(
                "leader store {} for region {} is not running",
                leader_store_id, region.id
            ));
        };
        let Some(shard_stats) = cluster.get_kvengine(node_id).get_shard_stat_opt(region.id) else {
            return Err(format!(
                "upstream shard stats missing for region {} on leader store {}",
                region.id, leader_store_id
            ));
        };
        stats.insert(region.id, shard_stats.into());
    }
    Ok(stats)
}

pub(crate) fn collect_downstream_shard_stats(
    rep_kv: &Engine,
    region_ids: &BTreeSet<u64>,
) -> std::result::Result<BTreeMap<u64, StableShardStats>, String> {
    let mut stats = BTreeMap::new();
    let mut missing = Vec::new();
    for &region_id in region_ids {
        let Some(shard_stats) = rep_kv.get_shard_stat_opt(region_id) else {
            missing.push(region_id);
            continue;
        };
        stats.insert(region_id, shard_stats.into());
    }
    if !missing.is_empty() {
        return Err(format!(
            "downstream shard stats missing for regions {:?}",
            missing
        ));
    }
    Ok(stats)
}

pub(crate) fn assert_downstream_shards_absent(rep_kv: &Engine, region_ids: &BTreeSet<u64>) {
    let present = region_ids
        .iter()
        .copied()
        .filter(|region_id| rep_kv.get_shard_stat_opt(*region_id).is_some())
        .collect::<Vec<_>>();
    assert!(
        present.is_empty(),
        "unexpected downstream shards already loaded: {:?}",
        present
    );
}

pub(crate) fn assert_shard_stats_eventually_match(
    cluster: &ServerCluster,
    rep_kv: &Engine,
    regions: &[RawRegion],
) {
    let expected_region_ids = regions
        .iter()
        .map(|region| region.id)
        .collect::<BTreeSet<_>>();
    must_wait_result(
        || {
            let upstream = collect_upstream_shard_stats(cluster, regions)?;
            let downstream = collect_downstream_shard_stats(rep_kv, &expected_region_ids)?;
            if upstream != downstream {
                return Err(format!(
                    "shard stats mismatch, upstream: {:?}, downstream: {:?}",
                    upstream, downstream
                ));
            }
            Ok(())
        },
        WAIT_SECS,
        || "wait whole-keyspace shard stats to match".to_owned(),
    );
}

pub(crate) fn assert_shard_stats_eventually_differ(
    cluster: &ServerCluster,
    rep_kv: &Engine,
    regions: &[RawRegion],
) {
    let expected_region_ids = regions
        .iter()
        .map(|region| region.id)
        .collect::<BTreeSet<_>>();
    must_wait_result(
        || {
            let upstream = collect_upstream_shard_stats(cluster, regions)?;
            let downstream = collect_downstream_shard_stats(rep_kv, &expected_region_ids)?;
            if upstream == downstream {
                return Err(format!(
                    "shard stats still match unexpectedly, upstream: {:?}, downstream: {:?}",
                    upstream, downstream
                ));
            }
            Ok(())
        },
        WAIT_SECS,
        || "wait whole-keyspace shard stats to differ".to_owned(),
    );
}

pub(crate) fn take_lightweight_backup(
    runtime: &tokio::runtime::Runtime,
    dfs_config: &kvengine::dfs::DFSConfig,
    pd_client: Arc<dyn PdClient>,
) {
    let backup_worker = backup_worker::BackupWorker::new(
        backup::BackupConfig {
            dfs: dfs_config.clone(),
            skip_keyspace_meta: true,
            backup_delay: ReadableDuration::millis(1),
            ..Default::default()
        },
        pd_client,
        Duration::ZERO,
        Duration::from_secs(1),
    );
    runtime.block_on(backup_worker.instant_backup()).unwrap();
}