// Copyright 2026 PingCAP, Inc.

use std::{str::FromStr, sync::Arc, time::Duration};

use api_version::ApiV2;
use cloud_worker::CloudWorker;
use grpcio::{ChannelBuilder, EnvBuilder, Environment};
use http::Uri;
use hyper::{Body, Method, Request};
use kvproto::{
    kvrpcpb::{Context, GetRequest, ScanRequest},
    tikvpb::TikvClient,
};
use native_br::{backup, backup_worker};
use pd_client::Config as PdConfig;
use rstest::rstest;
use test_cloud_server::{
    ServerClusterBuilder, alloc_node_id_vec, must_wait_result, oss::prepare_dfs,
    ticdc::KeyspaceTestService,
};
use tikv::config::TikvConfig;
use tikv_util::{
    codec::{bytes::encode_bytes, number::NumberEncoder},
    config::ReadableDuration,
};

const KEYSPACE_ID: u32 = 1;

const TIDB_SCHEMA_VERSION_KEY: &[u8] = b"SchemaVersionKey";
const TIDB_META_KEY_PREFIX: u8 = b'm';
const STRING_DATA_TYPE: u8 = b's';

fn schema_version_key(keyspace_id: u32) -> Vec<u8> {
    let mut key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
    let encoded_key = encode_bytes(TIDB_SCHEMA_VERSION_KEY);
    key.reserve(encoded_key.len() + 1 + 8);
    key.push(TIDB_META_KEY_PREFIX);
    key.extend_from_slice(&encoded_key);
    key.encode_u64(STRING_DATA_TYPE as u64).unwrap();
    key
}

