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

use std::time::{Duration, Instant};

use bytes::Bytes;
use kvproto::kvrpcpb::{self, PrewriteRequestPessimisticAction, WriteConflictReason};
use rand::{RngCore, SeedableRng, rngs::StdRng};
use rand_xorshift::XorShiftRng;
use test_cloud_server::{
    ServerCluster, alloc_node_id_vec,
    client::{ClusterClient, Error, PrewriteOptions, TxnMutations},
    util::Mutation,
};
use tikv::storage::config::ExtraCfScanMode;
use tikv_util::config::ReadableDuration;
use txn_types::TimeStamp;

use super::{
    helper::{new_lock_mutation, new_put_mutation},
    util::ts,
};
use crate::{i_to_key, i_to_val};

const EXTRA_CF_TXN_FUZZ_DURATION_ENV: &str = "CSE_EXTRA_CF_TXN_FUZZ_DURATION";
const EXTRA_CF_TXN_DEFAULT_ROUNDS: usize = 3;
const EXTRA_CF_TXN_REPORT_INTERVAL: Duration = Duration::from_secs(10);

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

    let duration_override = extra_cf_fuzz_duration_from_env(EXTRA_CF_TXN_FUZZ_DURATION_ENV);
    let nodes = alloc_node_id_vec(1);
    let root_seed = rand::thread_rng().next_u64();
    println!(
        "extra cf txn differential fuzz start: root_seed={root_seed:#x} nodes={nodes:?} duration_override={duration_override:?} default_rounds={}",
        EXTRA_CF_TXN_DEFAULT_ROUNDS,
    );

    let start = Instant::now();
    let mut next_report = EXTRA_CF_TXN_REPORT_INTERVAL;
    let mut completed_rounds = 0;
    let mut round_seed_rng = StdRng::seed_from_u64(root_seed);
    while should_continue_extra_cf_fuzz(
        duration_override,
        start,
        completed_rounds,
        EXTRA_CF_TXN_DEFAULT_ROUNDS,
    ) {
        let seed = round_seed_rng.next_u64();
        run_seeded_scenarios(&nodes, seed);
        completed_rounds += 1;

        let elapsed = start.elapsed();
        while duration_override.is_some() && elapsed >= next_report {
            println!(
                "extra cf txn differential fuzz progress: root_seed={root_seed:#x} elapsed={elapsed:?} rounds={completed_rounds} scenarios={}",
                completed_rounds * 7,
            );
            next_report += EXTRA_CF_TXN_REPORT_INTERVAL;
        }
    }

    println!(
        "extra cf txn differential fuzz done: root_seed={root_seed:#x} elapsed={:?} rounds={completed_rounds} scenarios={}",
        start.elapsed(),
        completed_rounds * 7,
    );
}

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

    // Keep both stores alive so constructing the pruned store would overwrite
    // the legacy store's mode if the setting were still process-global.
    let (mut legacy_cluster, mut legacy_client) =
        new_cluster_with_extra_cf_scan_mode(ExtraCfScanMode::Legacy);
    let (mut pruned_cluster, mut pruned_client) =
        new_cluster_with_extra_cf_scan_mode(ExtraCfScanMode::Pruned);

    let legacy_key = i_to_key(9_001);
    let pruned_key = i_to_key(9_002);
    prepare_rollback_and_newer_lock(&mut legacy_client, &legacy_key);
    prepare_rollback_and_newer_lock(&mut pruned_client, &pruned_key);

    // Legacy preserves ordered conflict selection and sees the newer LOCK
    // before the exact rollback record.
    let legacy_error = acquire_pessimistic_lock_error(&mut legacy_client, &legacy_key);
    assert!(legacy_error.has_conflict(), "{:?}", legacy_error);
    let conflict = legacy_error.get_conflict();
    assert_eq!(conflict.get_reason(), WriteConflictReason::PessimisticRetry);
    assert_eq!(conflict.get_conflict_ts(), ts(3_000).into_inner());
    assert_eq!(conflict.get_conflict_commit_ts(), ts(4_000).into_inner());

    // Pruned mode uses point-GET-first selection beyond the one-second gap and
    // therefore reports the exact rollback record.
    let pruned_error = acquire_pessimistic_lock_error(&mut pruned_client, &pruned_key);
    assert!(
        pruned_error
            .get_abort()
            .contains("PessimisticLockRolledBack"),
        "{:?}",
        pruned_error
    );

    legacy_cluster.stop();
    pruned_cluster.stop();
}

