// Copyright 2025 PingCAP, Inc.

use std::{convert::TryInto, time::Duration};

use bytes::Bytes;
use futures::{executor::block_on, future::try_join_all};
use kvengine::ia::util::IaConfig;
use kvproto::kvrpcpb;
use matches::assert_matches;
use rstest::rstest;
use schema::schema::{StorageClass, StorageClassSpec};
use test_cloud_server::{
    MajorCompactionTarget, ServerClusterBuilder, ServerClusterExt, TryWaiter, alloc_node_id_vec,
    client::{
        ClusterClientOptions, CommitAction, Error as ClientError, MutateOptions, TxnMutations,
        TxnWriteMethod,
    },
    keyspace::{CreateKeyspaceOptions, make_row_key},
    try_wait_result_async,
    util::Mutation,
};
use tikv_client::{Transaction, TransactionOptions};
use tikv_util::{config::ReadableSize, info, time::Instant};

use crate::random_value;

const DATA_COUNT: usize = 120;
const TXN_FILE_MAX_CHUNK_SIZE: usize = 1024;

#[rstest]
#[case::local(false, TxnWriteMethod::Normal)]
#[case::ia(true, TxnWriteMethod::Normal)]
#[case::txn_file(true, TxnWriteMethod::FileBased)]
fn test_get_put(#[case] enable_ia: bool, #[case] write_method: TxnWriteMethod) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let (cluster, keyspace_id, table_id) = prepare_cluster(enable_ia);
    let gen_key = i_to_key(keyspace_id, table_id);

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(TXN_FILE_MAX_CHUNK_SIZE),
        ..Default::default()
    });

    let key = Bytes::from(gen_key(40));
    let val = Bytes::from("v40");

    let ts0 = client.get_ts().into_inner();
    let v0 = client.get_key_version(&key, ts0).unwrap().unwrap();

    // Prewrite:
    let start_ts = client.get_ts();
    let mutation = make_put(key.clone(), val.clone());
    let txn_muts = block_on(TxnMutations::build(
        vec![mutation.clone()],
        vec![],
        write_method,
        client.txn_file_helper(),
    ))
    .unwrap();
    client
        .kv_prewrite(key.clone(), None, txn_muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();

    // "get" advance the min_commit_ts:
    let ts2 = client.get_ts().into_inner();
    let v = client.get_key_version(&key, ts2).unwrap();
    assert_eq!(v.as_ref(), Some(&v0));

    // Commit:
    let new_commit_ts = client
        .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
        .unwrap();
    assert!(new_commit_ts > commit_ts);

    // Verify:
    let v = client.get_key_version(&key, ts2).unwrap();
    assert_eq!(v.as_ref(), Some(&v0));
    let v = client
        .get_key_version(&key, new_commit_ts.into_inner())
        .unwrap();
    assert_eq!(v.as_deref(), Some(val.as_ref()));

    client.put_kv_in_ref_store(vec![mutation]);
    client.verify_data_with_ref_store();

    // Exercise the TiKV client scan and batch-get paths on the same fixture.
    // These used to be separate cases whose dominant cost was rebuilding this
    // identical cluster and compacted base data.
    let ref_store = client.ref_store();
    let ref_store = ref_store.lock().unwrap();
    let keys = vec![gen_key(0), gen_key(40), gen_key(80)];
    let values = keys
        .iter()
        .map(|key| ref_store.get(key.as_slice()).cloned().unwrap().unwrap())
        .collect::<Vec<_>>();
    let task = async {
        let mut txn_client = cluster.new_txn_client().await;
        let count = txn_client
            .verify_data_by_scan(&ref_store, None)
            .await
            .unwrap();
        assert_eq!(count, DATA_COUNT);

        let start_ts = txn_client.current_timestamp().await.unwrap();
        let mut snapshot = txn_client.snapshot(start_ts, TransactionOptions::default().read_only());
        let kvs = snapshot.batch_get(keys.clone()).await.unwrap();
        let mut count = 0;
        for (index, kv) in kvs.enumerate() {
            assert_eq!(kv.key(), keys[index].as_ref());
            assert_eq!(kv.value(), &values[index]);
            count += 1;
        }
        assert_eq!(count, keys.len());
    };
    rt.block_on(task);
    drop(ref_store);

    // Exercise pessimistic get-for-update and commit on the same fixture.
    let key = gen_key(41);
    let val = b"v41".to_vec();
    let ts0 = client.get_ts().into_inner();
    let v0 = client.get_key_version(&key, ts0).unwrap().unwrap();
    rt.block_on(async {
        let txn_client = cluster.new_txn_client().await;
        let mut tx = txn_client.begin_pessimistic().await.unwrap();
        let v = tx.get_for_update(key.clone()).await.unwrap().unwrap();
        assert_eq!(v, v0);
        let ts2 = client.get_ts().into_inner();
        let v = client.get_key_version(&key, ts2).unwrap();
        assert_eq!(v.as_ref(), Some(&v0));
        tx.put(key.clone(), val.clone()).await.unwrap();
        tx.commit().await.unwrap();
    });
    client.put_kv_in_ref_store(vec![make_put(key.into(), val.into())]);

    // Exercise optimistic write-conflict detection without rebuilding the
    // fixture and its compacted base data.
    let key = Bytes::from(gen_key(42));
    let val = Bytes::from("v42");
    let ts0 = client.get_ts();
    let mutation = make_put(key.clone(), val.clone());
    let txn_muts = block_on(TxnMutations::build(
        vec![mutation.clone()],
        vec![],
        write_method,
        client.txn_file_helper(),
    ))
    .unwrap();
    let start_ts = client.get_ts();
    client
        .kv_prewrite(key.clone(), None, txn_muts.clone(), start_ts)
        .unwrap();
    let commit_ts = client.get_ts();
    client
        .kv_commit(txn_muts.clone(), start_ts, commit_ts, false)
        .unwrap();

    // Another write:
    let err = client
        .kv_prewrite(key.clone(), None, txn_muts.clone(), ts0)
        .unwrap_err();
    assert_matches!(err, ClientError::WriteConflict { .. });
    client.put_kv_in_ref_store(vec![mutation]);

    // Cover insert's existing-key error and successful new-key path here too.
    {
        let key = Bytes::from(gen_key(43));
        let val = Bytes::from("v43");
        let mutation = make_insert(key.clone(), val.clone());
        let txn_muts = block_on(TxnMutations::build(
            vec![mutation.clone()],
            vec![],
            write_method,
            client.txn_file_helper(),
        ))
        .unwrap();
        let err = client
            .kv_prewrite(key.clone(), None, txn_muts, client.get_ts())
            .unwrap_err();
        assert_matches!(err, ClientError::AlreadyExist { .. });
    }

    {
        let key = Bytes::from(gen_key(1000));
        let val = Bytes::from("v1000");
        let mutation = make_insert(key.clone(), val.clone());
        let txn_muts = block_on(TxnMutations::build(
            vec![mutation.clone()],
            vec![],
            write_method,
            client.txn_file_helper(),
        ))
        .unwrap();
        let start_ts = client.get_ts();
        client
            .kv_prewrite(key.clone(), None, txn_muts.clone(), start_ts)
            .unwrap();
        client
            .kv_commit(txn_muts.clone(), start_ts, client.get_ts(), false)
            .unwrap();

        client.put_kv_in_ref_store(vec![mutation]);
    }

    client.verify_data_with_ref_store();
}

