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

use kvproto::{
    kvrpcpb::{
        self, CheckSecondaryLocksRequest, CleanupRequest, PessimisticRollbackRequest,
        PrewriteRequest, PrewriteRequestPessimisticAction, TxnHeartBeatRequest,
    },
    tikvpb::TikvClient,
};
use test_cloud_server::{
    ServerCluster,
    client::{ClusterClient, TxnMutations},
};
use txn_types::{TimeStamp, WriteType};

use super::helper::{
    must_not_written, must_unlocked, must_written_with_properties, new_lock_mutation,
    new_shared_lock_mutation,
};
use crate::{alloc_node_id_vec, i_to_key};

fn new_single_node_cluster() -> (ServerCluster, ClusterClient) {
    let cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
        cfg.pessimistic_txn.pipelined = false;
    });
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client();
    let keyspace_id = api_version::ApiV2::get_u32_keyspace_id_by_key(&i_to_key(0)).unwrap();
    client.split_keyspace(keyspace_id);
    (cluster, client)
}

fn prepare_region_ctx(client: &mut ClusterClient, key: &[u8]) -> (kvrpcpb::Context, TikvClient) {
    let region_id = client.get_region_id(key);
    let ctx = client.new_rpc_ctx(region_id, key).unwrap();
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());
    (ctx, kv_client)
}

fn shared_holder_start_tss(lock_info: &kvrpcpb::LockInfo) -> Vec<u64> {
    let mut start_tss = lock_info
        .get_shared_lock_infos()
        .iter()
        .map(|lock_info| lock_info.get_lock_version())
        .collect::<Vec<_>>();
    start_tss.sort_unstable();
    start_tss
}

