// Copyright 2025 PingCAP, Inc.

use kvproto::kvrpcpb;
use test_cloud_server::{
    ServerCluster,
    client::{Error, PrewriteExt, TxnMutations},
    util::Mutation,
};
use tikv_util::info;
use txn_types::{LockType, WriteType};

use super::helper::{
    must_locked_with_properties, must_not_written, must_unlocked, must_written_with_properties,
    new_delete_mutation, new_pessimistic_lock_mutation, new_put_mutation, new_shared_lock_mutation,
};
use crate::{alloc_node_id_vec, i_to_key, i_to_val};

#[test]
fn test_commit_ok() {
    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(b"xkey").unwrap();
    client.split_keyspace(keyspace_id);

    let k1 = i_to_key(1);
    let v1 = i_to_val(1);
    let k2 = i_to_key(2); // For Lock type
    let k3 = i_to_key(3); // For Delete type
    let primary_key = k1.clone();

    let start_ts = client.get_ts();

    // Prewrite all keys in one transaction
    let mutations = vec![
        Mutation {
            key: k1.to_vec().into(),
            value: v1.to_vec().into(),
            op: kvrpcpb::Op::Put,
            ..Default::default()
        },
        Mutation {
            key: k2.to_vec().into(),
            value: vec![].into(), // Lock has no value usually
            op: kvrpcpb::Op::Lock,
            ..Default::default()
        },
        Mutation {
            key: k3.to_vec().into(),
            value: vec![].into(),
            op: kvrpcpb::Op::Del,
            ..Default::default()
        },
    ];

    client
        .kv_prewrite(
            primary_key.to_vec().into(),
            None, // Let client figure out secondaries
            TxnMutations::from_normal(mutations.clone()),
            start_ts,
        )
        .expect("prewrite failed");

    // Check locks
    must_locked_with_properties(&mut client, &k1, &primary_key, start_ts, LockType::Put);
    must_locked_with_properties(&mut client, &k2, &primary_key, start_ts, LockType::Lock);
    must_locked_with_properties(&mut client, &k3, &primary_key, start_ts, LockType::Delete);

    let commit_ts = client.get_ts();

    // Commit all keys
    client
        .kv_commit(
            TxnMutations::from_normal(mutations.clone()),
            start_ts,
            commit_ts,
            false,
        )
        .expect("commit failed");

    // Check writes and that locks are gone
    must_unlocked(&mut client, &k1);
    must_unlocked(&mut client, &k2);
    must_unlocked(&mut client, &k3);
    must_written_with_properties(&mut client, &k1, start_ts, commit_ts, WriteType::Put);
    must_written_with_properties(&mut client, &k2, start_ts, commit_ts, WriteType::Lock);
    must_written_with_properties(&mut client, &k3, start_ts, commit_ts, WriteType::Delete);

    // Idempotency check: Commit again should succeed
    client
        .kv_commit(
            TxnMutations::from_normal(mutations.clone()),
            start_ts,
            commit_ts,
            false,
        )
        .expect("idempotent commit failed");

    // Re-Check writes to ensure commit didn't change anything
    must_written_with_properties(&mut client, &k1, start_ts, commit_ts, WriteType::Put);
    must_written_with_properties(&mut client, &k2, start_ts, commit_ts, WriteType::Lock);
    must_written_with_properties(&mut client, &k3, start_ts, commit_ts, WriteType::Delete);

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

#[test]
fn test_commit_mixed_put_primary_and_shared_secondary() {
    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 primary_key = i_to_key(101);
    let shared_key = i_to_key(102);
    let value = vec![b'v'; 512];

    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    client
        .kv_pessimistic_lock(
            primary_key.clone().into(),
            vec![primary_key.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    client
        .kv_shared_pessimistic_lock(
            primary_key.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_key.clone().into(),
            Some(&vec![shared_key.clone().into()]),
            TxnMutations::from_normal(vec![
                new_put_mutation(primary_key.clone(), value),
                new_shared_lock_mutation(shared_key.clone()),
            ]),
            &mut txn,
            PrewriteExt {
                for_update_ts,
                ..Default::default()
            },
        )
        .unwrap();

    must_locked_with_properties(
        &mut client,
        &primary_key,
        &primary_key,
        start_ts,
        LockType::Put,
    );

    let commit_ts = client.get_ts();
    client
        .kv_commit(
            TxnMutations::from_normal(vec![
                new_put_mutation(primary_key.clone(), vec![]),
                new_shared_lock_mutation(shared_key.clone()),
            ]),
            start_ts,
            commit_ts,
            false,
        )
        .unwrap();

    must_unlocked(&mut client, &primary_key);
    must_unlocked(&mut client, &shared_key);
    must_written_with_properties(
        &mut client,
        &primary_key,
        start_ts,
        commit_ts,
        WriteType::Put,
    );
    must_written_with_properties(
        &mut client,
        &shared_key,
        start_ts,
        commit_ts,
        WriteType::Lock,
    );

    cluster.stop();
}

#[test]
fn test_commit_mixed_delete_primary_and_shared_secondary() {
    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 primary_key = i_to_key(103);
    let shared_key = i_to_key(104);

    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    client
        .kv_pessimistic_lock(
            primary_key.clone().into(),
            vec![primary_key.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    client
        .kv_shared_pessimistic_lock(
            primary_key.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_key.clone().into(),
            Some(&vec![shared_key.clone().into()]),
            TxnMutations::from_normal(vec![
                new_delete_mutation(primary_key.clone()),
                new_shared_lock_mutation(shared_key.clone()),
            ]),
            &mut txn,
            PrewriteExt {
                for_update_ts,
                ..Default::default()
            },
        )
        .unwrap();

    let commit_ts = client.get_ts();
    client
        .kv_commit(
            TxnMutations::from_normal(vec![
                new_delete_mutation(primary_key.clone()),
                new_shared_lock_mutation(shared_key.clone()),
            ]),
            start_ts,
            commit_ts,
            false,
        )
        .unwrap();

    must_unlocked(&mut client, &primary_key);
    must_unlocked(&mut client, &shared_key);
    must_written_with_properties(
        &mut client,
        &primary_key,
        start_ts,
        commit_ts,
        WriteType::Delete,
    );
    must_written_with_properties(
        &mut client,
        &shared_key,
        start_ts,
        commit_ts,
        WriteType::Lock,
    );

    cluster.stop();
}

#[test]
fn test_commit_mixed_put_and_stale_pessimistic_lock() {
    test_util::init_log_for_test();
    let mut 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);

    let primary_key = i_to_key(111);
    let stale_lock_key = i_to_key(112);
    let value = i_to_val(111);

    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    client
        .kv_pessimistic_lock(
            primary_key.clone().into(),
            vec![primary_key.clone().into(), stale_lock_key.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();

    must_locked_with_properties(
        &mut client,
        &primary_key,
        &primary_key,
        start_ts,
        LockType::Pessimistic,
    );
    must_locked_with_properties(
        &mut client,
        &stale_lock_key,
        &primary_key,
        start_ts,
        LockType::Pessimistic,
    );

    let mut txn = client.begin_transaction(Some(start_ts));
    client
        .kv_prewrite_ext(
            primary_key.clone().into(),
            None,
            TxnMutations::from_normal(vec![new_put_mutation(primary_key.clone(), value)]),
            &mut txn,
            PrewriteExt::with_for_update_ts(for_update_ts),
        )
        .unwrap();

    must_locked_with_properties(
        &mut client,
        &primary_key,
        &primary_key,
        start_ts,
        LockType::Put,
    );
    must_locked_with_properties(
        &mut client,
        &stale_lock_key,
        &primary_key,
        start_ts,
        LockType::Pessimistic,
    );

    let commit_ts = client.get_ts();
    client
        .kv_commit(
            TxnMutations::from_normal(vec![
                new_put_mutation(primary_key.clone(), vec![]),
                new_pessimistic_lock_mutation(stale_lock_key.clone()),
            ]),
            start_ts,
            commit_ts,
            false,
        )
        .unwrap();

    must_unlocked(&mut client, &primary_key);
    must_unlocked(&mut client, &stale_lock_key);
    must_written_with_properties(
        &mut client,
        &primary_key,
        start_ts,
        commit_ts,
        WriteType::Put,
    );
    let stale_lock_info = super::helper::get_mvcc_info(&mut client, &stale_lock_key);
    assert!(
        stale_lock_info
            .as_ref()
            .is_none_or(|info| info.get_writes().is_empty()),
        "stale pessimistic commit should only remove the lock: {:?}",
        stale_lock_info
    );

    cluster.stop();
}

#[test]
fn test_commit_err() {
    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(b"xkey").unwrap();
    client.split_keyspace(keyspace_id);

    let k = i_to_key(1);
    let v = i_to_val(1);
    let primary_key = k.clone();

    let ts1 = client.get_ts();
    let ts2 = client.get_ts();
    info!("ts1: {}", ts1);
    info!("ts2: {}", ts2);

    // 1. Not prewrite yet
    let commit_muts_k = TxnMutations::from_normal(vec![Mutation {
        key: k.to_vec().into(),
        value: vec![].into(),
        op: kvrpcpb::Op::Put,
        ..Default::default()
    }]);
    let err = client
        .kv_commit(commit_muts_k.clone(), ts1, ts2, false)
        .unwrap_err();
    match err {
        Error::KeyError(key_error) => {
            assert!(
                key_error.has_txn_not_found()
                    || key_error.get_retryable().contains("TxnLockNotFound")
            );
        }
        _ => panic!("Expected KeyError(TxnLockNotFound), got {:?}", err),
    }
    must_unlocked(&mut client, &k);

    // 2. Prewrite, then commit with wrong start_ts
    let start_ts_ok = client.get_ts();
    let start_ts_wrong = client.get_ts();
    let commit_ts_ok = client.get_ts();
    let prewrite_muts = vec![Mutation {
        key: k.to_vec().into(),
        value: v.to_vec().into(),
        op: kvrpcpb::Op::Put,
        ..Default::default()
    }];
    client
        .kv_prewrite(
            primary_key.to_vec().into(),
            None,
            TxnMutations::from_normal(prewrite_muts.clone()),
            start_ts_ok,
        )
        .expect("prewrite failed");

    let _lock_info =
        must_locked_with_properties(&mut client, &k, &primary_key, start_ts_ok, LockType::Put);

    let err_wrong_ts = client
        .kv_commit(
            TxnMutations::from_normal(prewrite_muts.clone()),
            start_ts_wrong,
            commit_ts_ok,
            false,
        )
        .unwrap_err();
    // Should still be TxnLockNotFound because the lock's start_ts doesn't match
    // start_ts_wrong
    match err_wrong_ts {
        Error::KeyError(key_error) => {
            assert!(
                key_error.has_txn_not_found()
                    || key_error.get_retryable().contains("TxnLockNotFound")
            );
        }
        _ => panic!(
            "Expected KeyError(TxnLockNotFound) for wrong start_ts, got {:?}",
            err_wrong_ts
        ),
    }
    // Lock should still exist
    let _lock_info =
        must_locked_with_properties(&mut client, &k, &primary_key, start_ts_ok, LockType::Put);

    // 3. Rollback, then commit
    client
        .kv_rollback(
            TxnMutations::from_normal(prewrite_muts.clone()),
            start_ts_ok,
        )
        .expect("rollback failed");

    must_unlocked(&mut client, &k); // Lock should be gone
    must_written_with_properties(
        &mut client,
        &k,
        start_ts_ok,
        start_ts_ok,
        WriteType::Rollback,
    );

    let commit_ts_after_rollback = client.get_ts();
    let err_after_rollback = client
        .kv_commit(
            TxnMutations::from_normal(prewrite_muts.clone()),
            start_ts_ok,
            commit_ts_after_rollback,
            false,
        )
        .unwrap_err();
    // Commit after rollback should also result in TxnLockNotFound or a conflict
    // error indicating the rollback.
    match err_after_rollback {
        Error::KeyError(key_error) => {
            assert!(
                key_error.has_txn_not_found()
                    || key_error.get_retryable().contains("TxnLockNotFound")
            );
        }
        _ => panic!(
            "Expected KeyError(TxnLockNotFound) after rollback, got {:?}",
            err_after_rollback
        ),
    }

    // Verify final state (should be unlocked and no write record)
    must_unlocked(&mut client, &k);
    must_written_with_properties(
        &mut client,
        &k,
        start_ts_ok,
        start_ts_ok,
        WriteType::Rollback,
    );
    client.verify_data_with_ref_store();
    cluster.stop();
}

#[test]
fn test_commit_rolls_back_stale_pessimistic_lock() {
    let mut 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(b"xkey").unwrap();
    client.split_keyspace(keyspace_id);

    let k1 = i_to_key(11);
    let v1 = i_to_val(11);
    let k2 = i_to_key(12);
    let primary_key = k1.clone();

    let start_ts = client.get_ts();
    let for_update_ts = start_ts;
    let (_, key_errors) = client
        .kv_pessimistic_lock(
            primary_key.clone().into(),
            vec![k1.clone().into(), k2.clone().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .expect("pessimistic lock should succeed");
    assert!(
        key_errors.is_empty(),
        "unexpected key errors: {:?}",
        key_errors
    );

    must_locked_with_properties(
        &mut client,
        &k1,
        &primary_key,
        start_ts,
        LockType::Pessimistic,
    );
    must_locked_with_properties(
        &mut client,
        &k2,
        &primary_key,
        start_ts,
        LockType::Pessimistic,
    );

    let mut txn = client.begin_transaction(Some(start_ts));
    client
        .kv_prewrite_ext(
            primary_key.clone().into(),
            None,
            TxnMutations::from_normal(vec![new_put_mutation(k1.clone(), v1.clone())]),
            &mut txn,
            PrewriteExt {
                resolve_lock: false,
                for_update_ts,
                ..Default::default()
            },
        )
        .expect("prewrite should succeed");

    must_locked_with_properties(&mut client, &k1, &primary_key, start_ts, LockType::Put);
    must_locked_with_properties(
        &mut client,
        &k2,
        &primary_key,
        start_ts,
        LockType::Pessimistic,
    );

    let commit_ts = client.get_ts();
    client
        .kv_commit(
            TxnMutations::from_normal(vec![
                new_put_mutation(k1.clone(), v1.clone()),
                new_pessimistic_lock_mutation(k2.clone()),
            ]),
            start_ts,
            commit_ts,
            false,
        )
        .expect("commit should succeed");

    must_unlocked(&mut client, &k1);
    must_unlocked(&mut client, &k2);
    must_written_with_properties(&mut client, &k1, start_ts, commit_ts, WriteType::Put);
    must_not_written(&mut client, &k2, start_ts);
    must_not_written(&mut client, &k2, commit_ts);

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

// TODO: migrate test_min_commit_ts, which requires adding min_commit_ts to
// MvccInfo