// Copyright 2026 PingCAP, Inc.

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

use cloud_worker::CloudWorker;
use http::Uri;
use hyper::{Body, Method, Request};
use native_br::{backup, backup_worker};
use pd_client::Config as PdConfig;
use test_cloud_server::{
    ServerClusterBuilder, alloc_node_id_vec,
    client::{ClusterClient, CommitAction, MutateOptions, PrewriteExt, TxnMutations},
    keyspace::make_row_key,
    must_wait_result,
    oss::prepare_dfs,
    ticdc::KeyspaceTestService,
    try_wait_result,
    util::Mutation,
};
use test_util::alloc_port;
use tikv::config::TikvConfig;
use tikv_client::TimestampExt as _;
use tikv_util::config::ReadableDuration;
use txn_types::{LockType, TimeStamp, WriteType};

const KEYSPACE_ID: u32 = 1;
const TABLE_ID: i64 = 42;

fn put_mutation(key: Vec<u8>, value: Vec<u8>) -> Mutation {
    Mutation {
        op: kvproto::kvrpcpb::Op::Put,
        key: key.into(),
        value: value.into(),
        ..Default::default()
    }
}

use crate::transaction::helper::{
    get_mvcc_info, must_locked_with_properties, must_unlocked, must_written_with_properties,
    new_delete_mutation, new_lock_mutation, new_shared_lock_mutation,
};

fn has_write(info: &kvproto::kvrpcpb::MvccInfo, commit_ts: u64, op: kvproto::kvrpcpb::Op) -> bool {
    info.get_writes()
        .iter()
        .any(|w| w.commit_ts == commit_ts && w.get_type() == op)
}

fn has_write_with_start_ts(info: &kvproto::kvrpcpb::MvccInfo, start_ts: TimeStamp) -> bool {
    let start_ts = start_ts.into_inner();
    info.get_writes()
        .iter()
        .any(|write| write.get_start_ts() == start_ts)
}

fn shared_holder_start_tss(lock_info: &kvproto::kvrpcpb::LockInfo) -> Vec<u64> {
    let mut start_tss = lock_info
        .get_shared_lock_infos()
        .iter()
        .map(|lock_info| lock_info.get_lock_version())
        .collect::<Vec<_>>();
    start_tss.sort_unstable();
    start_tss
}

fn shared_lock_holder_start_tss_from_conflict(
    client: &mut ClusterClient,
    key: &[u8],
) -> Result<Vec<u64>, String> {
    // MvccInfo can prove the shared key still has a lock, but the conflict
    // response is the RPC surface that exposes the full shared_lock_infos list.
    let start_ts = client.get_ts();
    let for_update_ts = client.get_ts();
    let (_results, key_errors) = client
        .kv_pessimistic_lock(
            key.to_vec().into(),
            vec![key.to_vec().into()],
            start_ts,
            3000,
            for_update_ts,
            Some(1),
        )
        .map_err(|err| format!("pessimistic lock probe failed: {err:?}"))?;
    if key_errors.len() != 1 || !key_errors[0].has_locked() {
        let rollback_errors = if key_errors.is_empty() {
            client
                .kv_pessimistic_rollback(vec![key.to_vec()], start_ts, for_update_ts)
                .map_err(|err| format!("rollback successful probe lock failed: {err:?}"))?
        } else {
            vec![]
        };
        return Err(format!(
            "expected one shared-lock conflict on {:?}, got {:?}, probe rollback errors {:?}",
            key, key_errors, rollback_errors
        ));
    }

    let locked = key_errors[0].get_locked();
    if locked.get_lock_type() != kvproto::kvrpcpb::Op::SharedLock {
        return Err(format!(
            "expected shared-lock conflict on {:?}, got {:?}",
            key, locked
        ));
    }
    if locked.get_key() != key {
        return Err(format!(
            "shared-lock conflict key mismatch, expected {:?}, got {:?}",
            key,
            locked.get_key()
        ));
    }
    Ok(shared_holder_start_tss(locked))
}

