// Copyright 2023 PingCAP, Inc.

use std::time::Duration;

// Copyright 2025 TiKV Project Authors. Licensed under Apache-2.0.
use health_controller::types::{InspectFactor, LatencyInspector};
use rfstore::store::StoreMsg;
use test_cloud_server::{ServerCluster, alloc_node_id};
use tikv_util::{
    config::{ReadableDuration, ReadableSize},
    info,
    time::Instant,
};

#[test]
#[ignore = "It is mainly used for manual inspection"]
fn test_rfstore_latency() {
    let nodes = vec![1, 2, 3];
    let cluster = ServerCluster::new(nodes, |_, config| {
        config.server.grpc_keepalive_time = ReadableDuration::secs(5);
        config.raft_store.raft_base_tick_interval = ReadableDuration::secs(1);
        config.raft_store.raft_election_timeout_ticks = 10;
        config.raft_store.raft_store_max_leader_lease = ReadableDuration::secs(9);
        config.raft_store.split_region_check_tick_interval = ReadableDuration::secs(10);
        config.raft_store.raft_log_gc_tick_interval = ReadableDuration::secs(3);
        config.raft_store.pd_heartbeat_tick_interval = ReadableDuration::secs(60);
        config.raft_store.pd_store_heartbeat_tick_interval = ReadableDuration::secs(10);
        config.rocksdb.writecf.write_buffer_size = ReadableSize::kb(512);
        config.rocksdb.writecf.block_size = ReadableSize::kb(16);
        config.rocksdb.writecf.target_file_size_base = ReadableSize::kb(512);
        config.rfengine.target_file_size = ReadableSize::mb(16);
    });
    cluster.wait_region_replicated(&[], 3);
    cluster.get_pd_client().disable_default_operator();
    let mut client = cluster.new_client();
    let mut keyspace_id = 1;
    for _ in 0..10 {
        info!("spit keyspace {}", keyspace_id);
        let keyspaces = keyspace_id..(keyspace_id + 1000);
        client.split_keyspaces(keyspaces);
        keyspace_id += 1000;
    }
    std::thread::sleep(Duration::from_secs(5));
    for i in 0..10000 {
        let start = Instant::now();
        client.put_kv(i..i + 1, i_to_key, i_to_key);
        let latency = start.saturating_elapsed();
        if latency > Duration::from_millis(30) {
            info!("latency: {:?} at {}", start.saturating_elapsed(), i);
        }
    }
}

fn i_to_key(i: usize) -> Vec<u8> {
    let keyspace = i % 100 + 1;
    let mut key = api_version::ApiV2::get_txn_keyspace_prefix(keyspace as u32);
    key.extend_from_slice(format!("key_{:05}", i).as_bytes());
    key
}

#[test]
fn test_latency_inspect() {
    let node_id = alloc_node_id();
    let mut cluster = ServerCluster::new(vec![node_id], |_, config| {
        config.raft_store.store_io_pool_size = 2;
        config.raft_store.inspect_kvdb_interval = ReadableDuration::millis(500);
    });
    cluster.wait_region_replicated(&[], 1);

    let (tx, rx) = std::sync::mpsc::sync_channel(10);
    // Inspect different disk_factors.
    for factor in [
        InspectFactor::RaftDisk,
        InspectFactor::KvDisk,
        InspectFactor::Network,
    ] {
        let cloned_tx = tx.clone();
        let inspector = LatencyInspector::new(
            1,
            Box::new(move |_, duration| {
                let dur = duration.sum(true);
                cloned_tx.send(dur).unwrap();
            }),
        );
        let msg = StoreMsg::LatencyInspect {
            factor,
            send_time: Instant::now(),
            ttl: std::time::Duration::from_millis(500),
            inspector,
        };
        match factor {
            InspectFactor::RaftDisk => {
                cluster.send_store_msg(node_id, msg);
                rx.recv_timeout(std::time::Duration::from_secs(2)).unwrap();
            }
            InspectFactor::KvDisk | InspectFactor::Network => {
                // Since kvengine and rfengine use the same mount path in the
                // test environment, KvDisk inspection is disabled. Network
                // inspection is not sent to rfstore in production. Both cases
                // should not be handled by StoreFsm.
                cluster.send_store_msg(node_id, msg);
                rx.recv_timeout(std::time::Duration::from_secs(1))
                    .unwrap_err();
            }
        }
    }
    cluster.stop();
}