// Copyright 2026 PingCAP, Inc.

use std::{
    env,
    ops::Deref,
    path::Path,
    sync::{Arc, atomic::AtomicU64},
    thread,
    time::Duration,
    vec,
};

use api_version::{ApiV2, api_v2::KEYSPACE_PREFIX_LEN};
use async_trait::async_trait;
use bytes::Bytes;
#[cfg(test)]
use cloud_encryption::KmsVendor;
use cloud_encryption::MasterKey;
use file_system::IoRateLimiter;
use kvenginepb as pb;
use security::SecurityManager;
use tempfile::TempDir;
use tikv_util::mpsc;
use util::test_util::KeyBuilder;

use crate::{
    dfs::InMemFs,
    ia::manager::{IaManager, IaManagerOptions},
    limiter::StoreLimiter,
    prepare::PrepareOpts,
    shard::{ShardCfBuilder, ShardDataBuilder},
    table::{
        BIT_DELETE, ChecksumType,
        file::InMemFile,
        sstable::{BlockCache, SsTable},
    },
    *,
};

macro_rules! unwrap_or_return {
    ($e:expr, $m:expr) => {
        match $e {
            Ok(x) => x,
            Err(y) => {
                error!("{:?} {:?}", y, $m);
                return;
            }
        }
    };
}

pub const KEYSPACE_ID: u32 = 1;

pub const DEF_BLOCK_SIZE: usize = 4 << 10;
pub const DEF_MIN_BLOB_SIZE: u32 = 64;

pub const TABLE_KEY_PREFIX: &str = "t_";

/// Wrap `Engine` to make sure that it will be closed after the test, and not
/// interfere with other tests.
pub struct TestEngine {
    pub engine: Engine,
    pub key_builder: KeyBuilder,
}

impl std::ops::Deref for TestEngine {
    type Target = Engine;

    fn deref(&self) -> &Self::Target {
        &self.engine
    }
}

impl Drop for TestEngine {
    fn drop(&mut self) {
        self.engine.close();
    }
}

impl TestEngine {
    pub fn key_builder(&self) -> &KeyBuilder {
        &self.key_builder
    }

    pub fn new_ia_manager(&self, opts: IaManagerOptions) -> IaManager {
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .thread_name("ia-mgr")
            .enable_all()
            .worker_threads(2)
            .build()
            .unwrap();
        IaManager::new(opts, self.fs.clone(), None, runtime.into()).unwrap()
    }
}

pub fn new_test_engine() -> (TestEngine, mpsc::Sender<ApplyTask>) {
    new_test_engine_opt(false, DEF_BLOCK_SIZE, "")
}

#[cfg(test)]
pub(crate) fn new_test_engine_with_master_key(
    master_key: MasterKey,
) -> (TestEngine, mpsc::Sender<ApplyTask>) {
    let tester = EngineTester::new(false, DEF_BLOCK_SIZE);
    build_test_engine_from_tester(tester, false, "", None, master_key)
}

pub fn new_test_engine_api_v2() -> (TestEngine, mpsc::Sender<ApplyTask>) {
    new_test_engine_opt(true, DEF_BLOCK_SIZE, "")
}

pub fn new_test_engine_default_options(
    enable_inner_key_off: bool,
) -> (TestEngine, mpsc::Sender<ApplyTask>) {
    let tester = EngineTester::new_with_default_options(enable_inner_key_off);
    build_test_engine_from_tester(
        tester,
        enable_inner_key_off,
        "",
        None,
        MasterKey::new(&[1u8; 32]),
    )
}

pub fn new_test_engine_opt(
    enable_inner_key_off: bool,
    block_size: usize,
    key_prefix: &str,
) -> (TestEngine, mpsc::Sender<ApplyTask>) {
    let tester = EngineTester::new(enable_inner_key_off, block_size);
    build_test_engine_from_tester(
        tester,
        enable_inner_key_off,
        key_prefix,
        None,
        MasterKey::new(&[1u8; 32]),
    )
}

