// Copyright 2026 PingCAP, Inc.

use std::time::Duration;

use kvproto::{kvrpcpb, kvrpcpb::KeyRange, metapb, pdpb};
use matches::assert_matches;
use pd_client::{BucketMeta, PdClient};
use rfstore::store::RegionSnapshot;
use test_cloud_server::{ServerCluster, alloc_node_id_vec};
use test_pd_client::PdClientExt;
use test_raftstore::get_tso;
use test_util::eventually;
use tikv_kv::{Engine, Error, ErrorInner, ExtraRegionOverride, SnapContext, Snapshot, SnapshotExt};
use tikv_util::{codec::bytes::encode_bytes, store::find_peer};

#[test]
fn test_extra_snapshot_override() {
    let node_ids = alloc_node_id_vec(3);
    let cluster = ServerCluster::new(node_ids.clone(), |_, _| {});
    cluster.wait_region_replicated(&[], 3);

    let pd_client = cluster.get_pd_client();
    pd_client.disable_default_operator();

    // Create two regions and wait until the split is visible on every replica.
    let left_key = encode_bytes(b"a");
    let right_key = encode_bytes(b"c");
    let region = pd_client.get_region(&right_key).unwrap();
    pd_client.must_split_region(region, pdpb::CheckPolicy::Usekey, vec![right_key.clone()]);
    cluster.wait_keys_split(&left_key, &right_key);
    cluster.wait_region_replicated(&right_key, 3);

    // Put both region leaders on one store so one RaftKv exercises the override.
    let left = pd_client.get_region_info(&left_key).unwrap();
    let left_leader = left.leader.as_ref().unwrap();
    let mut right = pd_client.get_region_info(&right_key).unwrap();
    right.leader = Some(
        find_peer(&right, left_leader.get_store_id())
            .unwrap()
            .clone(),
    );
    let right_leader = right.leader.as_ref().unwrap();
    pd_client.transfer_leader(right.get_id(), right_leader.clone(), vec![]);
    pd_client.region_leader_must_be(right.get_id(), right_leader.clone());
    let leader_node_id = *node_ids
        .iter()
        .find(|&&node_id| cluster.get_store_id(node_id) == left_leader.get_store_id())
        .unwrap();

    let pb_ctx = {
        let mut ctx = kvrpcpb::Context::default();
        ctx.set_region_id(left.get_id());
        ctx.set_region_epoch(left.get_region_epoch().clone());
        ctx.set_peer(left_leader.clone());
        ctx.set_term(10000);
        ctx
    };
    let mut snap_ctx = SnapContext {
        pb_ctx: &pb_ctx,
        start_ts: Some(get_tso(pd_client.as_ref()).into()),
        key_ranges: vec![KeyRange {
            start_key: encode_bytes(b"c1"),
            end_key: encode_bytes(b"c2"),
            ..Default::default()
        }],
        ..Default::default()
    };
    snap_ctx.extra_region_override = Some(ExtraRegionOverride {
        region_id: right.get_id(),
        region_epoch: right.get_region_epoch().clone(),
        peer: right_leader.clone(),
        check_term: None,
    });
    let mut raft_kv = cluster.get_raft_kv(leader_node_id).unwrap();

    // Wait for the transferred leader's local-read delegate before checking errors.
    let mut snapshot = None;
    eventually(
        Duration::from_millis(10),
        Duration::from_secs(10),
        || match raft_kv.snapshot(snap_ctx.clone()) {
            Ok(snap) => {
                snapshot = Some(snap);
                true
            }
            Err(_) => false,
        },
    );
    let snapshot = snapshot.unwrap();
    let right_term = snapshot.ext().get_term().unwrap().get();
    check_snapshot(snapshot, &right.region);

    // A stale override term must be rejected, while the current term succeeds.
    assert!(right_term >= 2);
    snap_ctx.extra_region_override.as_mut().unwrap().check_term = Some(right_term - 2);
    let err = raft_kv.snapshot(snap_ctx.clone()).err().unwrap();
    assert_matches!(err, Error(box ErrorInner::Request(header)) if header.stale_command.is_some());

    snap_ctx.extra_region_override.as_mut().unwrap().check_term = Some(right_term);
    let snapshot = raft_kv.snapshot(snap_ctx.clone()).unwrap();
    check_snapshot(snapshot, &right.region);
    snap_ctx.extra_region_override.as_mut().unwrap().check_term = None;

    // A stale epoch must be rejected by the overridden region rather than the base
    // region.
    let mut stale_epoch = right.get_region_epoch().clone();
    stale_epoch.version -= 1;
    snap_ctx
        .extra_region_override
        .as_mut()
        .unwrap()
        .region_epoch = stale_epoch;
    let err = raft_kv.snapshot(snap_ctx.clone()).err().unwrap();
    assert_matches!(err, Error(box ErrorInner::Request(header)) if header.epoch_not_match.is_some());
    snap_ctx
        .extra_region_override
        .as_mut()
        .unwrap()
        .region_epoch = right.get_region_epoch().clone();

    // After leadership moves, wait until the old peer deterministically rejects the
    // read.
    let old_leader = right_leader.clone();
    let new_leader = right
        .peers
        .iter()
        .find(|peer| peer.get_id() != old_leader.get_id())
        .unwrap()
        .clone();
    pd_client.transfer_leader(right.get_id(), new_leader.clone(), vec![]);
    pd_client.region_leader_must_be(right.get_id(), new_leader);
    snap_ctx.extra_region_override.as_mut().unwrap().peer = old_leader;
    assert!(!snap_ctx.pb_ctx.replica_read);
    eventually(Duration::from_millis(10), Duration::from_secs(10), || {
        matches!(
            raft_kv.snapshot(snap_ctx.clone()),
            Err(Error(box ErrorInner::Request(header))) if header.not_leader.is_some()
        )
    });
}

fn check_snapshot(snapshot: RegionSnapshot, region: &metapb::Region) {
    let BucketMeta {
        region_id,
        region_epoch,
        ..
    } = snapshot.bucket_meta.unwrap().as_ref().clone();
    assert_eq!(region_id, region.get_id());
    assert_eq!(region_epoch, region.get_region_epoch().clone());
}