fn must_shared_lock_holder_start_tss(
    client: &mut ClusterClient,
    key: &[u8],
    expected_start_tss: &[TimeStamp],
) {
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    let (_results, key_errors) = client
        .kv_pessimistic_lock(
            key.to_vec().into(),
            vec![key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            Some(1),
        )
        .unwrap();
    assert_eq!(
        key_errors.len(),
        1,
        "expected one shared-lock conflict on {:?}, got {:?}",
        key,
        key_errors
    );
    assert!(
        key_errors[0].has_locked(),
        "expected shared-lock conflict on {:?}, got {:?}",
        key,
        key_errors[0]
    );

    let locked = key_errors[0].get_locked();
    assert_eq!(locked.get_lock_type(), kvrpcpb::Op::SharedLock);

    let mut expected = expected_start_tss
        .iter()
        .map(|ts| ts.into_inner())
        .collect::<Vec<_>>();
    expected.sort_unstable();
    assert_eq!(shared_holder_start_tss(locked), expected);
}

fn acquire_shared_pessimistic_secondary(
    client: &mut ClusterClient,
    primary_key: &[u8],
    shared_key: &[u8],
    start_ts: TimeStamp,
    for_update_ts: TimeStamp,
) {
    client
        .kv_pessimistic_lock(
            primary_key.to_vec().into(),
            vec![primary_key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    client
        .kv_shared_pessimistic_lock(
            primary_key.to_vec().into(),
            vec![shared_key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
}

fn prewrite_shared_secondary(
    client: &mut ClusterClient,
    primary_key: &[u8],
    shared_key: &[u8],
    start_ts: TimeStamp,
    for_update_ts: TimeStamp,
) {
    acquire_shared_pessimistic_secondary(client, primary_key, shared_key, start_ts, for_update_ts);

    let (ctx, kv_client) = prepare_region_ctx(client, primary_key);
    let mut req = build_shared_pessimistic_prewrite_req(
        ctx,
        primary_key,
        shared_key,
        start_ts,
        for_update_ts,
    );
    req.set_secondaries(vec![shared_key.to_vec()].into());
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
}

fn build_shared_pessimistic_prewrite_req(
    ctx: kvrpcpb::Context,
    primary_key: &[u8],
    shared_key: &[u8],
    start_ts: TimeStamp,
    for_update_ts: TimeStamp,
) -> PrewriteRequest {
    let muts = TxnMutations::from_normal(vec![
        new_lock_mutation(primary_key.to_vec()),
        new_shared_lock_mutation(shared_key.to_vec()),
    ]);
    let mut req = PrewriteRequest::default();
    req.set_context(ctx);
    muts.set_prewrite_req(&mut req);
    req.set_primary_lock(primary_key.to_vec());
    req.set_start_version(start_ts.into_inner());
    req.set_lock_ttl(3000);
    req.set_txn_size(2);
    req.set_min_commit_ts(start_ts.into_inner() + 1);
    req.set_for_update_ts(for_update_ts.into_inner());
    req.pessimistic_actions =
        vec![PrewriteRequestPessimisticAction::DoPessimisticCheck; req.get_mutations().len()];
    req
}

#[test]
fn test_shared_lock_prewrite_rejects_async_commit_and_1pc_via_kv_rpc() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_single_node_cluster();

    let async_primary = i_to_key(200);
    let async_shared = i_to_key(201);
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    acquire_shared_pessimistic_secondary(
        &mut client,
        &async_primary,
        &async_shared,
        start_ts,
        for_update_ts,
    );

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &async_primary);
    let mut req = build_shared_pessimistic_prewrite_req(
        ctx,
        &async_primary,
        &async_shared,
        start_ts,
        for_update_ts,
    );
    req.use_async_commit = true;
    req.set_secondaries(vec![async_shared.clone()].into());
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert_eq!(resp.get_errors().len(), 1, "resp {:?}", resp);
    assert!(
        resp.get_errors()[0]
            .get_abort()
            .contains("shared lock prewrite cannot use async-commit"),
        "resp {:?}",
        resp
    );
    must_not_written(&mut client, &async_shared, start_ts);
    must_shared_lock_holder_start_tss(&mut client, &async_shared, &[start_ts]);

    let one_pc_primary = i_to_key(202);
    let one_pc_shared = i_to_key(203);
    let start_ts_2 = client.get_ts();
    let for_update_ts_2 = client.get_ts();
    acquire_shared_pessimistic_secondary(
        &mut client,
        &one_pc_primary,
        &one_pc_shared,
        start_ts_2,
        for_update_ts_2,
    );

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &one_pc_primary);
    let mut req = build_shared_pessimistic_prewrite_req(
        ctx,
        &one_pc_primary,
        &one_pc_shared,
        start_ts_2,
        for_update_ts_2,
    );
    req.try_one_pc = true;
    let resp = kv_client.kv_prewrite(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert_eq!(resp.get_errors().len(), 1, "resp {:?}", resp);
    assert!(
        resp.get_errors()[0]
            .get_abort()
            .contains("shared lock prewrite is incompatible with one-phase commit"),
        "resp {:?}",
        resp
    );
    must_not_written(&mut client, &one_pc_shared, start_ts_2);
    must_shared_lock_holder_start_tss(&mut client, &one_pc_shared, &[start_ts_2]);

    cluster.stop();
}

#[test]
fn test_cleanup_shared_lock_removes_only_matching_holder_via_kv_rpc() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_single_node_cluster();

    let shared_key = i_to_key(210);
    let primary_1 = i_to_key(211);
    let primary_2 = i_to_key(212);

    let start_ts_1 = client.get_ts();
    let for_update_ts_1 = client.get_ts();
    prewrite_shared_secondary(
        &mut client,
        &primary_1,
        &shared_key,
        start_ts_1,
        for_update_ts_1,
    );

    let start_ts_2 = client.get_ts();
    let for_update_ts_2 = client.get_ts();
    prewrite_shared_secondary(
        &mut client,
        &primary_2,
        &shared_key,
        start_ts_2,
        for_update_ts_2,
    );

    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_1, start_ts_2]);

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &shared_key);
    let mut req = CleanupRequest::default();
    req.set_context(ctx);
    req.set_key(shared_key.clone());
    req.set_start_version(start_ts_1.into_inner());
    req.set_current_ts(0);
    let resp = kv_client.kv_cleanup(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(!resp.has_error(), "resp {:?}", resp);

    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_2]);
    must_written_with_properties(
        &mut client,
        &shared_key,
        start_ts_1,
        start_ts_1,
        WriteType::Rollback,
    );

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &shared_key);
    let mut req = CleanupRequest::default();
    req.set_context(ctx);
    req.set_key(shared_key.clone());
    req.set_start_version(start_ts_2.into_inner());
    req.set_current_ts(0);
    let resp = kv_client.kv_cleanup(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(!resp.has_error(), "resp {:?}", resp);
    must_unlocked(&mut client, &shared_key);

    cluster.stop();
}