fn new_cluster_with_extra_cf_scan_mode(mode: ExtraCfScanMode) -> (ServerCluster, ClusterClient) {
    let nodes = alloc_node_id_vec(1);
    let cluster = ServerCluster::new(nodes, |_, cfg| {
        cfg.storage.check_backup_ts = false;
        cfg.storage.extra_cf_scan_mode = mode.as_str().to_owned();
    });
    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_rollback_and_newer_lock(client: &mut ClusterClient, key: &[u8]) {
    rollback_lock_record(client, key, ts(10));
    commit_lock_record(client, key, ts(3_000), ts(4_000));
}

fn acquire_pessimistic_lock_error(client: &mut ClusterClient, key: &[u8]) -> kvrpcpb::KeyError {
    let (_, mut errors) = client
        .kv_pessimistic_lock(
            Bytes::copy_from_slice(key),
            vec![Bytes::copy_from_slice(key)],
            ts(10),
            20_000,
            ts(2_000),
            None,
        )
        .unwrap();
    assert_eq!(errors.len(), 1, "unexpected key errors: {errors:?}");
    errors.remove(0)
}

fn run_seeded_scenarios(nodes: &[u16], seed: u64) {
    let mut cluster = ServerCluster::new(nodes.to_vec(), |_, cfg| {
        cfg.storage.check_backup_ts = false;
        cfg.storage.extra_cf_scan_mode = ExtraCfScanMode::Both.as_str().to_owned();
    });
    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 mut rng = rng_from_seed(seed);
    run_optimistic_extra_cf_conflict(&mut client, &mut rng);
    run_optimistic_success(&mut client, &mut rng);
    run_self_rollback_conflict(&mut client, &mut rng);
    run_acquire_pessimistic_lock_conflict(&mut client, &mut rng);
    run_acquire_pessimistic_lock_success(&mut client, &mut rng);
    run_amend_pessimistic_lock_conflict(&mut client, &mut rng);
    run_mixed_multi_key_conflict(&mut client, &mut rng);
    client.verify_data_with_ref_store();
    cluster.stop();
}

fn should_continue_extra_cf_fuzz(
    duration_override: Option<Duration>,
    start: Instant,
    completed_rounds: usize,
    default_rounds: usize,
) -> bool {
    match duration_override {
        Some(duration) => completed_rounds == 0 || start.elapsed() < duration,
        None => completed_rounds < default_rounds,
    }
}

fn extra_cf_fuzz_duration_from_env(env_var: &str) -> Option<Duration> {
    match std::env::var(env_var) {
        Ok(value) => Some(Duration::from(
            value
                .parse::<ReadableDuration>()
                .unwrap_or_else(|err| panic!("invalid {}={value:?}: {}", env_var, err)),
        )),
        Err(std::env::VarError::NotPresent) => None,
        Err(err) => panic!("failed to read {}: {}", env_var, err),
    }
}

fn run_optimistic_extra_cf_conflict(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key = seeded_key(rng, "opt_conflict", 0);
    commit_lock_record(client, &key, ts(10), ts(15));

    let mutations = single_mutation_txn_muts(new_put_mutation(key.clone(), i_to_val(1)));
    let error = client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.clone()),
            None,
            mutations,
            ts(12),
            PrewriteOptions::default(),
        )
        .expect_err("optimistic prewrite should conflict with the newer LOCK record");
    assert_write_conflict(
        error,
        WriteConflictReason::Optimistic,
        ts(12),
        ts(10),
        ts(15),
    );
}

fn run_optimistic_success(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key = seeded_key(rng, "opt_success", 0);
    commit_lock_record(client, &key, ts(10), ts(15));

    let mutations = single_mutation_txn_muts(new_put_mutation(key.clone(), i_to_val(2)));
    client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.clone()),
            None,
            mutations.clone(),
            ts(20),
            PrewriteOptions::default(),
        )
        .expect("optimistic prewrite should ignore the older LOCK record");
    client
        .kv_commit(mutations, ts(20), ts(25), false)
        .expect("optimistic transaction should commit");
}

