// Copyright 2023 PingCAP, Inc.

use std::{sync::mpsc, thread, time::Duration};

use bytes::Bytes;
use kvproto::{
    deadlock::WaitForEntry,
    kvrpcpb::{self, ApiVersion, Context, GetLockWaitInfoRequest, WriteConflictReason},
};
use test_cloud_server::{
    ServerCluster,
    client::{ClusterClient, PrewriteExt, TxnMutations},
    util::Mutation,
};
use tikv_util::{config::ReadableDuration, time::Instant};
use txn_types::{TimeStamp, WriteType};

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

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

    // Split to test for resolve locks on multiple regions.
    let split_keys = [5, 15, 33, 35];
    for i in split_keys {
        client.split(&i_to_key(i));
    }
    cluster.wait_pd_region_min_count(split_keys.len() + 1);

    // Test for kv_prewrite.
    {
        let mut mutations = vec![];
        for i in 0..10 {
            mutations.push(Mutation {
                key: Bytes::from(i_to_key(i)),
                value: Bytes::from(i_to_val(i)),
                op: kvrpcpb::Op::Put,
                ..Default::default()
            });
        }
        let start_ts = client.get_ts();
        // prewrite but no commit.
        let txn_muts = TxnMutations::from_normal(mutations);
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts, start_ts)
            .expect("kv_prewrite");

        // prewrite should meet locks of previous prewrite.
        client.put_kv(0..20, i_to_key, i_to_val);
        client.verify_data_with_ref_store();
    }

    // Test for kv_get.
    {
        let mut mutations = vec![];
        for i in 30..40 {
            mutations.push(Mutation {
                key: Bytes::from(i_to_key(i)),
                value: Bytes::from(i_to_val(i)),
                op: kvrpcpb::Op::Put,
                ..Default::default()
            });
        }
        let start_ts = client.get_ts();
        let put_time = Instant::now();
        let txn_muts = TxnMutations::from_normal(mutations);
        client
            .kv_prewrite(txn_muts.primary(), None, txn_muts.clone(), start_ts)
            .expect("kv_prewrite");

        let mut client1 = cluster.new_client();
        let handle = thread::spawn(move || {
            // `get_key` should wait until committed, and get the committed value.
            let (val, _) = client1.must_get_key(&i_to_key(35), put_time);
            assert_eq!(val, i_to_val(35));
        });

        thread::sleep(Duration::from_secs(1));
        let commit_ts = client.get_ts();
        client
            .kv_commit(txn_muts, start_ts, commit_ts, false)
            .expect("kv_commit");

        handle.join().unwrap();
        client.verify_data_with_ref_store();
    }

    cluster.stop();
}

