// Copyright 2026 PingCAP, Inc.

mod test_pd_worker;
mod test_peer_fsm;
mod test_startup_batch;

use std::{
    collections::{HashMap, HashSet},
    sync::{Arc, Mutex, atomic::AtomicU64},
};

use api_version::ApiV2;
use cloud_encryption::MasterKey;
use file_system::IoRateLimiter;
use health_controller::types::InspectFactor;
use kvengine::limiter::StoreLimiter;
use kvproto::{
    kvrpcpb::ApiVersion,
    metapb::{self},
    raft_serverpb::RaftMessage,
};
use raftstore::coprocessor::CoprocessorHost;
use rfengine::{RfEngine, RfEngineConfig};
use security::SecurityManager;
use sst_importer::SstImporter;
use tempfile::TempDir;
use test_pd_client::TestPdClient;
use tikv_util::{
    codec::bytes::encode_bytes, config::VersionTrack, store::new_peer, worker::LazyWorker,
};

use crate::{
    RaftRouter,
    store::{Config, initial_region, prepare_bootstrap_cluster, *},
};

// Creates new engines (`kvengine` and `rf_engine`) for use in tests.
pub(crate) fn new_test_engines() -> (Engines, TempDir) {
    let region = initial_region(1, 1, 1);
    new_test_engines_with_region(&region)
}

fn cold_test_region() -> metapb::Region {
    let mut region = initial_region(1, 1, 1);
    let (start, end) = ApiV2::get_txn_keyspace_range(1);
    region.set_start_key(encode_bytes(&start));
    region.set_end_key(encode_bytes(&end));
    region
}

fn new_test_engines_with_region(region: &metapb::Region) -> (Engines, TempDir) {
    let tmp_dir = ::test_util::temp_dir("test_node_", true);
    let data_dir = tmp_dir.path();

    let mut rfengine_cfg = RfEngineConfig::default();
    rfengine_cfg.wal_sync_dir = format!("{}/wal", data_dir.display());

    let rf_engine = RfEngine::open(
        data_dir.join("rfengine").as_path(),
        &rfengine_cfg,
        Some(data_dir),
        None,
    )
    .unwrap();

    let rate_limiter = Arc::new(IoRateLimiter::new_for_test());
    let store_limiter = Arc::new(StoreLimiter::dummy());

    let (sender, _) = tikv_util::mpsc::unbounded();
    let meta_change_listener = Box::new(MetaChangeListener {
        sender: sender.clone(),
    });

    let dfs = Arc::new(kvengine::dfs::InMemFs::new());
    let recoverer = crate::store::RecoverHandler::new(rf_engine.clone());
    let mut meta_iter = recoverer.clone();
    let pd_client = Arc::new(TestPdClient::new(1, false));
    let id_allocator = Arc::new(PdIdAllocator::new(pd_client.clone()));

    let kv_cfg = kvengine::config::Config::default();
    let mut kv_opts = kvengine::Options::default();
    kv_opts.local_dirs = vec![data_dir.to_path_buf()];

    let kvengine = kvengine::Engine::open(
        dfs,
        Arc::new(kv_opts),
        kv_cfg,
        &mut meta_iter,
        recoverer,
        id_allocator,
        meta_change_listener,
        rate_limiter,
        store_limiter,
        None,
        MasterKey::new(&[1u8; 32]),
        Arc::new(SecurityManager::default()),
        true,
    )
    .unwrap();

    let engines = Engines::new(kvengine, rf_engine, tikv_util::mpsc::unbounded(), None);

    let peer_id = region
        .get_peers()
        .iter()
        .find(|peer| peer.get_store_id() == 1)
        .unwrap()
        .get_id();
    prepare_bootstrap_cluster(&engines, region, peer_id).unwrap();

    (engines, tmp_dir)
}

// A mock transport implementation that captures Raft messages sent by peers.
#[derive(Clone)]
pub struct MockTransport {
    messages: Arc<Mutex<Vec<RaftMessage>>>,
}

impl MockTransport {
    pub fn new() -> Self {
        MockTransport {
            messages: Arc::new(Mutex::new(Vec::new())),
        }
    }
}

impl Transport for MockTransport {
    fn send(&mut self, msg: RaftMessage) -> crate::Result<()> {
        self.messages.lock().unwrap().push(msg);
        Ok(())
    }

    fn need_flush(&self) -> bool {
        false
    }

    fn flush(&mut self) {}
}