#[rstest]
#[case::local(false)]
#[case::ia(true)]
fn test_write_conflict_pessimistic(#[case] enable_ia: bool) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let (cluster, keyspace_id, table_id) = prepare_cluster(enable_ia);
    let gen_key = i_to_key(keyspace_id, table_id);

    let key = gen_key(1000);
    const CONCURRENCY: usize = 8;

    let mut client = cluster.new_client();
    let val = CONCURRENCY.to_be_bytes().to_vec();
    client.put_kv_in_ref_store(vec![make_put(key.clone().into(), val.into())]);
    let ref_store = client.dump_ref_store();

    let task = async {
        let mut txn_client = cluster.new_txn_client().await;

        let mut tx = txn_client.begin_pessimistic().await.unwrap();
        tx.insert(key.clone(), 0u64.to_be_bytes().to_vec())
            .await
            .unwrap();
        tx.commit().await.unwrap();

        // Concurrent "v = v + 1":
        let mut handles = Vec::with_capacity(CONCURRENCY);
        for _ in 0..CONCURRENCY {
            let cli = txn_client.inner.clone();
            let key = key.clone();

            let t = async move {
                loop {
                    let mut tx = cli.begin_pessimistic().await.unwrap();
                    match inc(&mut tx, &key).await {
                        Ok(()) => {
                            tx.commit().await.unwrap();
                            break;
                        }
                        Err(err) if is_write_conflict_err(&err) => {
                            tx.rollback().await.unwrap();
                            continue;
                        }
                        Err(err) => {
                            let _ = tx.rollback().await;
                            panic!("unexpected error: {:?}", err)
                        }
                    }
                }
            };
            handles.push(tokio::spawn(t));
        }
        try_join_all(handles).await.unwrap();

        // Verify:
        let ts = txn_client.current_timestamp().await.unwrap();
        let mut snapshot = txn_client.snapshot(ts, TransactionOptions::default());
        let v = snapshot.get(key.clone()).await.unwrap().unwrap();
        let v = u64::from_be_bytes(v.try_into().unwrap());
        assert_eq!(v, CONCURRENCY as u64);

        let cnt = txn_client
            .verify_data_by_scan(&ref_store, None)
            .await
            .unwrap();
        assert_eq!(cnt, DATA_COUNT + 1);
    };
    rt.block_on(task);
}

