// Copyright 2026 PingCAP, Inc.

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

use api_version::{ApiV2, api_v2::TIDB_META_KEY_PREFIX};
use cloud_worker::CloudWorker;
use futures::{SinkExt, StreamExt};
use grpcio::{ChannelBuilder, EnvBuilder, Environment, WriteFlags};
use http::Uri;
use hyper::{Body, Method, Request};
use kvproto::{
    cdcpb::{ChangeDataClient, ChangeDataRequest, ChangeDataRequestRegister, EventLogType},
    kvrpcpb::{Context, ScanLockRequest},
    tikvpb::TikvClient,
};
use native_br::{backup, backup_worker};
use pd_client::Config as PdConfig;
use test_cloud_server::{
    ServerClusterBuilder, alloc_node_id_vec, keyspace::make_row_key, must_wait_result,
    oss::prepare_dfs, ticdc::KeyspaceTestService,
};
use test_util::alloc_port;
use tikv::config::TikvConfig;
use tikv_util::{codec::bytes::encode_bytes, config::ReadableDuration};

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

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

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

    // Split at the end of TiDB meta prefix so meta keys (`m...`) are isolated.
    // The source/target regions in this test then stay non-pinned in lazy mode.
    let mut meta_end_key_full = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID);
    meta_end_key_full.push(TIDB_META_KEY_PREFIX + 1); // 'n'
    client.split(&meta_end_key_full);
    cluster.wait_region_replicated(&meta_end_key_full, 3);

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

    // Split inside the keyspace to build a merge source/target pair.
    let split_key = key(50);
    client.split(&split_key);
    cluster.wait_region_replicated(&split_key, 3);

    client.put_kv(
        0..100,
        |i| key(i as u32),
        |i| format!("v{i:08}").into_bytes(),
    );
    client.verify_data_with_ref_store();

    let source_key = key(10);
    let target_key = key(90);
    let keyspace_prefix_len = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID).len();
    let source_key_inner = source_key[keyspace_prefix_len..].to_vec();
    let target_key_inner = target_key[keyspace_prefix_len..].to_vec();

    // 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 (lazy shard load).
    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.lazy_load.enabled = 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.enable_resolve_lock = true;
    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 keyspace_pd_client = keyspace_svc.pd_client_ext();
    let source_region = keyspace_pd_client
        .get_region(&encode_bytes(&source_key_inner))
        .expect("get source region");
    let target_region = keyspace_pd_client
        .get_region(&encode_bytes(&target_key_inner))
        .expect("get target region");
    assert_ne!(
        source_region.get_id(),
        target_region.get_id(),
        "source/target must be different regions"
    );

    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.clone());
    let cdc = ChangeDataClient::new(channel);

    // In lazy mode, keyspace add should not eagerly load shards, so `kv_scan_lock`
    // should return region_not_found before any register.
    let region_id = source_region.get_id();
    let mut ctx = Context::default();
    ctx.set_region_id(region_id);
    ctx.mut_region_epoch()
        .set_version(source_region.get_region_epoch().get_version());
    ctx.mut_region_epoch()
        .set_conf_ver(source_region.get_region_epoch().get_conf_ver());

    let mut scan_lock_req = ScanLockRequest::default();
    scan_lock_req.set_context(ctx.clone());
    scan_lock_req.set_max_version(0);
    scan_lock_req.set_limit(1);
    let scan_lock_resp = tikv.kv_scan_lock(&scan_lock_req).unwrap();
    assert!(scan_lock_resp.has_region_error());
    assert!(scan_lock_resp.get_region_error().has_region_not_found());

    // Register the source region via CDC event feed. This should trigger on-demand
    // shard load.
    let (mut req_tx, mut resp_rx) = cdc.event_feed().unwrap();
    let request_id = 1u64;
    let mut reg_req = ChangeDataRequest::default();
    reg_req.set_region_id(region_id);
    reg_req.set_request_id(request_id);
    reg_req.set_region_epoch(source_region.get_region_epoch().clone());
    reg_req.set_checkpoint_ts(0);
    reg_req.set_register(ChangeDataRequestRegister::new());

    rt.block_on(async {
        req_tx.send((reg_req, WriteFlags::default())).await.unwrap();
    });

    // Wait for the Initialized event (and fail fast on any error event).
    rt.block_on(async {
        let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
        loop {
            let next = tokio::time::timeout_at(deadline, resp_rx.next())
                .await
                .expect("wait Initialized timeout");
            let Some(ev_res) = next else {
                panic!("event feed closed unexpectedly");
            };
            let ev = ev_res.expect("event feed recv");
            for event in ev.get_events() {
                if event.get_region_id() != region_id || event.get_request_id() != request_id {
                    continue;
                }
                if event.has_error() {
                    panic!("unexpected register error: {:?}", event.get_error());
                }
                let rows = event.get_entries().get_entries();
                if rows
                    .iter()
                    .any(|row| row.get_type() == EventLogType::Initialized)
                {
                    return;
                }
            }
        }
    });

    // The shard should now be loaded, so `kv_scan_lock` should not return a region
    // error.
    let scan_lock_resp = tikv.kv_scan_lock(&scan_lock_req).unwrap();
    assert!(
        !scan_lock_resp.has_region_error(),
        "unexpected region_error: {:?}",
        scan_lock_resp.get_region_error()
    );

    // Merge the source region into the target region on upstream. The merge target
    // shard is not registered (and thus stays unloaded), so replication_worker
    // must rely on merged_engine topology-change destroy signals to send
    // region_not_found for the source subscription.
    assert!(
        client.try_merge_and_wait(&source_key, &target_key, 30),
        "merge did not finish in time"
    );

    rt.block_on(async {
        let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
        loop {
            let next = tokio::time::timeout_at(deadline, resp_rx.next())
                .await
                .expect("wait region_not_found timeout");
            let Some(ev_res) = next else {
                panic!("event feed closed unexpectedly");
            };
            let ev = ev_res.expect("event feed recv");
            for event in ev.get_events() {
                if event.get_region_id() == region_id
                    && event.get_request_id() == request_id
                    && event.has_error()
                {
                    let err = event.get_error();
                    assert!(
                        err.has_region_not_found(),
                        "unexpected error event: {:?}",
                        err
                    );
                    return;
                }
            }
        }
    });

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