// Copyright 2025 PingCAP, Inc.

use std::time::Duration;

use test_cloud_server::{
    ServerCluster,
    client::{PrewriteExt, TxnMutations},
};
use txn_types::LockType;

use super::helper::*;
use crate::{alloc_node_id_vec, i_to_key, i_to_val};

#[test]
fn test_verify_is_primary() {
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, _| {});
    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);

    let key = i_to_key(1);
    let primary = i_to_key(2);
    let value = i_to_val(3);
    let start_ts = client.get_ts();

    // Test with pessimistic lock
    let res = client
        .kv_pessimistic_lock(
            primary.clone().into(),
            vec![primary.clone().into(), key.clone().into()],
            start_ts,
            20000,
            start_ts,
            None,
        )
        .expect("pessimistic lock should succeed");
    println!("res: {:?}", res);
    // Verify the lock exists with correct properties
    // sleep 1s
    std::thread::sleep(Duration::from_secs(1));
    let _ =
        must_locked_with_properties(&mut client, &key, &primary, start_ts, LockType::Pessimistic);

    // Check txn status with verify_is_primary=true should fail with PrimaryMismatch
    let result = client.kv_check_txn_status(
        &key, start_ts, start_ts, start_ts, true, false, false, false,
    );
    assert!(result.unwrap().get_error().has_primary_mismatch());

    let mut txn = client.begin_transaction(Some(start_ts));
    // Convert pessimistic lock to prewrite
    client
        .kv_prewrite_ext(
            primary.to_vec().into(),
            Some(&vec![key.to_vec().into()]),
            TxnMutations::from_normal(vec![
                new_put_mutation(primary.to_vec(), value.to_vec()),
                new_put_mutation(key.to_vec(), value.to_vec()),
            ]),
            &mut txn,
            PrewriteExt {
                for_update_ts: start_ts,
                ..Default::default()
            },
        )
        .expect("pessimistic prewrite should succeed");

    // Verify the lock exists with correct properties after prewrite
    let _ = must_locked_with_properties(&mut client, &key, &primary, start_ts, LockType::Put);

    // Check txn status with verify_is_primary=true should still fail with
    // PrimaryMismatch
    let result = client.kv_check_txn_status(
        &key, start_ts, start_ts, start_ts, true, false, false, false,
    );
    assert!(result.unwrap().get_error().has_primary_mismatch());

    client.verify_data_with_ref_store();
    cluster.stop();
}

#[test]
fn test_check_txn_status_rejects_shared_secondary_lock() {
    test_util::init_log_for_test();
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, _| {});
    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);

    let shared_key = i_to_key(10);
    let primary = i_to_key(11);
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();

    client
        .kv_pessimistic_lock(
            primary.clone().into(),
            vec![primary.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    client
        .kv_shared_pessimistic_lock(
            primary.clone().into(),
            vec![shared_key.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();

    let mut txn = client.begin_transaction(Some(start_ts));
    client
        .kv_prewrite_ext(
            primary.clone().into(),
            Some(&vec![shared_key.clone().into()]),
            TxnMutations::from_normal(vec![
                new_lock_mutation(primary.clone()),
                new_shared_lock_mutation(shared_key.clone()),
            ]),
            &mut txn,
            PrewriteExt {
                for_update_ts,
                ..Default::default()
            },
        )
        .unwrap();

    let current_ts = client.get_ts();
    let result = client.kv_check_txn_status(
        &shared_key,
        start_ts,
        start_ts,
        current_ts,
        true,
        false,
        false,
        false,
    );
    assert!(result.unwrap().get_error().has_primary_mismatch());

    cluster.stop();
}