#[rstest]
#[case::local(false, TxnWriteMethod::Normal, false)]
#[case::ia(true, TxnWriteMethod::Normal, false)]
#[case::ia_cleanup(true, TxnWriteMethod::Normal, true)]
#[case::txn_file(true, TxnWriteMethod::FileBased, false)]
fn test_resolve_lock(
    #[case] enable_ia: bool,
    #[case] write_method: TxnWriteMethod,
    #[case] use_cleanup: bool,
) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let (cluster, keyspace_id, table_id) = prepare_cluster(enable_ia);
    let gen_key = i_to_key(keyspace_id, table_id);

    let mut client = cluster.new_client_opt(ClusterClientOptions {
        txn_file_max_chunk_size: Some(TXN_FILE_MAX_CHUNK_SIZE),
        ..Default::default()
    });

    let cases = vec![
        (
            0,
            40,
            CommitAction::AsyncCommitSecondaryKeys(Duration::ZERO),
        ),
        (
            80,
            120,
            CommitAction::AsyncCommitSecondaryKeys(Duration::MAX),
        ),
        (160, 240, CommitAction::NoCommit),
    ];
    for (start, end, commit_action) in cases {
        client
            .try_put_kv(
                start..end,
                &gen_key,
                |_| random_value(8),
                MutateOptions {
                    commit_action,
                    write_method,
                    ..Default::default()
                },
            )
            .unwrap();
    }

    if use_cleanup {
        // TxnClient rust use "cleanup" interface to rollback locks.
        let ref_store = client.dump_ref_store();
        let task = async {
            let mut txn_client = cluster.new_txn_client().await;
            let cnt = txn_client
                .verify_data_by_scan(&ref_store, None)
                .await
                .unwrap();
            assert_eq!(cnt, DATA_COUNT);
        };
        rt.block_on(task);
    }

    client.verify_data_with_ref_store();
}