#[cfg(test)]
pub(crate) fn new_test_engine_opt_with_custom_options(
    enable_inner_key_off: bool,
    block_size: usize,
    key_prefix: &str,
    customize: impl FnOnce(&mut Options),
) -> (TestEngine, mpsc::Sender<ApplyTask>) {
    let tester = EngineTester::new_with_custom_options(enable_inner_key_off, block_size, customize);
    build_test_engine_from_tester(
        tester,
        enable_inner_key_off,
        key_prefix,
        None,
        MasterKey::new(&[1u8; 32]),
    )
}

#[cfg(test)]
pub(crate) fn new_test_engine_opt_with_legacy_gc_safe_point_cache(
    enable_inner_key_off: bool,
    block_size: usize,
    key_prefix: &str,
    legacy_gc_safe_point_cache: Arc<dashmap::DashMap<u32, u64>>,
    customize: impl FnOnce(&mut Options),
) -> (TestEngine, mpsc::Sender<ApplyTask>) {
    let tester = EngineTester::new_with_custom_options(enable_inner_key_off, block_size, customize);
    build_test_engine_from_tester(
        tester,
        enable_inner_key_off,
        key_prefix,
        Some(legacy_gc_safe_point_cache),
        MasterKey::new(&[1u8; 32]),
    )
}

/// Replaces shard 1's LOCK_CF with one deterministic SST table.
///
/// Each `(key_index, value)` entry maps to
/// `key_builder.i_to_inner_key(key_index)`. Empty values are encoded as
/// internal tombstones.
pub fn install_lock_cf_test_table(
    engine: &TestEngine,
    table_id: u64,
    entries: &[(usize, Vec<u8>)],
) {
    let shard = engine.get_shard(1).unwrap();
    let block_size = engine.opts.table_builder_options.block_size;
    let comp_tp = engine.opts.table_builder_options.compression_tps[0];
    let comp_lvl = engine.opts.table_builder_options.compression_lvl;
    let mut table_builder = table::sstable::builder::Builder::new(
        table_id,
        block_size,
        comp_tp,
        comp_lvl,
        ChecksumType::Crc32,
        None,
    );
    for (key_index, value) in entries {
        let key = engine.key_builder.i_to_inner_key(*key_index);
        let value = if value.is_empty() {
            table::Value::new_with_meta_version(BIT_DELETE, 1, 0, &[])
        } else {
            table::Value::new_with_meta_version(0, 1, 0, value)
        };
        table_builder.add(key.as_ref(), &value, None).unwrap();
    }
    let mut data_buf = Vec::new();
    table_builder.finish(0, &mut data_buf).unwrap();
    let data = Bytes::from(data_buf);
    let file = InMemFile::new(table_id, data);
    let table = SsTable::new(
        Arc::new(file),
        BlockCache::None,
        cloud_encryption::EncryptionState::Disabled,
    )
    .unwrap();

    let mut cf_builder = ShardCfBuilder::new(LOCK_CF);
    cf_builder.add_table(table, 1);
    let mut cfs = shard.get_data().cfs.clone();
    cfs[LOCK_CF] = cf_builder.build();
    let mut data_builder = ShardDataBuilder::new(shard.get_data());
    data_builder.set_cfs(cfs);
    shard.set_data(data_builder.build());
}

fn build_test_engine_from_tester(
    tester: EngineTester,
    enable_inner_key_off: bool,
    key_prefix: &str,
    legacy_gc_safe_point_cache: Option<Arc<dashmap::DashMap<u32, u64>>>,
    master_key: MasterKey,
) -> (TestEngine, mpsc::Sender<ApplyTask>) {
    // Shared test constructor for the different EngineTester/master-key entry
    // points so they all reuse the same Engine::open and listener wiring.
    let (listener_tx, listener_rx) = mpsc::unbounded();
    let meta_change_listener = Box::new(TestMetaChangeListener {
        sender: listener_tx,
    });
    let rate_limiter = Arc::new(IoRateLimiter::new_for_test());
    let store_limiter = Arc::new(StoreLimiter::dummy());
    let mut meta_iter = tester.clone();
    let engine = Engine::open(
        tester.fs.clone(),
        tester.opts.clone(),
        tester.config.clone(),
        &mut meta_iter,
        tester.clone(),
        tester.core.clone(),
        meta_change_listener,
        rate_limiter,
        store_limiter,
        legacy_gc_safe_point_cache,
        master_key,
        Arc::new(SecurityManager::default()),
        true,
    )
    .unwrap();
    {
        let shard = engine.get_shard(1).unwrap();
        store_bool(&shard.active, true);
    }
    let (applier_tx, applier_rx) = mpsc::unbounded();
    let meta_listener = MetaListener::new(listener_rx, applier_tx.clone());
    thread::spawn(move || {
        meta_listener.run();
    });
    let applier = Applier::new(engine.clone(), applier_rx);
    thread::spawn(move || {
        applier.run();
    });
    let keyspace_id = if enable_inner_key_off { 1 } else { 0 };
    let key_builder = KeyBuilder::new(keyspace_id, key_prefix);
    (
        TestEngine {
            engine,
            key_builder,
        },
        applier_tx,
    )
}