fn prewrite_shared_holder(
    client: &mut ClusterClient,
    primary: &[u8],
    key: &[u8],
) -> (TimeStamp, TimeStamp) {
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    let (_, primary_errors) = client
        .kv_pessimistic_lock(
            primary.to_vec().into(),
            vec![primary.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    assert!(primary_errors.is_empty(), "{:?}", primary_errors);
    let (_, shared_errors) = client
        .kv_shared_pessimistic_lock(
            primary.to_vec().into(),
            vec![key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            None,
        )
        .unwrap();
    assert!(shared_errors.is_empty(), "{:?}", shared_errors);

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

fn get_lock_wait_info_entries(
    client: &ClusterClient,
    store_id: u64,
    keyspace_id: u32,
) -> Vec<WaitForEntry> {
    let kv_client = client.get_kv_client(store_id);
    let mut ctx = Context::default();
    ctx.set_api_version(ApiVersion::V2);
    ctx.set_keyspace_id(keyspace_id);
    let mut req = GetLockWaitInfoRequest::default();
    req.set_context(ctx);
    kv_client
        .get_lock_wait_info(&req)
        .unwrap()
        .take_entries()
        .into_vec()
}

fn wait_for_lock_wait_count(
    client: &ClusterClient,
    store_id: u64,
    keyspace_id: u32,
    key: &[u8],
    expected: usize,
) {
    for attempt in 0..50 {
        let entries = get_lock_wait_info_entries(client, store_id, keyspace_id);
        let actual = entries.iter().filter(|entry| entry.key == key).count();
        if actual == expected {
            return;
        }
        if attempt == 49 {
            panic!(
                "expected {} waiters for {:?}, got {:?}",
                expected, key, entries
            );
        }
        thread::sleep(Duration::from_millis(100));
    }
}

fn assert_pessimistic_retry(errors: &[kvrpcpb::KeyError]) {
    assert_eq!(errors.len(), 1, "{:?}", errors);
    assert!(errors[0].has_conflict(), "{:?}", errors);
    assert_eq!(
        errors[0].get_conflict().get_reason(),
        WriteConflictReason::PessimisticRetry
    );
}

#[test]
fn test_resolve_lock_batch_handles_upgrade_hint_before_final_release() {
    test_util::init_log_for_test();
    let mut cluster = ServerCluster::new(alloc_node_id_vec(1), |_, cfg| {
        cfg.pessimistic_txn.wake_up_delay_duration = ReadableDuration::secs(30);
    });
    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(40);
    let primary_u = i_to_key(41);
    let primary_h1 = i_to_key(42);
    let primary_h2 = i_to_key(43);
    let primary_x = i_to_key(44);

    // X has the oldest start_ts, but enters the wait queue after U. If the
    // upgrade hint is not consumed, the final release wakes X and leaves U for
    // the delayed ordinary-waiter fallback.
    let start_x = client.get_ts();
    let (start_u, for_update_u) = prewrite_shared_holder(&mut client, &primary_u, &key);
    let (start_h1, _) = prewrite_shared_holder(&mut client, &primary_h1, &key);
    let (start_h2, _) = prewrite_shared_holder(&mut client, &primary_h2, &key);

    // U first marks the wrapper shrink-only, then installs the pending upgrader.
    let (_, first_errors) = client
        .kv_pessimistic_lock(
            primary_u.clone().into(),
            vec![key.clone().into()],
            start_u,
            3000,
            for_update_u,
            None,
        )
        .unwrap();
    assert_pessimistic_retry(&first_errors);

    let mut upgrade_client = client.clone();
    let upgrade_primary = primary_u.clone();
    let upgrade_key = key.clone();
    let (upgrade_tx, upgrade_rx) = mpsc::sync_channel(1);
    let upgrade_task = thread::spawn(move || {
        let result = upgrade_client.kv_pessimistic_lock(
            upgrade_primary.into(),
            vec![upgrade_key.into()],
            start_u,
            3000,
            for_update_u,
            Some(30_000),
        );
        upgrade_tx.send(result).unwrap();
    });
    let store_id = client.get_stores()[0];
    wait_for_lock_wait_count(&client, store_id, keyspace_id, &key, 1);

    // Add an ordinary exclusive waiter behind the pending upgrader.
    let for_update_x = client.get_ts();
    let (_, primary_errors) = client
        .kv_pessimistic_lock(
            primary_x.clone().into(),
            vec![primary_x.clone().into()],
            start_x,
            3000,
            for_update_x,
            None,
        )
        .unwrap();
    assert!(primary_errors.is_empty(), "{:?}", primary_errors);
    let mut waiter_client = client.clone();
    let waiter_primary = primary_x.clone();
    let waiter_key = key.clone();
    let (waiter_tx, waiter_rx) = mpsc::sync_channel(1);
    let waiter_task = thread::spawn(move || {
        let result = waiter_client.kv_pessimistic_lock(
            waiter_primary.into(),
            vec![waiter_key.into()],
            start_x,
            3000,
            for_update_x,
            Some(30_000),
        );
        waiter_tx.send(result).unwrap();
    });
    wait_for_lock_wait_count(&client, store_id, keyspace_id, &key, 2);

    // Resolving all holders makes the same wrapper pass through 3 -> 2 -> 1 -> 0.
    // The intermediate sole holder is intentionally not part of this assertion.
    let commit_ts_h1 = client.get_ts();
    client
        .kv_resolve_lock_batch(
            &key,
            &[
                (start_u, TimeStamp::zero()),
                (start_h1, commit_ts_h1),
                (start_h2, TimeStamp::zero()),
            ],
        )
        .unwrap();

    let (_, upgrade_errors) = upgrade_rx
        .recv_timeout(Duration::from_secs(5))
        .expect("the upgrade hint should complete U before the delayed fallback")
        .unwrap();
    assert_eq!(upgrade_errors.len(), 1, "{:?}", upgrade_errors);
    if upgrade_errors[0].has_shared_lock_lost() {
        let lost = upgrade_errors[0].get_shared_lock_lost();
        assert_eq!(lost.get_start_ts(), start_u.into_inner());
        assert_eq!(lost.get_key(), key);
    } else {
        assert_pessimistic_retry(&upgrade_errors);
    }
    upgrade_task.join().unwrap();

    let (_, waiter_errors) = waiter_rx
        .recv_timeout(Duration::from_secs(5))
        .expect("the final release should retry the ordinary waiter")
        .unwrap();
    assert_pessimistic_retry(&waiter_errors);
    waiter_task.join().unwrap();
    wait_for_lock_wait_count(&client, store_id, keyspace_id, &key, 0);

    must_unlocked(&mut client, &key);
    must_written_with_properties(&mut client, &key, start_u, start_u, WriteType::Rollback);
    must_written_with_properties(&mut client, &key, start_h1, commit_ts_h1, WriteType::Lock);
    must_written_with_properties(&mut client, &key, start_h2, start_h2, WriteType::Rollback);

    // U's rollback status fences any later commit based on its old shared holder.
    let late_commit_ts = client.get_ts();
    let late_commit = client.kv_commit(
        TxnMutations::from_normal(vec![new_shared_lock_mutation(key.clone())]),
        start_u,
        late_commit_ts,
        false,
    );
    late_commit.expect_err("terminal U unexpectedly committed");

    let cleanup_errors = client
        .kv_pessimistic_rollback(vec![primary_x], start_x, for_update_x)
        .unwrap();
    assert!(cleanup_errors.is_empty(), "{:?}", cleanup_errors);

    cluster.stop();
}

#[test]
fn test_resolve_lock_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(50);
    let shared_key = i_to_key(51);
    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();

    let commit_ts = client.get_ts();
    client
        .kv_resolve_lock_batch(&primary_key, &[(start_ts, commit_ts)])
        .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_resolve_lock_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(52);
    let shared_key = i_to_key(53);

    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_resolve_lock_batch(&primary_key, &[(start_ts, commit_ts)])
        .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();
}

fn send_async_region_resolve_lock(
    client: &mut ClusterClient,
    representative_key: &[u8],
    start_ts: u64,
    commit_ts: Option<u64>,
) {
    let region = client.get_region_by_key(representative_key);
    let ctx = client.new_rpc_ctx(region.id, representative_key).unwrap();
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());

    let txn_info = kvrpcpb::TxnInfo {
        txn: start_ts,
        status: commit_ts.unwrap_or_default(),
        ..Default::default()
    };
    let mut req = kvrpcpb::ResolveLockRequest::default();
    req.set_context(ctx);
    req.set_is_async(true);
    req.set_txn_infos(vec![txn_info].into());

    let resp = kv_client.kv_resolve_lock(&req).unwrap();
    assert!(
        !resp.has_region_error() && !resp.has_error(),
        "resp {:?}",
        resp
    );
}

fn wait_until_no_locks_in_range(client: &mut ClusterClient, start_key: &[u8], end_key: &[u8]) {
    let deadline = std::time::Instant::now() + Duration::from_secs(10);
    loop {
        let max_ts = client.get_ts();
        let locks = scan_locks_in_range(client, start_key, end_key, max_ts.into_inner(), 100);
        if locks.is_empty() {
            return;
        }
        assert!(
            std::time::Instant::now() < deadline,
            "locks are not cleared before timeout: {:?}",
            locks
        );
        thread::sleep(Duration::from_millis(100));
    }
}

fn scan_locks_in_range(
    client: &mut ClusterClient,
    start_key: &[u8],
    end_key: &[u8],
    max_ts: u64,
    limit: u32,
) -> Vec<kvrpcpb::LockInfo> {
    let region = client.get_region_by_key(start_key);
    let ctx = client.new_rpc_ctx(region.id, start_key).unwrap();
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());

    let mut req = kvrpcpb::ScanLockRequest::default();
    req.set_context(ctx);
    req.set_start_key(start_key.to_vec());
    req.set_end_key(end_key.to_vec());
    req.set_max_version(max_ts);
    req.set_limit(limit);

    let mut resp = kv_client.kv_scan_lock(&req).unwrap();
    assert!(
        !resp.has_region_error() && !resp.has_error(),
        "resp {:?}",
        resp
    );
    resp.take_locks().into_vec()
}

#[test]
fn test_async_region_resolve_lock_merges_requests() {
    test_util::init_log_for_test();
    let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, _| {});
    cluster.wait_region_replicated(&[], 3);
    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_commit = i_to_key(1);
    let key_rollback = i_to_key(2);
    let key_pending = i_to_key(3);

    let start_ts_commit = client.get_ts();
    let commit_ts = client.get_ts();
    let commit_txn = TxnMutations::from_normal(vec![Mutation {
        key: Bytes::from(key_commit.clone()),
        value: Bytes::from(i_to_val(1)),
        op: kvrpcpb::Op::Put,
        assertion: kvrpcpb::Assertion::None,
    }]);
    client
        .kv_prewrite_with_retry(commit_txn.primary(), None, commit_txn, start_ts_commit)
        .expect("kv_prewrite");

    let start_ts_rollback = client.get_ts();
    let rollback_txn = TxnMutations::from_normal(vec![Mutation {
        key: Bytes::from(key_rollback.clone()),
        value: Bytes::from(i_to_val(2)),
        op: kvrpcpb::Op::Put,
        assertion: kvrpcpb::Assertion::None,
    }]);
    client
        .kv_prewrite_with_retry(
            rollback_txn.primary(),
            None,
            rollback_txn,
            start_ts_rollback,
        )
        .expect("kv_prewrite");

    let start_ts_pending = client.get_ts();
    let pending_txn = TxnMutations::from_normal(vec![Mutation {
        key: Bytes::from(key_pending.clone()),
        value: Bytes::from(i_to_val(3)),
        op: kvrpcpb::Op::Put,
        assertion: kvrpcpb::Assertion::None,
    }]);
    client
        .kv_prewrite_with_retry(pending_txn.primary(), None, pending_txn, start_ts_pending)
        .expect("kv_prewrite");

    let max_ts = client.get_ts();
    let locks = scan_locks_in_range(
        &mut client,
        &i_to_key(0),
        &i_to_key(10),
        max_ts.into_inner(),
        100,
    );
    assert_eq!(locks.len(), 3, "should find three locks before resolve");

    send_async_region_resolve_lock(
        &mut client,
        &key_commit,
        start_ts_commit.into_inner(),
        Some(commit_ts.into_inner()),
    );
    send_async_region_resolve_lock(
        &mut client,
        &key_rollback,
        start_ts_rollback.into_inner(),
        None,
    );

    // `scan_lock` uses a half-open key range [start_key, end_key), so this
    // waits for the commit and rollback transactions but excludes key_pending.
    wait_until_no_locks_in_range(&mut client, &i_to_key(0), &i_to_key(3));

    // Verify LOCK CF directly: only the transaction that was not included in
    // the async resolve requests should still have a lock.
    let max_ts = client.get_ts();
    let locks = scan_locks_in_range(
        &mut client,
        &i_to_key(0),
        &i_to_key(10),
        max_ts.into_inner(),
        100,
    );
    assert_eq!(locks.len(), 1, "only the unresolved txn should remain");
    assert_eq!(
        locks[0].get_key(),
        key_pending,
        "the remaining lock should be on the txn without resolve request"
    );
    assert_eq!(
        locks[0].get_lock_version(),
        start_ts_pending.into_inner(),
        "the remaining lock should belong to the txn without resolve request"
    );

    // Verify user-visible MVCC results as a supplement to the LOCK CF check
    // above. `get_key_version` observes committed writes, not remaining locks.
    assert_eq!(
        client
            .get_key_version(&key_commit, u64::MAX)
            .expect("get committed key should succeed"),
        Some(i_to_val(1))
    );
    assert_eq!(
        client
            .get_key_version(&key_rollback, u64::MAX)
            .expect("get rolled back key should succeed"),
        None
    );
    assert_eq!(
        client
            .get_key_version(&key_pending, u64::MAX)
            .expect("get unresolved key should succeed"),
        None
    );

    cluster.stop();
}