// Initializes a `PeerFsm` instance for testing purposes.
fn new_test_peer_fsm(engines: Engines, region: &metapb::Region) -> Result<PeerFsm, crate::Error> {
    let read_worker = tikv_util::worker::Worker::new("test-read-worker");
    let read_scheduler = read_worker.start(
        "test-read-worker",
        crate::store::worker::ReadRunner::new(
            engines.raft.clone(),
            RaftRouter::new(
                crate::raft_router_peer_channel(0).0,
                tikv_util::mpsc::unbounded().0,
            ),
        ),
    );
    PeerFsm::create(1, &Config::default(), engines, region, read_scheduler)
}

// Builds a `RaftContext` for testing.
fn new_test_raft_ctx(
    engines: Engines,
    pd_scheduler: Option<tikv_util::worker::Scheduler<PdTask>>,
) -> RaftContext {
    let rf_store_cfg = Config::default();
    let cfg = Arc::new(VersionTrack::new(rf_store_cfg));

    let tempfile = tempfile::tempdir().unwrap();
    let sst_import_dir = tempfile.path();
    let importer = Arc::new(
        SstImporter::new(
            &sst_importer::Config::default(),
            sst_import_dir,
            None,
            ApiVersion::V2,
        )
        .unwrap(),
    );

    let pd_scheduler =
        pd_scheduler.unwrap_or_else(|| LazyWorker::new("test-pd-worker").scheduler());
    let gc_worker = LazyWorker::new("test-gc-worker");

    let (peer_sender, _) = crate::raft_router_peer_channel(0);
    let (store_sender, _) = tikv_util::mpsc::unbounded();
    let router = RaftRouter::new(peer_sender, store_sender);

    let read_worker = tikv_util::worker::Worker::new("test-read-worker");
    let read_scheduler = read_worker.start(
        "test-read-worker",
        crate::store::worker::ReadRunner::new(engines.raft.clone(), router.clone()),
    );

    let background_worker = tikv_util::worker::Worker::new("test-background-worker");
    let split_check_scheduler = SplitCheckScheduler::start(
        &background_worker,
        "test-split-check",
        engines.kv.clone(),
        router.clone(),
    );

    // Create schedulers for inspecting latency jitters of disk I/O operations.
    // Since under testing env, the mount path for kvengine and rfengine should
    // be same root path, it's no need to inspect latency jitters of kvengine.
    let mut disk_check_schedulers = HashMap::new();
    {
        // Inspect disks of rfengine.
        let (rfdb_dir, rf_wal_dir) = (engines.raft.db_dir(), engines.raft.wal_dir());
        assert!(rfdb_dir.is_some() && rf_wal_dir.is_some());
        let mut disk_check_runner = DiskCheckRunner::new(rfdb_dir.unwrap().to_path_buf());
        disk_check_runner.bind_background_worker(background_worker.clone());
        let disk_check_scheduler =
            background_worker.start("raft-disk-check-worker", disk_check_runner);
        disk_check_schedulers.insert(InspectFactor::RaftDisk as u8, disk_check_scheduler);
    }

    let trans = MockTransport::new();
    let trans_idle = MockTransport::new();
    let schema_worker = LazyWorker::new("test-schema-worker");

    let global_ctx: GlobalContext = GlobalContext {
        cfg,
        engines,
        store: metapb::Store::default(),
        readers: StoreMeta::new(PENDING_MSG_CAP).readers.clone(),
        router,
        cold_store_liveness: Arc::default(),
        trans: Box::new(trans),
        trans_idle: Box::new(trans_idle),
        pd_scheduler,
        gc_scheduler: gc_worker.scheduler(),
        schema_scheduler: schema_worker.scheduler(),
        split_check_scheduler,
        read_scheduler,
        txn_file_lock_resolve_scheduler: TxnFileLockResolveScheduler::dummy(),
        coprocessor_host: CoprocessorHost::default(),
        importer,
        disk_check_schedulers: Arc::new(disk_check_schedulers),
        destroying: HashSet::default(),
        engine_total_bytes_written: Arc::new(AtomicU64::new(0)),
        engine_total_keys_written: Arc::new(AtomicU64::new(0)),
        last_raft_append_success_at_millis: Arc::new(AtomicU64::new(0)),
        busy_apply_status: Arc::new(Mutex::new(BusyApplyStatus {
            busy_apply_peers: HashSet::default(),
            completed_apply_peers_count: Some(0),
        })),
    };
    RaftContext::new(global_ctx, WorkerType::Main)
}