#[derive(Clone)]
struct TestMetaChangeListener {
    sender: mpsc::Sender<pb::ChangeSet>,
}

impl MetaChangeListener for TestMetaChangeListener {
    fn on_change_set(&self, cs: pb::ChangeSet) {
        info!("on meta change listener");
        self.sender.send(cs).unwrap();
    }
}

#[derive(Clone)]
struct EngineTester {
    core: Arc<EngineTesterCore>,
}

impl Deref for EngineTester {
    type Target = EngineTesterCore;
    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

impl EngineTester {
    fn new(enable_inner_key_off: bool, block_size: usize) -> Self {
        Self::new_with_custom_options(enable_inner_key_off, block_size, |_| {})
    }

    fn new_with_default_options(enable_inner_key_off: bool) -> Self {
        let tmp_dir = TempDir::new().unwrap();
        let mut opts = Options::default();
        opts.local_dirs = vec![tmp_dir.path().to_path_buf()];
        Self::new_with_options(enable_inner_key_off, tmp_dir, opts)
    }

    fn new_with_custom_options(
        enable_inner_key_off: bool,
        block_size: usize,
        customize: impl FnOnce(&mut Options),
    ) -> Self {
        let tmp_dir = TempDir::new().unwrap();
        let mut opts = new_test_options(tmp_dir.path(), enable_inner_key_off, block_size);
        customize(&mut opts);
        Self::new_with_options(enable_inner_key_off, tmp_dir, opts)
    }

    fn new_with_options(enable_inner_key_off: bool, tmp_dir: TempDir, opts: Options) -> Self {
        let initial_cs = new_initial_cs(enable_inner_key_off);
        let initial_meta = ShardMeta::new(1, &initial_cs);
        let metas = dashmap::DashMap::new();
        metas.insert(1, Arc::new(initial_meta));
        let config = KvEngineConfig::default();

        Self {
            core: Arc::new(EngineTesterCore {
                _tmp_dir: tmp_dir,
                metas,
                fs: Arc::new(InMemFs::new()),
                opts: Arc::new(opts),
                config,
                id: AtomicU64::new(0),
            }),
        }
    }
}

struct EngineTesterCore {
    _tmp_dir: TempDir,
    metas: dashmap::DashMap<u64, Arc<ShardMeta>>,
    fs: Arc<dfs::InMemFs>,
    opts: Arc<Options>,
    config: KvEngineConfig,
    id: AtomicU64,
}

impl MetaIterator for EngineTester {
    fn iterate<F>(&mut self, mut f: F) -> Result<()>
    where
        F: FnMut(kvenginepb::ChangeSet),
    {
        for meta in &self.metas {
            f(meta.value().to_change_set())
        }
        Ok(())
    }

    fn engine_id(&self) -> u64 {
        1
    }
}

impl RecoverHandler for EngineTester {
    fn recover(
        &self,
        _engine: &Engine,
        _shard: &Arc<Shard>,
        _info: &ShardMeta,
        _is_parent: bool,
    ) -> Result<()> {
        Ok(())
    }
}

#[async_trait]
impl IdAllocator for EngineTesterCore {
    fn alloc_id(&self, count: usize) -> Result<Vec<u64>> {
        let start_id = self
            .id
            .fetch_add(count as u64, std::sync::atomic::Ordering::Relaxed)
            + 1;
        let end_id = start_id + count as u64;
        let mut ids = Vec::with_capacity(count);
        for id in start_id..end_id {
            ids.push(id);
        }
        Ok(ids)
    }