#[rstest]
#[case(false)]
#[case::lazy(true)]
fn test_replication_worker_kv_get_and_scan(#[case] lazy_load_shards: bool) {
    let (_tmp_dir, mut oss, dfs_config) = prepare_dfs("test_replication_worker_kv_get_scan_");

    let nodes = alloc_node_id_vec(3);
    // Node IDs are locked across nextest processes, and this cluster doesn't
    // start the corresponding workers, so their worker ports are collision-free.
    let worker_port = 17000 + nodes[0];
    let grpc_port = 17000 + nodes[1];
    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 mut meta_low_key_full = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID);
    let meta_low_encoded = encode_bytes(b"A");
    meta_low_key_full.reserve(meta_low_encoded.len() + 1 + 8);
    meta_low_key_full.push(TIDB_META_KEY_PREFIX);
    meta_low_key_full.extend_from_slice(&meta_low_encoded);
    meta_low_key_full
        .encode_u64(STRING_DATA_TYPE as u64)
        .unwrap();
    let meta_key_full = schema_version_key(KEYSPACE_ID);

    let meta_low_value = b"111".to_vec();
    let meta_value = b"123".to_vec();
    client.put_kv(
        0..1,
        |_| meta_low_key_full.clone(),
        |_| meta_low_value.clone(),
    );
    client.put_kv(0..1, |_| meta_key_full.clone(), |_| meta_value.clone());
    client.verify_data_with_ref_store();
    let committed_read_ts = client.get_ts().into_inner();

    // 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 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.lazy_load.enabled = lazy_load_shards;
    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();

    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");

    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 keyspace_pd_client = keyspace_svc.pd_client_ext();

    let env: Arc<Environment> = Arc::new(EnvBuilder::new().cq_count(1).build());
    let channel = ChannelBuilder::new(env)
        .keepalive_time(Duration::from_secs(1))
        .keepalive_timeout(Duration::from_secs(1))
        .connect(&worker_conf.replication_worker.grpc_addr);
    let tikv = TikvClient::new(channel);

    let prefix_len = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID).len();
    let meta_low_key_inner = meta_low_key_full[prefix_len..].to_vec();
    let meta_key_inner = meta_key_full[prefix_len..].to_vec();

    let meta_region = keyspace_pd_client
        .get_region(&encode_bytes(&meta_key_inner))
        .expect("get meta region from pd");
    let mut meta_ctx = Context::default();
    meta_ctx.set_region_id(meta_region.get_id());
    meta_ctx
        .mut_region_epoch()
        .set_version(meta_region.get_region_epoch().get_version());
    meta_ctx
        .mut_region_epoch()
        .set_conf_ver(meta_region.get_region_epoch().get_conf_ver());

    let meta_low_region = keyspace_pd_client
        .get_region(&encode_bytes(&meta_low_key_inner))
        .expect("get meta-low region from pd");
    let mut meta_low_ctx = Context::default();
    meta_low_ctx.set_region_id(meta_low_region.get_id());
    meta_low_ctx
        .mut_region_epoch()
        .set_version(meta_low_region.get_region_epoch().get_version());
    meta_low_ctx
        .mut_region_epoch()
        .set_conf_ver(meta_low_region.get_region_epoch().get_conf_ver());

    must_wait_result(
        || {
            let mut req = GetRequest::default();
            req.set_context(meta_ctx.clone());
            req.set_key(meta_key_inner.clone());
            req.set_version(committed_read_ts);
            let resp = tikv.kv_get(&req).map_err(|e| format!("{e:?}"))?;
            if resp.has_region_error() {
                return Err(format!("region_error: {:?}", resp.get_region_error()));
            }
            if resp.get_not_found() {
                return Err("kv_get not_found".to_string());
            }
            if resp.get_value() != meta_value.as_slice() {
                return Err(format!("unexpected value: {:?}", resp.get_value()));
            }
            Ok(())
        },
        60,
        || "wait kv_get meta key success".into(),
    );

    must_wait_result(
        || {
            let mut req = GetRequest::default();
            req.set_context(meta_low_ctx.clone());
            req.set_key(meta_low_key_inner.clone());
            req.set_version(committed_read_ts);
            let resp = tikv.kv_get(&req).map_err(|e| format!("{e:?}"))?;
            if resp.has_region_error() {
                return Err(format!("region_error: {:?}", resp.get_region_error()));
            }
            if resp.get_not_found() {
                return Err("kv_get not_found".to_string());
            }
            if resp.get_value() != meta_low_value.as_slice() {
                return Err(format!("unexpected value: {:?}", resp.get_value()));
            }
            Ok(())
        },
        60,
        || "wait kv_get meta-low key success".into(),
    );

    let mut scan_req = ScanRequest::default();
    scan_req.set_context(meta_low_ctx);
    scan_req.set_start_key(vec![b'm']);
    scan_req.set_end_key(vec![b'n']);
    scan_req.set_limit(128);
    scan_req.set_version(committed_read_ts);
    let scan_resp = tikv.kv_scan(&scan_req).unwrap();
    assert!(
        !scan_resp.has_region_error(),
        "region_error: {:?}",
        scan_resp.get_region_error()
    );
    assert!(
        !scan_resp.has_error(),
        "scan key error: {:?}",
        scan_resp.get_error()
    );
    assert!(
        scan_resp
            .get_pairs()
            .iter()
            .any(|pair| pair.get_key() == meta_key_inner && pair.get_value() == meta_value),
        "schema version kv not found in scan pairs: {:?}",
        scan_resp.get_pairs()
    );
    assert!(
        scan_resp
            .get_pairs()
            .iter()
            .any(|pair| pair.get_key() == meta_low_key_inner && pair.get_value() == meta_low_value),
        "meta-low kv not found in scan pairs: {:?}",
        scan_resp.get_pairs()
    );

    let mut scan_key_only_req = scan_req.clone();
    scan_key_only_req.set_key_only(true);
    let scan_key_only_resp = tikv.kv_scan(&scan_key_only_req).unwrap();
    assert!(
        !scan_key_only_resp.has_region_error(),
        "region_error: {:?}",
        scan_key_only_resp.get_region_error()
    );
    assert!(
        !scan_key_only_resp.has_error(),
        "scan key error: {:?}",
        scan_key_only_resp.get_error()
    );
    assert!(
        scan_key_only_resp
            .get_pairs()
            .iter()
            .any(|pair| pair.get_key() == meta_key_inner),
        "schema version key not found in key-only scan pairs: {:?}",
        scan_key_only_resp.get_pairs()
    );
    assert!(
        scan_key_only_resp
            .get_pairs()
            .iter()
            .any(|pair| pair.get_key() == meta_low_key_inner),
        "meta-low key not found in key-only scan pairs: {:?}",
        scan_key_only_resp.get_pairs()
    );

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