fn run_self_rollback_conflict(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key = seeded_key(rng, "self_rollback", 0);
    rollback_lock_record(client, &key, ts(10));

    let mutations = single_mutation_txn_muts(new_put_mutation(key.clone(), i_to_val(3)));
    let error = client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.clone()),
            None,
            mutations,
            ts(10),
            PrewriteOptions::default(),
        )
        .expect_err("prewrite should reject a transaction that was already rolled back");
    assert_write_conflict(
        error,
        WriteConflictReason::SelfRolledBack,
        ts(10),
        ts(10),
        ts(10),
    );
}

fn run_acquire_pessimistic_lock_conflict(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key = seeded_key(rng, "pes_lock_conflict", 0);
    commit_lock_record(client, &key, ts(10), ts(15));

    let (_, mut key_errors) = client
        .kv_pessimistic_lock(
            Bytes::from(key.clone()),
            vec![Bytes::from(key.clone())],
            ts(12),
            20_000,
            ts(12),
            None,
        )
        .expect("pessimistic-lock RPC should return a key-level conflict");
    assert_eq!(key_errors.len(), 1, "unexpected key errors: {key_errors:?}");
    let mut key_error = key_errors.remove(0);
    assert!(
        key_error.has_conflict(),
        "unexpected key error: {:?}",
        key_error
    );
    let conflict = key_error.take_conflict();
    assert_write_conflict_fields(
        &conflict,
        WriteConflictReason::PessimisticRetry,
        ts(12),
        ts(10),
        ts(15),
    );
}

fn run_acquire_pessimistic_lock_success(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key = seeded_key(rng, "pes_lock_success", 0);
    commit_lock_record(client, &key, ts(10), ts(15));

    let (_, key_errors) = client
        .kv_pessimistic_lock(
            Bytes::from(key.clone()),
            vec![Bytes::from(key.clone())],
            ts(12),
            20_000,
            ts(15),
            None,
        )
        .expect("pessimistic-lock RPC should succeed at the conflict threshold");
    assert!(
        key_errors.is_empty(),
        "unexpected key errors: {:?}",
        key_errors
    );
}

fn run_amend_pessimistic_lock_conflict(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key = seeded_key(rng, "amend_conflict", 0);
    commit_put_record(client, &key, ts(20), ts(25));

    let mutations = single_mutation_txn_muts(new_put_mutation(key.clone(), i_to_val(4)));
    let error = client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.clone()),
            None,
            mutations,
            ts(10),
            PrewriteOptions {
                for_update_ts: ts(22),
                pessimistic_action: PrewriteRequestPessimisticAction::DoConstraintCheck,
                ..Default::default()
            },
        )
        .expect_err("amending the lost pessimistic lock should fail");
    assert_key_error_abort(error, "PessimisticLockNotFound");
}

fn run_mixed_multi_key_conflict(client: &mut ClusterClient, rng: &mut XorShiftRng) {
    let key_a = seeded_key(rng, "mixed_batch", 0);
    let key_b = seeded_key(rng, "mixed_batch", 1);
    commit_lock_record(client, &key_b, ts(30), ts(35));

    let mutations = TxnMutations::from_normal(vec![
        new_put_mutation(key_a.clone(), i_to_val(5)),
        new_put_mutation(key_b.clone(), i_to_val(6)),
    ]);
    let error = client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key_a.clone()),
            None,
            mutations,
            ts(32),
            PrewriteOptions::default(),
        )
        .expect_err("multi-key prewrite should report the conflicting key");
    assert_write_conflict(
        error,
        WriteConflictReason::Optimistic,
        ts(32),
        ts(30),
        ts(35),
    );
}

fn commit_lock_record(
    client: &mut ClusterClient,
    key: &[u8],
    start_ts: TimeStamp,
    commit_ts: TimeStamp,
) {
    let mutations = single_mutation_txn_muts(new_lock_mutation(key.to_vec()));
    client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.to_vec()),
            None,
            mutations.clone(),
            start_ts,
            PrewriteOptions::default(),
        )
        .unwrap();
    client
        .kv_commit(mutations, start_ts, commit_ts, false)
        .unwrap();
}