    async fn alloc_id_async(&self, count: usize) -> Result<Vec<u64>> {
        self.alloc_id(count)
    }
}

struct MetaListener {
    applier_tx: mpsc::Sender<ApplyTask>,
    meta_rx: mpsc::Receiver<pb::ChangeSet>,
}

impl MetaListener {
    fn new(meta_rx: mpsc::Receiver<pb::ChangeSet>, applier_tx: mpsc::Sender<ApplyTask>) -> Self {
        Self {
            meta_rx,
            applier_tx,
        }
    }

    fn run(&self) {
        loop {
            let cs = unwrap_or_return!(self.meta_rx.recv(), "meta_listener_a");
            let (tx, rx) = mpsc::bounded(1);
            let task = ApplyTask::new_cs(cs, tx);
            self.applier_tx.send(task).unwrap();
            let res = unwrap_or_return!(rx.recv(), "meta_listener_b");
            unwrap_or_return!(res, "meta_listener_c");
        }
    }
}

struct Applier {
    engine: Engine,
    task_rx: mpsc::Receiver<ApplyTask>,
}

impl Applier {
    fn new(engine: Engine, task_rx: mpsc::Receiver<ApplyTask>) -> Self {
        Self { engine, task_rx }
    }

    fn run(&self) {
        let mut seq = 2;
        loop {
            let mut task = unwrap_or_return!(self.task_rx.recv(), "apply recv task");
            seq += 1;
            if let Some(wb) = task.wb.as_mut() {
                wb.set_sequence(seq);
                self.engine.write(wb, &[]);
            }
            if let Some(mut cs) = task.cs.take() {
                cs.set_sequence(seq);
                if cs.has_split() {
                    let mut ids = vec![];
                    for new_shard in cs.get_split().get_new_shards() {
                        ids.push(new_shard.shard_id);
                    }
                    let skip_new_shard_ids = Vec::new();
                    unwrap_or_return!(self.engine.split(cs, 1, &skip_new_shard_ids), "apply split");
                    for id in ids {
                        let shard = self.engine.get_shard(id).unwrap();
                        shard.set_active(true);
                    }
                    info!("applier executed split");
                } else {
                    self.engine.meta_committed(&cs, false);
                    unwrap_or_return!(
                        self.engine.apply_change_set(
                            &self
                                .engine
                                .prepare_change_set(cs, PrepareOpts::default())
                                .unwrap()
                        ),
                        "applier apply changeset"
                    );
                }
            }
            task.result_tx.send(Ok(seq)).unwrap();
        }
    }
}

pub struct ApplyTask {
    wb: Option<WriteBatch>,
    cs: Option<pb::ChangeSet>,
    result_tx: mpsc::Sender<Result<u64 /* write_sequence */>>,
}

impl ApplyTask {
    /// Creates a task to apply a change set.
    pub(crate) fn new_cs(cs: pb::ChangeSet, result_tx: mpsc::Sender<Result<u64>>) -> Self {
        Self {
            wb: None,
            cs: Some(cs),
            result_tx,
        }
    }

    pub fn new_wb(wb: WriteBatch, result_tx: mpsc::Sender<Result<u64>>) -> Self {
        Self {
            wb: Some(wb),
            cs: None,
            result_tx,
        }
    }
}

pub struct Splitter {
    apply_sender: mpsc::Sender<ApplyTask>,
    keys: Vec<Vec<u8>>,
    current_shard_id: u64,
    shard_ver: u64,
    new_id: u64,
}

#[allow(dead_code)]
impl Splitter {
    pub fn new(
        keys: Vec<Vec<u8>>,
        current_shard_id_ver: IdVer,
        new_id_base: u64,
        apply_sender: mpsc::Sender<ApplyTask>,
    ) -> Self {
        Self {
            keys,
            apply_sender,
            current_shard_id: current_shard_id_ver.id,
            shard_ver: current_shard_id_ver.ver,
            new_id: new_id_base,
        }
    }

