// Copyright 2026 PingCAP, Inc.

use std::thread;

use api_version::ApiV2;
use bytes::Bytes;
use futures::executor::block_on;
use kvproto::kvrpcpb::{self, ResolveLockRequest, ResolveLockResponse};
use rfstore::store::test_util::assert_stale_prepare_txn_file_result_is_ignored_after_registration;
use security::SecurityConfig;
use test_cloud_server::{
    ServerClusterBuilder, TikvWorkerOptions, alloc_node_id_vec,
    client::{ClusterClient, ClusterClientOptions, RequestOptions, TxnMutations, TxnWriteMethod},
    oss::prepare_dfs,
    util::Mutation,
};
use test_pd_client::PdWrapper;
use tikv_util::time::Instant;
use txn_types::TimeStamp;

use crate::cases::cfg_notify_once;

const KEYSPACE_ID: u32 = 1;

#[test]
fn test_stale_prepare_txn_file_result_is_ignored_after_registration() {
    test_util::init_log_for_test();
    assert_stale_prepare_txn_file_result_is_ignored_after_registration();
}

#[test]
fn test_txn_file_resolve_lock_retries_after_stale_missing_snapshot() {
    test_util::init_log_for_test();

    let before_schedule_fp = "txn::before_schedule_txn_file_cmd";

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("txn_file_stale_resolve");
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(1), |_, conf| {
        conf.dfs = dfs_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 1);
    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(1024),
        ..Default::default()
    });
    client.split_keyspace(KEYSPACE_ID);

    let key = txn_file_key(0);
    let value = b"txn-file-stale-resolve-lock".to_vec();
    let start_ts = client.get_ts();
    let commit_ts = client.get_ts();

    let (rx, _before_schedule_guard) = cfg_notify_once(before_schedule_fp);
    let mut stale_resolve_client = cluster.new_client();
    let resolve_key = key.clone();
    let h = thread::spawn(move || {
        send_txn_file_resolve_lock(&mut stale_resolve_client, &resolve_key, start_ts, commit_ts)
    });

    let tx = rx.blocking_recv().unwrap();

    prewrite_txn_file_put(&mut client, key.clone(), value.clone(), start_ts);
    assert_ne!(
        get_key_without_resolving_locks(&mut client, &key),
        Some(value.clone()),
        "prewritten txn-file lock should block reads before resolve"
    );

    tx.send(()).unwrap();
    let resp = h.join().unwrap();
    assert!(
        resp.has_region_error() && resp.get_region_error().has_stale_command(),
        "resp {:?}",
        resp
    );
    assert_ne!(
        get_key_without_resolving_locks(&mut client, &key),
        Some(value.clone()),
        "stale resolve lock must not commit the txn-file lock"
    );

    let resp = send_txn_file_resolve_lock(&mut client, &key, start_ts, commit_ts);
    assert!(
        !resp.has_region_error() && !resp.has_error(),
        "resp {:?}",
        resp
    );
    assert_eq!(client.get_key_version(&key, u64::MAX).unwrap(), Some(value));

    fail::remove(before_schedule_fp);
    cluster.stop();
    oss.shutdown();
}

fn prewrite_txn_file_put(
    client: &mut ClusterClient,
    key: Vec<u8>,
    value: Vec<u8>,
    start_ts: TimeStamp,
) {
    let txn_muts = block_on(TxnMutations::build(
        vec![Mutation {
            key: Bytes::from(key.clone()),
            value: Bytes::from(value),
            op: kvrpcpb::Op::Put,
            ..Default::default()
        }],
        vec![],
        TxnWriteMethod::FileBased,
        client.txn_file_helper(),
    ))
    .unwrap();
    client
        .kv_prewrite(txn_muts.primary(), None, txn_muts, start_ts)
        .unwrap();
}

fn send_txn_file_resolve_lock(
    client: &mut ClusterClient,
    key: &[u8],
    start_ts: TimeStamp,
    commit_ts: TimeStamp,
) -> ResolveLockResponse {
    let region = client.get_region_by_key(key);
    let ctx = client.new_rpc_ctx(region.id, key).unwrap();
    let kv_client = client.get_kv_client(ctx.get_peer().get_store_id());

    let mut txn_info = kvrpcpb::TxnInfo::default();
    txn_info.set_txn(start_ts.into_inner());
    txn_info.set_status(commit_ts.into_inner());
    txn_info.set_is_txn_file(true);

    let mut req = ResolveLockRequest::new();
    req.set_context(ctx);
    req.set_txn_infos(vec![txn_info].into());
    kv_client.kv_resolve_lock(&req).unwrap()
}

fn get_key_without_resolving_locks(client: &mut ClusterClient, key: &[u8]) -> Option<Vec<u8>> {
    client
        .get_key_version_opt(
            key,
            u64::MAX,
            Instant::now(),
            &RequestOptions::no_resolve_lock(),
        )
        .unwrap()
        .0
}

fn txn_file_key(i: usize) -> Vec<u8> {
    let mut key = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID);
    key.extend(format!("txn_file_key_{i:08}").into_bytes());
    key
}