#[rstest]
#[case::local(false)]
#[case::ia(true)]
fn test_resolve_pessimistic_lock(#[case] enable_ia: bool) {
    let rt = tokio::runtime::Runtime::new().unwrap();
    let _enter = rt.enter();

    let (cluster, keyspace_id, table_id) = prepare_cluster(enable_ia);
    let gen_key = i_to_key(keyspace_id, table_id);

    let key_ids = vec![40usize, 80, 120];
    let keys = key_ids.iter().map(|&k| gen_key(k)).collect::<Vec<_>>();

    let mut client = cluster.new_client();
    client.set_req_timeout(Duration::from_millis(300));

    let task = async {
        let txn_client = cluster.new_txn_client().await;
        let start_ts = client.get_ts();
        let for_update_ts = client.get_ts();
        let (_results, key_errors) = client
            .kv_pessimistic_lock(
                keys[0].clone().into(),
                keys.iter().cloned().map(Into::into).collect(),
                start_ts,
                1000,
                for_update_ts,
                None,
            )
            .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);

        // A short timeout is enough to prove the keys are locked. Cleanup cases
        // wait for the actual lock expiry below.
        let err = client
            .try_put_kv(
                key_ids.clone(),
                &gen_key,
                |_| random_value(8),
                MutateOptions::default(),
            )
            .unwrap_err();
        assert_matches!(err, ClientError::KeyErrors(errs) if errs.iter().all(|e| e.has_locked()));

        let key_errors = client
            .kv_pessimistic_rollback(keys.to_vec(), start_ts, for_update_ts)
            .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);
        // FIXME: resolve pessimistic lock is not implemented.
        client
            .try_put_kv(
                key_ids,
                &gen_key,
                |_| random_value(8),
                MutateOptions::default(),
            )
            .unwrap();

        // Exercise the TTL-cleanup path on the same fixture with distinct
        // keys, after the explicit rollback path above has completed.
        let cleanup_key_ids = vec![41usize, 81, 121];
        let cleanup_keys = cleanup_key_ids
            .iter()
            .map(|&key_id| gen_key(key_id))
            .collect::<Vec<_>>();
        let cleanup_start_ts = client.get_ts();
        let cleanup_for_update_ts = client.get_ts();
        let (_results, key_errors) = client
            .kv_pessimistic_lock(
                cleanup_keys[0].clone().into(),
                cleanup_keys.iter().cloned().map(Into::into).collect(),
                cleanup_start_ts,
                1000,
                cleanup_for_update_ts,
                None,
            )
            .unwrap();
        assert!(key_errors.is_empty(), "key_errors {:?}", key_errors);

        let err = client
            .try_put_kv(
                cleanup_key_ids,
                &gen_key,
                |_| random_value(8),
                MutateOptions::default(),
            )
            .unwrap_err();
        assert_matches!(err, ClientError::KeyErrors(errs) if errs.iter().all(|e| e.has_locked()));

        // Retry until the lock TTL expires and cleanup succeeds.
        try_wait_result_async(
            || {
                let txn_client = txn_client.clone();
                let cleanup_keys = cleanup_keys.clone();
                Box::pin(async move {
                    let mut tx = txn_client.begin_pessimistic().await.unwrap();
                    tx.lock_keys(cleanup_keys).await?;
                    tx.rollback().await
                })
            },
            3,
        )
        .await
        .unwrap();
    };
    rt.block_on(task);

    client.verify_data_with_ref_store();
}

// TODO: test_unsafe_destroy_range
// TODO: test_mvcc_by_key
// TODO: test_mvcc_by_start_ts