    pub fn run(&mut self) {
        let keys = self.keys.clone();
        for key in keys {
            thread::sleep(Duration::from_millis(200));
            self.new_id += 1;
            self.split(key.clone(), vec![self.new_id, self.current_shard_id]);
        }
    }

    pub fn send_task(&mut self, cs: pb::ChangeSet) {
        let (tx, rx) = mpsc::bounded(1);
        let task = ApplyTask {
            cs: Some(cs),
            wb: None,
            result_tx: tx,
        };
        self.apply_sender.send(task).unwrap();
        let res = unwrap_or_return!(rx.recv(), "splitter recv");
        res.unwrap();
    }

    pub fn split(&mut self, key: Vec<u8>, new_ids: Vec<u64>) {
        let mut cs = pb::ChangeSet::new();
        cs.set_shard_id(self.current_shard_id);
        cs.set_shard_ver(self.shard_ver);
        let mut finish_split = pb::Split::new();
        finish_split.set_keys(protobuf::RepeatedField::from_vec(vec![key]));
        let mut new_shards = Vec::new();
        for new_id in &new_ids {
            let mut new_shard = pb::Properties::new();
            new_shard.set_shard_id(*new_id);
            new_shards.push(new_shard);
        }
        finish_split.set_new_shards(protobuf::RepeatedField::from_vec(new_shards));
        cs.set_split(finish_split);
        self.send_task(cs);
        self.shard_ver += 1;
    }
}

pub fn new_initial_cs(enable_inner_key_off: bool) -> pb::ChangeSet {
    let mut cs = pb::ChangeSet::new();
    cs.set_shard_id(1);
    cs.set_shard_ver(1);
    cs.set_sequence(1);
    let mut snap = pb::Snapshot::new();
    snap.set_base_version(1);
    if enable_inner_key_off {
        let (start, end) = ApiV2::get_txn_keyspace_range(KEYSPACE_ID);
        snap.set_outer_start(start);
        snap.set_outer_end(end);
        snap.set_inner_key_off(KEYSPACE_PREFIX_LEN as u32);
    } else {
        snap.set_outer_end(GLOBAL_SHARD_END_KEY.to_vec());
    }
    let props = snap.mut_properties();
    props.shard_id = 1;
    cs.set_snapshot(snap);
    cs
}

pub fn new_test_options(
    path: impl AsRef<Path>,
    _enable_inner_key_off: bool,
    block_size: usize,
) -> Options {
    let min_blob_size: u32 = match env::var("MIN_BLOB_SIZE") {
        Ok(val) => match val.trim().parse() {
            Ok(n) => n,
            Err(e) => {
                warn!("MIN_BLOB_SIZE=<number>, got {}", e);
                DEF_MIN_BLOB_SIZE
            }
        },
        Err(_) => DEF_MIN_BLOB_SIZE,
    };
    info!("MIN_BLOB_SIZE={}", min_blob_size);
    let mut opts = Options::default();
    opts.local_dirs = vec![path.as_ref().to_path_buf()];
    opts.base_size = 64 << 10;
    opts.table_builder_options.block_size = block_size;
    opts.table_builder_options.max_table_size = 8 << 10;
    opts.table_builder_options.flush_split_l0 = true;
    opts.columnar_build_options.max_columnar_table_size = 1024;
    opts.columnar_build_options.pack_max_row_count = 9;
    opts.set_build_fts_index(true);
    opts.max_mem_table_size = 32 << 10; // mem-table size should be much larger than max_table_size.
    opts.num_compactors = 2;
    opts.blob_table_build_options.min_blob_size = min_blob_size;
    opts.max_del_range_delay = Duration::from_secs(1);
    opts.read_columnar = true;
    opts
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_engine_bootstraps_legacy_mode_from_global_master_key_vendor() {
        let master_key = MasterKey::new_with_kms_vendor(&[9u8; 32], KmsVendor::Aws);
        let (engine, _) = new_test_engine_with_master_key(master_key);
        assert_eq!(
            engine.get_encryption_key_manager().current_mode(),
            Some(cloud_encryption::EncryptionMethod::Legacy)
        );
    }
}