fn must_shared_lock_holder_start_tss(
    client: &mut ClusterClient,
    key: &[u8],
    expected_start_tss: &[TimeStamp],
) {
    let mut expected = expected_start_tss
        .iter()
        .map(|ts| ts.into_inner())
        .collect::<Vec<_>>();
    expected.sort_unstable();
    let actual = shared_lock_holder_start_tss_from_conflict(client, key).unwrap();
    assert_eq!(actual, expected);
}

#[test]
fn test_resolve_lock() {
    let (_tmp_dir, mut oss, dfs_config) = prepare_dfs("test_resolve_lock_");

    let nodes = alloc_node_id_vec(3);
    let mut cluster = ServerClusterBuilder::new(nodes, |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    })
    .pd_server_cnt(1)
    .build();
    cluster.wait_region_replicated(&[], 3);

    let pd_client = cluster.get_pure_pd_client();
    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let key = |i: u32| make_row_key(KEYSPACE_ID, TABLE_ID, format!("k{i:08}").as_bytes());

    // Split inside the keyspace to make sure secondary keys can be on different
    // regions.
    for split_i in [50u32, 150u32] {
        client.split(&key(split_i));
    }
    cluster.wait_pd_region_min_count(4);

    // Ensure the keyspace shows up in backup.
    client.put_kv(
        0..10,
        |i| key(i as u32),
        |i| format!("v{i:08}").into_bytes(),
    );
    client.verify_data_with_ref_store();

    // Take an initial backup for replication worker to recover from.
    let rt = tokio::runtime::Runtime::new().unwrap();
    let backup_config = backup::BackupConfig {
        dfs: dfs_config.clone(),
        skip_keyspace_meta: true,
        backup_delay: ReadableDuration::millis(1),
        ..Default::default()
    };
    let backup_worker = backup_worker::BackupWorker::new(
        backup_config,
        pd_client.clone(),
        Duration::ZERO,
        Duration::from_secs(1),
    );
    rt.block_on(backup_worker.instant_backup())
        .expect("instant_backup");

    // Start a CloudWorker with replication worker enabled.
    let worker_port = alloc_port();
    let grpc_port = alloc_port();
    let data_dir = tempfile::Builder::new()
        .prefix("rep_worker_")
        .tempdir()
        .unwrap();

    let mut worker_conf = cloud_worker::Config::default();
    worker_conf.addr = format!("127.0.0.1:{worker_port}");
    worker_conf.pd = PdConfig::new(cluster.pd_endpoints().to_vec());
    worker_conf.dfs = dfs_config.clone();
    worker_conf.data_dir = data_dir.path().display().to_string();
    worker_conf.replication_worker.enabled = true;
    worker_conf.replication_worker.enable_resolve_lock = true;
    worker_conf.replication_worker.grpc_addr = format!("127.0.0.1:{grpc_port}");
    worker_conf.replication_worker.advertise_addr =
        worker_conf.replication_worker.grpc_addr.clone();
    worker_conf.replication_worker.sync_interval = ReadableDuration::secs(1);
    worker_conf.replication_worker.report_region_interval = ReadableDuration::secs(1);
    worker_conf.replication_worker.merged_engine.mem_table_size = cluster.get_mem_table_size();
    let merged_store_id = worker_conf.replication_worker.merged_engine.merged_store_id;

    let mut worker = CloudWorker::new(worker_conf.clone(), None, 2, pd_client.clone());
    worker.start();

    // Add keyspace to replication worker using an in-process test PD + mock TiCDC.
    let keyspace_svc = KeyspaceTestService::new(
        KEYSPACE_ID,
        merged_store_id,
        worker_conf.replication_worker.advertise_addr.clone(),
    );

    let http = hyper::Client::new();
    let add_keyspace_uri = Uri::from_str(&format!(
        "http://{}/cdc/keyspace?keyspace_id={KEYSPACE_ID}",
        worker.addr()
    ))
    .unwrap();
    let add_keyspace_body = format!(
        r#"{{"pd_url":"{}","cdc_addr":"{}"}}"#,
        keyspace_svc.pd_url(),
        keyspace_svc.cdc_addr()
    );
    let req = Request::builder()
        .method(Method::POST)
        .uri(add_keyspace_uri)
        .body(Body::from(add_keyspace_body))
        .unwrap();
    let resp = rt.block_on(http.request(req)).unwrap();
    assert!(resp.status().is_success(), "add keyspace failed");

    // Wait until the replication worker store is registered to the test PD.
    must_wait_result(
        || {
            let stores = keyspace_svc
                .pd_client_ext()
                .get_all_stores(true)
                .map_err(|e| format!("{e:?}"))?;
            if stores.iter().any(|s| s.get_id() == merged_store_id) {
                Ok(())
            } else {
                Err(format!("store {merged_store_id} not registered yet"))
            }
        },
        10,
        || "wait replication worker store registered".into(),
    );
    let rep_txn_client = rt.block_on(keyspace_svc.new_txn_client());

    // Create transaction states on upstream:
    // 1) not committed, with both Put and Delete locks
    // 2) partially committed (primary committed, secondary locks left)
    // 3) an uncommitted shared-lock transaction
    // 4) a later stale Put lock in the same region as the shared lock
    // 5) fully committed
    let no_commit_put_keys = [10u32, 60, 160].map(|i| key(i));
    let no_commit_delete_keys = [12u32, 62, 162].map(|i| key(i));
    let no_commit_keys = no_commit_put_keys
        .iter()
        .chain(no_commit_delete_keys.iter())
        .cloned()
        .collect::<Vec<_>>();
    let partial_keys = [20u32, 70, 170].map(|i| key(i));
    let shared_primary_key = key(190);
    let shared_key = key(191);
    let later_no_commit_key = key(192);
    let full_commit_keys = [30u32, 80, 180].map(|i| key(i));

    let lock_ttl = PrewriteExt {
        lock_ttl: Duration::from_millis(1),
        ..Default::default()
    };

    let start_ts_no_commit = client.get_ts();
    let no_commit_muts = no_commit_put_keys
        .iter()
        .enumerate()
        .map(|(idx, k)| put_mutation(k.clone(), format!("no_commit_v{idx}").into_bytes()))
        .chain(
            no_commit_delete_keys
                .iter()
                .map(|k| new_delete_mutation(k.clone())),
        )
        .collect::<Vec<_>>();
    client
        .prewrite_and_commit(
            no_commit_muts,
            MutateOptions {
                start_ts: Some(start_ts_no_commit),
                commit_action: CommitAction::NoCommit,
                ..Default::default()
            },
            lock_ttl,
        )
        .unwrap();

    let start_ts_partial = client.get_ts();
    let partial_muts = partial_keys
        .iter()
        .enumerate()
        .map(|(idx, k)| put_mutation(k.clone(), format!("partial_v{idx}").into_bytes()))
        .collect::<Vec<_>>();
    let commit_ts_partial = client
        .prewrite_and_commit(
            partial_muts,
            MutateOptions {
                start_ts: Some(start_ts_partial),
                commit_action: CommitAction::AsyncCommitSecondaryKeys(Duration::MAX),
                ..Default::default()
            },
            lock_ttl,
        )
        .unwrap();

    let start_ts_shared = client.get_ts();
    let for_update_ts_shared = client.get_ts();
    client
        .kv_pessimistic_lock(
            shared_primary_key.clone().into(),
            vec![shared_primary_key.clone().into()],
            start_ts_shared,
            3000,
            for_update_ts_shared,
            None,
        )
        .unwrap();
    client
        .kv_shared_pessimistic_lock(
            shared_primary_key.clone().into(),
            vec![shared_key.clone().into()],
            start_ts_shared,
            3000,
            for_update_ts_shared,
            None,
        )
        .unwrap();
    let mut shared_txn = client.begin_transaction(Some(start_ts_shared));
    client
        .kv_prewrite_ext(
            shared_primary_key.clone().into(),
            Some(&vec![shared_key.clone().into()]),
            TxnMutations::from_normal(vec![
                new_lock_mutation(shared_primary_key.clone()),
                new_shared_lock_mutation(shared_key.clone()),
            ]),
            &mut shared_txn,
            PrewriteExt {
                for_update_ts: for_update_ts_shared,
                ..Default::default()
            },
        )
        .unwrap();

    let start_ts_later_no_commit = client.get_ts();
    client
        .prewrite_and_commit(
            vec![put_mutation(
                later_no_commit_key.clone(),
                b"later_no_commit_v".to_vec(),
            )],
            MutateOptions {
                start_ts: Some(start_ts_later_no_commit),
                commit_action: CommitAction::NoCommit,
                ..Default::default()
            },
            lock_ttl,
        )
        .unwrap();

    let start_ts_full = client.get_ts();
    let full_muts = full_commit_keys
        .iter()
        .enumerate()
        .map(|(idx, k)| put_mutation(k.clone(), format!("full_v{idx}").into_bytes()))
        .collect::<Vec<_>>();
    let commit_ts_full = client
        .prewrite_and_commit(
            full_muts,
            MutateOptions {
                start_ts: Some(start_ts_full),
                commit_action: CommitAction::SyncCommit,
                ..Default::default()
            },
            lock_ttl,
        )
        .unwrap();
    let scan_ts = tikv_client::Timestamp::from_version(commit_ts_full.into_inner());

    // Sanity check initial states.
    for k in no_commit_put_keys.iter() {
        must_locked_with_properties(
            &mut client,
            k,
            &no_commit_keys[0],
            start_ts_no_commit,
            LockType::Put,
        );
    }
    for k in no_commit_delete_keys.iter() {
        must_locked_with_properties(
            &mut client,
            k,
            &no_commit_keys[0],
            start_ts_no_commit,
            LockType::Delete,
        );
    }
    {
        must_unlocked(&mut client, &partial_keys[0]);
        must_written_with_properties(
            &mut client,
            &partial_keys[0],
            start_ts_partial,
            commit_ts_partial,
            WriteType::Put,
        );
    }
    for k in partial_keys.iter().skip(1) {
        must_locked_with_properties(
            &mut client,
            k,
            &partial_keys[0],
            start_ts_partial,
            LockType::Put,
        );
    }
    must_locked_with_properties(
        &mut client,
        &shared_primary_key,
        &shared_primary_key,
        start_ts_shared,
        LockType::Lock,
    );
    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_shared]);
    {
        let shared_info =
            get_mvcc_info(&mut client, &shared_key).expect("shared key mvcc info missing");
        assert!(
            !has_write_with_start_ts(&shared_info, start_ts_shared),
            "shared key must not have a write before replication_worker scan: {:?}",
            shared_info
        );
    }
    must_locked_with_properties(
        &mut client,
        &later_no_commit_key,
        &later_no_commit_key,
        start_ts_later_no_commit,
        LockType::Put,
    );
    for k in full_commit_keys.iter() {
        must_unlocked(&mut client, k);
        must_written_with_properties(
            &mut client,
            k,
            start_ts_full,
            commit_ts_full,
            WriteType::Put,
        );
    }

    // Send scan-lock requests to replication worker and wait until it resolves them
    // upstream.
    let wait_start = Instant::now();
    try_wait_result(
        || -> Result<(), String> {
            let locks = rt
                .block_on(rep_txn_client.scan_locks(&scan_ts, .., 1024))
                .map_err(|e| format!("scan_locks failed: {e:?}"))?;
            assert!(locks.is_empty(), "locks: {:?}", locks); // Replication worker always return empty.

            for k in no_commit_keys.iter() {
                let Some(info) = get_mvcc_info(&mut client, k) else {
                    return Err("no-commit key mvcc info missing".into());
                };
                if info.has_lock() {
                    return Err("no-commit lock still exists".into());
                }
                if !has_write(
                    &info,
                    start_ts_no_commit.into_inner(),
                    kvproto::kvrpcpb::Op::Rollback,
                ) {
                    return Err("no-commit rollback record missing".into());
                }
            }

            for k in partial_keys.iter() {
                let Some(info) = get_mvcc_info(&mut client, k) else {
                    return Err("partial-commit key mvcc info missing".into());
                };
                if info.has_lock() {
                    return Err("partial-commit lock still exists".into());
                }
                if !has_write(
                    &info,
                    commit_ts_partial.into_inner(),
                    kvproto::kvrpcpb::Op::Put,
                ) {
                    return Err("partial-commit write record missing".into());
                }
            }

            let Some(later_info) = get_mvcc_info(&mut client, &later_no_commit_key) else {
                return Err("later no-commit key mvcc info missing".into());
            };
            if later_info.has_lock() {
                return Err("later no-commit lock still exists".into());
            }
            if !has_write(
                &later_info,
                start_ts_later_no_commit.into_inner(),
                kvproto::kvrpcpb::Op::Rollback,
            ) {
                return Err("later no-commit rollback record missing".into());
            }

            let Some(shared_info) = get_mvcc_info(&mut client, &shared_key) else {
                return Err("shared key mvcc info missing".into());
            };
            if !shared_info.has_lock() {
                return Err("shared lock was unexpectedly removed".into());
            }
            let mut expected_shared_holders = vec![start_ts_shared.into_inner()];
            expected_shared_holders.sort_unstable();
            let actual_shared_holders =
                shared_lock_holder_start_tss_from_conflict(&mut client, &shared_key)?;
            if actual_shared_holders != expected_shared_holders {
                return Err(format!(
                    "shared key has unexpected holders, expected {:?}, got {:?}",
                    expected_shared_holders, actual_shared_holders
                ));
            }
            if has_write_with_start_ts(&shared_info, start_ts_shared) {
                return Err(format!(
                    "shared key got unexpected write for start_ts {}: {:?}",
                    start_ts_shared.into_inner(),
                    shared_info
                ));
            }

            Ok(())
        },
        30,
    )
    .unwrap_or_else(|e| {
        panic!(
            "wait for scan locks timed out after {:?}: {:?}",
            wait_start.elapsed(),
            e
        )
    });

    // Final assertions (ref store verification ensures full-commit and
    // partial-commit apply).
    for k in no_commit_keys.iter() {
        must_unlocked(&mut client, k);
        must_written_with_properties(
            &mut client,
            k,
            start_ts_no_commit,
            start_ts_no_commit,
            WriteType::Rollback,
        );
    }
    for k in partial_keys.iter() {
        must_unlocked(&mut client, k);
        must_written_with_properties(
            &mut client,
            k,
            start_ts_partial,
            commit_ts_partial,
            WriteType::Put,
        );
    }
    must_unlocked(&mut client, &later_no_commit_key);
    must_written_with_properties(
        &mut client,
        &later_no_commit_key,
        start_ts_later_no_commit,
        start_ts_later_no_commit,
        WriteType::Rollback,
    );
    must_locked_with_properties(
        &mut client,
        &shared_primary_key,
        &shared_primary_key,
        start_ts_shared,
        LockType::Lock,
    );
    must_shared_lock_holder_start_tss(&mut client, &shared_key, &[start_ts_shared]);
    {
        let shared_info =
            get_mvcc_info(&mut client, &shared_key).expect("shared key mvcc info missing");
        assert!(
            !has_write_with_start_ts(&shared_info, start_ts_shared),
            "replication_worker must not create rollback/commit writes for shared key: {:?}",
            shared_info
        );
    }

    client
        .kv_resolve_lock_batch(&shared_primary_key, &[(start_ts_shared, TimeStamp::zero())])
        .unwrap();
    must_unlocked(&mut client, &shared_primary_key);
    must_unlocked(&mut client, &shared_key);

    client.verify_data_with_ref_store();

    worker.shutdown();
    cluster.stop();
    oss.shutdown();
}