// Copyright 2025 PingCAP, Inc.

use std::{fs, sync::Arc, time::Duration};

use api_version::ApiV2;
use futures::executor::block_on;
use http::{Request, StatusCode};
use hyper::{Body, Client};
use kvengine::{
    WRITE_CF, WRITE_CF_BOTTOM_LEVEL,
    table::{ChecksumType, ZSTD_COMPRESSION},
    util::new_table_create_pb,
};
use kvenginepb::ChangeSet;
use load_data::task::{LoadDataConfig, LoadDataContext};
use matches::assert_matches;
use pd_client::PdClient;
use protobuf::Message;
use rand::Rng;
use test_cloud_server::{
    ServerCluster, alloc_node_id_vec,
    client::RequestOptions,
    load_data::{Error as LoadDataError, build, cleanup, init_task, put_chunks, try_wait_finished},
    oss::prepare_dfs,
};
use tikv::config::TikvConfig;
use tikv_util::{AsyncSafeRuntime, codec::bytes::encode_bytes};

use crate::cases::{i_to_row_key, i_to_val};

const KEYSPACE_ID: u32 = 123;
const TABLE_ID: i64 = 5;
const DATA_COUNT: usize = 2000;
const WRITER_COUNT: usize = 5;
const DATA_BATCH_SIZE: usize = 10;
const COMPRESSION_TYPE: u8 = ZSTD_COMPRESSION;
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);

#[test]
fn test_load_data() {
    let low_space_fp = if rand::thread_rng().gen_bool(0.5) {
        "engine_is_low_space" // leader is low space.
    } else {
        "is_store_low_space" // follower is low space.
    };

    let (temp_dir, _oss, dfs_conf) = prepare_dfs("test");
    let base_dir = temp_dir.path();

    let runtime = Arc::new(AsyncSafeRuntime::new(
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(4)
            .thread_name("load_data_worker")
            .enable_all()
            .build()
            .unwrap(),
    ));
    let node_ids = alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(node_ids.clone(), |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_conf.clone();
    });
    cluster.wait_region_replicated(&[], 3);
    let pd_client = cluster.get_pd_client();
    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    // Init task.
    // Total data size is about 1.3MB = 10000 * (23 + 120)
    let load_data_config = LoadDataConfig {
        kvpairs_worker_num: 2,
        building_worker_num: 2,
        max_in_mem_size: 1024, // 1KB
        flush_batch_size: 512,
        block_size: 1024,
        sst_file_size: 4 * 1024,
        region_size: 16 * 1024,
        coarse_split_size: 128 * 1024,
        enable_checkpoint: false,
        rg_config: None,
        checksum_type: ChecksumType::Crc32,
        ..Default::default()
    };

    let dfs = kvengine::dfs::new_dfs_from_config(dfs_conf);
    let master_key = cluster.get_kvengine(node_ids[0]).get_master_key();
    let encryption_key_manager = cluster
        .get_kvengine(node_ids[0])
        .get_encryption_key_manager();
    let load_data_dir = base_dir.join("load_data");
    fs::create_dir_all(&load_data_dir).unwrap();
    let start_ts = block_on(pd_client.get_tso()).unwrap().into_inner();
    let commit_ts = block_on(pd_client.get_tso()).unwrap().into_inner();
    let load_data_ctx = LoadDataContext {
        dir: load_data_dir,
        dfs,
        pd: pd_client,
        runtime: runtime.clone(),
        master_key,
        encryption_key_manager,
    };

    let (scheduler, worker_handle) =
        init_task(load_data_config, load_data_ctx, start_ts, commit_ts);

    // Put chunks.
    let keyspace_prefix = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID);
    let i_to_key = move |i: usize| -> Vec<u8> { i_to_row_key(&keyspace_prefix, TABLE_ID, i) };
    let i_to_row_id = move |i: usize| -> Vec<u8> { i.to_be_bytes().to_vec() };
    let ref_store = put_chunks(
        &scheduler,
        WRITER_COUNT,
        DATA_COUNT,
        DATA_BATCH_SIZE,
        i_to_key,
        i_to_val,
        i_to_row_id,
        DEFAULT_TIMEOUT,
        |_| 0,
    );

    fail::cfg(low_space_fp, "return").unwrap();

    // Build. Must be timeout due to low space.
    let res = build(&scheduler, COMPRESSION_TYPE, Duration::from_secs(5));
    assert_matches!(res.unwrap_err(), LoadDataError::Timeout(_));

    // Remove fp and retry again.
    fail::remove(low_space_fp);
    try_wait_finished(&scheduler, DEFAULT_TIMEOUT).unwrap();

    // Cleanup.
    cleanup(&scheduler, worker_handle);

    // Verify data consistency.
    let verified_count = client
        .verify_data_with_given_ref_store(&ref_store, None, &RequestOptions::default())
        .expect("verify_data_with_given_ref_store");
    assert_eq!(verified_count, (DATA_COUNT, 0));

    cluster.stop();
}

#[test]
fn test_ingest_files_returns_not_leader_when_follower_store_stats_not_ready() {
    let _fp = super::FailpointGuard::new().cfg("mock_zero_follower_store_stats", "return");

    let (_temp_dir, _oss, dfs_conf) = prepare_dfs("test_zero_follower_store_stats");
    let node_ids = alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(node_ids, |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_conf.clone();
    });
    cluster.wait_region_replicated(&[], 3);

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

    let keyspace_prefix = ApiV2::get_txn_keyspace_prefix(KEYSPACE_ID);
    let target_key = i_to_row_key(&keyspace_prefix, TABLE_ID, 1);
    let region_info = pd_client
        .get_region_info(&encode_bytes(&target_key))
        .unwrap();
    let region = &region_info.region;
    let leader = region_info.leader.as_ref().unwrap();

    let mut cs = ChangeSet::default();
    cs.set_shard_id(region.get_id());
    cs.set_shard_ver(region.get_region_epoch().get_version());
    cs.mut_ingest_files().set_max_ts(1);
    cs.mut_ingest_files()
        .mut_table_creates()
        .push(new_table_create_pb(
            1,
            WRITE_CF_BOTTOM_LEVEL,
            WRITE_CF as i32,
            i_to_row_key(&keyspace_prefix, TABLE_ID, 1),
            i_to_row_key(&keyspace_prefix, TABLE_ID, 2),
            0,
        ));

    let leader_store = pd_client.get_store(leader.get_store_id()).unwrap();
    let req = Request::post(format!(
        "http://{}/ingest_files?cluster_id={}",
        leader_store.get_status_address(),
        pd_client.get_cluster_id().unwrap()
    ))
    .body(Body::from(cs.write_to_bytes().unwrap()))
    .unwrap();
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let resp = rt.block_on(Client::new().request(req)).unwrap();
    assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);

    let mut errpb = kvproto::errorpb::Error::default();
    errpb
        .merge_from_bytes(
            &rt.block_on(hyper::body::to_bytes(resp.into_body()))
                .unwrap(),
        )
        .unwrap();
    assert!(
        errpb.has_not_leader(),
        "expected not_leader, got: {:?}",
        errpb
    );
    assert!(
        !errpb.has_disk_full(),
        "unexpected disk_full error: {:?}",
        errpb
    );
    assert!(
        errpb.get_message().contains("stats not ready"),
        "unexpected not_leader message: {:?}",
        errpb
    );

    cluster.stop();
}