fn prepare_cluster(
    enable_ia: bool,
) -> (
    ServerClusterExt,
    u32, // keyspace_id
    i64, // table_id
) {
    let node_ids = alloc_node_id_vec(1);
    let cluster = ServerClusterBuilder::new(node_ids, |_, conf| {
        conf.rocksdb.writecf.write_buffer_size = ReadableSize(1024);
        conf.rocksdb.writecf.block_size = ReadableSize(256);
        conf.rocksdb.writecf.target_file_size_base = ReadableSize(2048);
        if enable_ia {
            conf.kvengine.ia = IaConfig {
                mem_cap: (64 * 1024).into(),
                disk_cap: (1024 * 1024).into(),
                segment_size: 1024,
                ..Default::default()
            }
        }
    })
    .pd_server_cnt(1)
    .tikv_worker_cnt(1)
    .build_ext();

    let ks_opts = CreateKeyspaceOptions {
        table_count: 1,
        storage_class_spec_fn: Box::new(move |_| {
            if enable_ia {
                StorageClass::Ia.into()
            } else {
                StorageClassSpec::default()
            }
        }),
        ..Default::default()
    };
    let keyspace_id = block_on(cluster.create_keyspace(&ks_opts, Duration::from_secs(10)));

    let ks_meta = cluster
        .keyspace_manager()
        .get_keyspace_meta(keyspace_id)
        .unwrap();
    let table_id = ks_meta.get_all_available_tables()[0];
    drop(ks_meta);
    let table_key = make_row_key(keyspace_id, table_id, &[]);
    cluster.wait_region_replicated(&table_key, 1);

    // Put base data to generate level 1+ tables.
    let mut client = cluster.new_client();
    client.put_kv(0..80, i_to_key(keyspace_id, table_id), i_to_value(64));
    cluster.request_major_compaction(
        MajorCompactionTarget::Table {
            keyspace_id,
            table_id,
        },
        true,
        Duration::from_secs(10),
    );
    let mut last_retry_time = Instant::now_coarse();
    TryWaiter::timeout(60)
        .interval_dur(Duration::from_millis(500))
        .must_wait(
            || {
                let Some(shard) = cluster.get_latest_shard_by_key(&table_key) else {
                    return false;
                };
                let stats = shard.get_stats();
                let compacted = !stats.write_cf_level_n_is_empty();
                if !compacted && last_retry_time.saturating_elapsed() >= Duration::from_secs(5) {
                    info!("retry major compaction, stats {:?}", stats);
                    cluster.request_major_compaction(
                        MajorCompactionTarget::Table {
                            keyspace_id,
                            table_id,
                        },
                        true,
                        Duration::from_secs(10),
                    );
                    last_retry_time = Instant::now_coarse();
                }
                compacted && shard.new_snap_access().is_sync() != enable_ia
            },
            || {
                let stats = cluster
                    .get_latest_shard_by_key(&table_key)
                    .map(|x| x.get_stats());
                format!("wait for compacted & async timeout: {:?}", stats)
            },
        );

    // Make mem-tables & L0 data.
    client.put_kv(
        40..DATA_COUNT,
        i_to_key(keyspace_id, table_id),
        i_to_value(32),
    );

    let (existed, deleted) = client.verify_data_with_ref_store();
    assert_eq!(existed, DATA_COUNT);
    assert_eq!(deleted, 0);
    info!(
        "shard: {:?}",
        cluster
            .get_latest_shard_by_key(&table_key)
            .map(|x| x.get_stats())
    );
    (cluster, keyspace_id, table_id)
}

fn i_to_key(keyspace_id: u32, table_id: i64) -> impl Fn(usize) -> Vec<u8> {
    move |i| make_row_key(keyspace_id, table_id, format!("k{i:08}").as_bytes())
}

fn i_to_value(value_len: usize) -> impl Fn(usize) -> Vec<u8> {
    move |_| random_value(value_len)
}

fn make_put(key: Bytes, value: Bytes) -> Mutation {
    Mutation {
        op: kvrpcpb::Op::Put,
        key,
        value,
        ..Default::default()
    }
}

fn make_insert(key: Bytes, value: Bytes) -> Mutation {
    Mutation {
        op: kvrpcpb::Op::Insert,
        key,
        value,
        ..Default::default()
    }
}

async fn inc(tx: &mut Transaction, key: &[u8]) -> tikv_client::Result<()> {
    let v = tx.get_for_update(key.to_vec()).await?.unwrap();
    let v = u64::from_be_bytes(v.try_into().unwrap());
    tx.put(key.to_vec(), (v + 1).to_be_bytes().to_vec()).await?;
    Ok(())
}

// Avoid exposing tikv-client’s private proto types.
fn is_write_conflict_err(err: &tikv_client::Error) -> bool {
    match err {
        tikv_client::Error::KeyError(e) => e.conflict.is_some(),
        tikv_client::Error::MultipleKeyErrors(errs) => errs.iter().all(is_write_conflict_err),
        tikv_client::Error::PessimisticLockError { inner, .. } => is_write_conflict_err(inner),
        _ => false,
    }
}