#[test]
fn test_pessimistic_rollback_shared_lock_removes_only_matching_holder_via_kv_rpc() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_single_node_cluster();

    let shared_key = i_to_key(220);
    let primary_1 = i_to_key(221);
    let primary_2 = i_to_key(222);

    let start_ts_1 = client.get_ts();
    let for_update_ts_1 = client.get_ts();
    acquire_shared_pessimistic_secondary(
        &mut client,
        &primary_1,
        &shared_key,
        start_ts_1,
        for_update_ts_1,
    );

    let start_ts_2 = client.get_ts();
    let for_update_ts_2 = client.get_ts();
    acquire_shared_pessimistic_secondary(
        &mut client,
        &primary_2,
        &shared_key,
        start_ts_2,
        for_update_ts_2,
    );

    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_1, start_ts_2]);

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &shared_key);
    let mut req = PessimisticRollbackRequest::default();
    req.set_context(ctx);
    req.set_keys(vec![shared_key.clone()].into());
    req.set_start_version(start_ts_1.into_inner());
    req.set_for_update_ts(for_update_ts_1.into_inner());
    let resp = kv_client.kv_pessimistic_rollback(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(resp.get_errors().is_empty(), "resp {:?}", resp);

    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_2]);
    must_not_written(&mut client, &shared_key, start_ts_1);

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &shared_key);
    let mut req = PessimisticRollbackRequest::default();
    req.set_context(ctx);
    req.set_keys(vec![shared_key.clone()].into());
    req.set_start_version(start_ts_2.into_inner());
    req.set_for_update_ts(for_update_ts_2.into_inner());
    let resp = kv_client.kv_pessimistic_rollback(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(resp.get_errors().is_empty(), "resp {:?}", resp);
    must_unlocked(&mut client, &shared_key);

    cluster.stop();
}

#[test]
fn test_resolve_lock_shared_commit_preserves_pending_holder() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_single_node_cluster();

    let shared_key = i_to_key(230);
    let primary_1 = i_to_key(231);
    let primary_2 = i_to_key(232);

    let start_ts_1 = client.get_ts();
    let for_update_ts_1 = client.get_ts();
    prewrite_shared_secondary(
        &mut client,
        &primary_1,
        &shared_key,
        start_ts_1,
        for_update_ts_1,
    );

    let start_ts_2 = client.get_ts();
    let for_update_ts_2 = client.get_ts();
    prewrite_shared_secondary(
        &mut client,
        &primary_2,
        &shared_key,
        start_ts_2,
        for_update_ts_2,
    );

    let commit_ts_1 = client.get_ts();
    client
        .kv_resolve_lock_batch(&shared_key, &[(start_ts_1, commit_ts_1)])
        .unwrap();

    must_written_with_properties(
        &mut client,
        &shared_key,
        start_ts_1,
        commit_ts_1,
        WriteType::Lock,
    );
    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_2]);

    client
        .kv_resolve_lock_batch(&shared_key, &[(start_ts_2, TimeStamp::zero())])
        .unwrap();

    must_unlocked(&mut client, &shared_key);
    must_written_with_properties(
        &mut client,
        &shared_key,
        start_ts_2,
        start_ts_2,
        WriteType::Rollback,
    );

    cluster.stop();
}

#[test]
fn test_txn_heart_beat_rejects_shared_lock_primary_invariant_via_kv_rpc() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_single_node_cluster();

    let shared_key = i_to_key(240);
    let primary_key = i_to_key(241);
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    prewrite_shared_secondary(
        &mut client,
        &primary_key,
        &shared_key,
        start_ts,
        for_update_ts,
    );

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &shared_key);
    let mut req = TxnHeartBeatRequest::default();
    req.set_context(ctx);
    req.set_primary_lock(shared_key.clone());
    req.set_start_version(start_ts.into_inner());
    req.set_advise_lock_ttl(5000);
    let resp = kv_client.kv_txn_heart_beat(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(resp.get_error().has_primary_mismatch(), "resp {:?}", resp);

    let lock_info = resp.get_error().get_primary_mismatch().get_lock_info();
    assert_eq!(lock_info.get_lock_type(), kvrpcpb::Op::SharedLock);
    assert_eq!(lock_info.get_key(), shared_key.as_slice());
    assert_eq!(
        shared_holder_start_tss(lock_info),
        vec![start_ts.into_inner()]
    );

    cluster.stop();
}

#[test]
fn test_check_secondary_locks_rejects_shared_lock_target_via_kv_rpc() {
    test_util::init_log_for_test();
    let (mut cluster, mut client) = new_single_node_cluster();

    let shared_key = i_to_key(250);
    let primary_key = i_to_key(251);
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    prewrite_shared_secondary(
        &mut client,
        &primary_key,
        &shared_key,
        start_ts,
        for_update_ts,
    );

    let (ctx, kv_client) = prepare_region_ctx(&mut client, &shared_key);
    let mut req = CheckSecondaryLocksRequest::default();
    req.set_context(ctx);
    req.set_keys(vec![shared_key.clone()].into());
    req.set_start_version(start_ts.into_inner());
    let resp = kv_client.kv_check_secondary_locks(&req).unwrap();
    assert!(!resp.has_region_error(), "resp {:?}", resp);
    assert!(
        resp.get_error()
            .get_abort()
            .contains("unexpected check-secondary-locks on shared lock"),
        "resp {:?}",
        resp
    );

    must_not_written(&mut client, &shared_key, start_ts);
    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts]);

    cluster.stop();
}