fn rollback_lock_record(client: &mut ClusterClient, key: &[u8], start_ts: TimeStamp) {
    let mutations = single_mutation_txn_muts(new_lock_mutation(key.to_vec()));
    client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.to_vec()),
            None,
            mutations.clone(),
            start_ts,
            PrewriteOptions::default(),
        )
        .unwrap();
    client.kv_rollback(mutations, start_ts).unwrap();
}

fn commit_put_record(
    client: &mut ClusterClient,
    key: &[u8],
    start_ts: TimeStamp,
    commit_ts: TimeStamp,
) {
    let mutations = single_mutation_txn_muts(new_put_mutation(key.to_vec(), i_to_val(7)));
    client
        .kv_prewrite_with_retry_opt_with_options(
            Bytes::from(key.to_vec()),
            None,
            mutations.clone(),
            start_ts,
            PrewriteOptions::default(),
        )
        .unwrap();
    client
        .kv_commit(mutations, start_ts, commit_ts, false)
        .unwrap();
}

fn single_mutation_txn_muts(mutation: Mutation) -> TxnMutations {
    TxnMutations::from_normal(vec![mutation])
}

fn assert_write_conflict(
    error: Error,
    expected_reason: WriteConflictReason,
    expected_start_ts: TimeStamp,
    expected_conflict_ts: TimeStamp,
    expected_conflict_commit_ts: TimeStamp,
) {
    let conflict = match error {
        Error::WriteConflict(conflict) => conflict,
        error => {
            let mut key_error = take_single_key_error(error);
            assert!(
                key_error.has_conflict(),
                "expected write conflict, got key error: {:?}",
                key_error
            );
            key_error.take_conflict()
        }
    };
    assert_write_conflict_fields(
        &conflict,
        expected_reason,
        expected_start_ts,
        expected_conflict_ts,
        expected_conflict_commit_ts,
    );
}

fn assert_write_conflict_fields(
    conflict: &kvrpcpb::WriteConflict,
    expected_reason: WriteConflictReason,
    expected_start_ts: TimeStamp,
    expected_conflict_ts: TimeStamp,
    expected_conflict_commit_ts: TimeStamp,
) {
    assert_eq!(conflict.get_reason(), expected_reason, "{conflict:?}");
    assert_eq!(
        conflict.get_start_ts(),
        expected_start_ts.into_inner(),
        "{conflict:?}"
    );
    assert_eq!(
        conflict.get_conflict_ts(),
        expected_conflict_ts.into_inner(),
        "{conflict:?}"
    );
    assert_eq!(
        conflict.get_conflict_commit_ts(),
        expected_conflict_commit_ts.into_inner(),
        "{conflict:?}"
    );
}

fn assert_key_error_abort(error: Error, expected_abort: &str) {
    let key_error = take_single_key_error(error);
    assert!(
        key_error.get_abort().contains(expected_abort),
        "expected abort containing {:?}, got {:?}",
        expected_abort,
        key_error
    );
}

fn take_single_key_error(error: Error) -> kvrpcpb::KeyError {
    match error {
        Error::KeyError(key_error) => key_error,
        Error::KeyErrors(mut key_errors) => {
            assert_eq!(key_errors.len(), 1, "unexpected key errors: {key_errors:?}");
            key_errors.remove(0)
        }
        error => panic!("expected key error, got {:?}", error),
    }
}

fn rng_from_seed(seed: u64) -> XorShiftRng {
    let mut seed_bytes = [0u8; 16];
    seed_bytes[..8].copy_from_slice(&seed.to_le_bytes());
    seed_bytes[8..].copy_from_slice(&(seed ^ 0x9e37_79b9_7f4a_7c15).to_le_bytes());
    XorShiftRng::from_seed(seed_bytes)
}

fn seeded_key(rng: &mut XorShiftRng, label: &str, idx: usize) -> Vec<u8> {
    let mut key = i_to_key(0)[..4].to_vec();
    key.extend(format!("ecfd_{}_{}_{}", label, idx, rng.next_u32()).as_bytes());
    key
}