// Copyright 2026 TiKV Project Authors. Licensed under Apache-2.0.

use std::sync::Arc;

use api_version::ApiV2;
use bytes::Bytes;
use cloud_encryption::{EncryptionKeyManager, MasterKey};
use kvengine::{
    GLOBAL_SHARD_END_KEY,
    dfs::{DFSConfig, new_dfs_from_config},
};
use kvenginepb::{Snapshot, TxnFileRef};
use kvproto::{
    metapb,
    raft_cmdpb::RaftCmdRequest,
    raft_serverpb::{PeerState, RegionLocalState},
};
use protobuf::Message;
use raft_proto::eraftpb;
use rfengine::{RfEngine, WriteBatch, raft_state_key};
use rfstore::store::{
    CustomBuilder, ProposalContext, RAFT_INIT_LOG_INDEX, state::RaftState, write_engine_meta,
    write_peer_state,
};
use security::{GetSecurityManager, SecurityConfig};
use tempfile::TempDir;
use tikv_util::codec::bytes::encode_bytes;
use txn_types::{LockType, TimeStamp};

use crate::{
    MergedEngine, MergedEngineConfig, MergedEngineContext, MergedEngineInitOpts, manifest::Manifest,
};

// ---------------------------------------------------------------------------
// Low-level entry / lock helpers
// ---------------------------------------------------------------------------

pub(crate) fn make_lock_prefix(primary: Vec<u8>, start_ts: u64) -> Vec<u8> {
    let mut lock = txn_types::Lock::new(
        LockType::Put,
        primary,
        start_ts.into(),
        3_000,
        None,
        0.into(),
        100,
        (start_ts + 1).into(),
    );
    lock.is_txn_file = true;
    lock.to_bytes()
}

pub(crate) fn build_preprocess_entry(cmd: &RaftCmdRequest) -> eraftpb::Entry {
    let mut entry = eraftpb::Entry::default();
    entry.set_entry_type(eraftpb::EntryType::EntryNormal);
    entry.set_context(ProposalContext::PRE_PROCESS.to_vec().into());
    entry.set_data(cmd.write_to_bytes().unwrap().into());
    entry
}

pub(crate) fn build_txn_file_ref_entry(
    log_index: u64,
    shard_ver: u64,
    start_ts: u64,
    lock_val_prefix: Vec<u8>,
    user_meta: Vec<u8>,
) -> eraftpb::Entry {
    let mut txn_file_ref = TxnFileRef::new();
    txn_file_ref.set_start_ts(start_ts);
    txn_file_ref.set_shard_ver(shard_ver);
    txn_file_ref.set_chunk_ids(vec![log_index]);
    txn_file_ref.set_inner_lower_bound(vec![]);
    txn_file_ref.set_inner_upper_bound(GLOBAL_SHARD_END_KEY.to_vec());
    if !lock_val_prefix.is_empty() {
        txn_file_ref.set_lock_val_prefix(lock_val_prefix);
    }
    if !user_meta.is_empty() {
        txn_file_ref.set_user_meta(user_meta);
    }

    let mut builder = CustomBuilder::new();
    builder.set_txn_file(&txn_file_ref);
    let mut cmd = RaftCmdRequest::default();
    cmd.set_custom_request(builder.build());

    let mut entry = build_preprocess_entry(&cmd);
    entry.set_index(log_index);
    entry.set_term(1);
    entry
}

pub(crate) fn encode_raft_state(raft_state: &RaftState) -> Vec<u8> {
    let mut buf = Vec::with_capacity(40);
    let hs = raft_state.get_hard_state();
    buf.extend_from_slice(&raft_state.get_term().to_le_bytes());
    buf.extend_from_slice(&hs.get_vote().to_le_bytes());
    buf.extend_from_slice(&raft_state.get_commit().to_le_bytes());
    buf.extend_from_slice(&raft_state.get_last_index().to_le_bytes());
    buf.extend_from_slice(&raft_state.get_last_preprocessed_index().to_le_bytes());
    buf
}

// ---------------------------------------------------------------------------
// Region / shard fixture builders
// ---------------------------------------------------------------------------

/// Describes a region to seed into a test engine.
#[derive(Clone, Copy)]
pub(crate) struct RegionSpec {
    pub region_id: u64,
    pub region_ver: u64,
    pub keyspace_id: u32,
    pub state: PeerState,
    /// If set, the shard meta will carry a parent shard with this region id.
    pub parent_id: Option<u64>,
}

/// Per-region data to write into the raft engine before startup.
#[derive(Clone, Default)]
pub(crate) struct RegionSeed {
    pub region_id: u64,
    /// Override `shard_meta.seq` / `data_sequence`.
    pub shard_meta_sequence: Option<u64>,
    /// Override `shard_meta.parent.seq` / `data_sequence`.
    pub parent_shard_meta_sequence: Option<u64>,
    pub raft_state: Option<RaftState>,
    pub raft_entries: Vec<eraftpb::Entry>,
}

/// Handle returned by [`open_test_engine_with_seed`].
pub(crate) struct TestEngineEnv {
    pub _dir: TempDir,
    pub engine: MergedEngine,
}

/// Minimal PD client stub for tests.
pub(crate) struct DummyPdClient;

impl GetSecurityManager for DummyPdClient {}

impl crate::PdClient for DummyPdClient {}

pub(crate) fn region_outer_start(keyspace_id: u32, region_id: u64) -> Vec<u8> {
    let mut start = ApiV2::get_txn_keyspace_prefix(keyspace_id);
    start.extend_from_slice(&region_id.to_be_bytes());
    start
}

pub(crate) fn region_outer_end(keyspace_id: u32, region_id: u64) -> Vec<u8> {
    let mut end = ApiV2::get_txn_keyspace_prefix(keyspace_id);
    end.extend_from_slice(&region_id.saturating_add(1).to_be_bytes());
    end
}

pub(crate) fn new_region_state(store_id: u64, spec: RegionSpec) -> RegionLocalState {
    let mut region = metapb::Region::default();
    region.set_id(spec.region_id);
    region.set_start_key(encode_bytes(&region_outer_start(
        spec.keyspace_id,
        spec.region_id,
    )));
    region.set_end_key(encode_bytes(&region_outer_end(
        spec.keyspace_id,
        spec.region_id,
    )));
    region.mut_region_epoch().set_version(spec.region_ver);
    region.mut_region_epoch().set_conf_ver(1);
    let mut peer = metapb::Peer::default();
    peer.set_id(spec.region_id);
    peer.set_store_id(store_id);
    region.set_peers(vec![peer].into());

    let mut state = RegionLocalState::default();
    state.set_region(region);
    state.set_state(spec.state);
    state
}

pub(crate) fn new_shard_meta(
    store_id: u64,
    region_id: u64,
    region_ver: u64,
    keyspace_id: u32,
) -> kvengine::ShardMeta {
    let mut snap = Snapshot::default();
    snap.set_outer_start(region_outer_start(keyspace_id, region_id));
    snap.set_outer_end(region_outer_end(keyspace_id, region_id));

    let mut cs = kvenginepb::ChangeSet::default();
    cs.set_shard_id(region_id);
    cs.set_shard_ver(region_ver);
    cs.set_sequence(RAFT_INIT_LOG_INDEX);
    cs.set_snapshot(snap);
    let mut meta = kvengine::ShardMeta::new(store_id, &cs);
    meta.seq = RAFT_INIT_LOG_INDEX;
    meta.data_sequence = RAFT_INIT_LOG_INDEX;
    meta
}

pub(crate) fn new_raft_state_with_progress(
    peer_id: u64,
    commit: u64,
    last_index: u64,
    last_preprocessed_index: u64,
) -> RaftState {
    let mut raft_state = RaftState::default();
    let mut hs = eraftpb::HardState::default();
    hs.set_term(7);
    hs.set_vote(peer_id);
    hs.set_commit(commit);
    raft_state.set_hard_state(&hs);
    raft_state.set_last_index(last_index);
    raft_state.set_last_preprocessed_index(last_preprocessed_index);
    raft_state
}

// ---------------------------------------------------------------------------
// Full engine fixture
// ---------------------------------------------------------------------------

/// Open a [`MergedEngine`] seeded with the given region specs and data.
///
/// - `store_id`: the store id to use for the engine and all region peers.
/// - `dir_prefix`: tempdir name prefix (for easier debugging).
/// - `specs`: regions to create.
/// - `dependents`: explicit `(parent_id, dependent_id)` pairs to register in
///   the raft engine (in addition to the parent links already encoded in
///   `specs`).
/// - `keyspaces`: keyspace ids to register in the manifest.
/// - `seeds`: per-region data overrides.
pub(crate) fn open_test_engine_with_seed(
    store_id: u64,
    dir_prefix: &str,
    specs: &[RegionSpec],
    dependents: &[(u64, u64)],
    keyspaces: &[u32],
    seeds: &[RegionSeed],
) -> TestEngineEnv {
    let dir = tempfile::Builder::new()
        .prefix(dir_prefix)
        .tempdir()
        .unwrap();
    let merged_dir = dir.path().join("merged");
    let manifest_dir = dir.path().join("manifest");
    let raft = RfEngine::open(
        merged_dir.as_path(),
        &rfengine::RfEngineConfig::default(),
        None,
        None,
    )
    .unwrap();
    raft.set_engine_id(store_id);

    for spec in specs {
        let region_state = new_region_state(store_id, *spec);
        let mut shard_meta =
            new_shard_meta(store_id, spec.region_id, spec.region_ver, spec.keyspace_id);
        if let Some(parent_id) = spec.parent_id {
            shard_meta.parent = Some(Box::new(new_shard_meta(
                store_id,
                parent_id,
                1,
                spec.keyspace_id,
            )));
        }
        let seed = seeds.iter().find(|s| s.region_id == spec.region_id);
        if let Some(sequence) = seed.and_then(|s| s.shard_meta_sequence) {
            shard_meta.seq = sequence;
            shard_meta.data_sequence = sequence;
        }
        if let Some(sequence) = seed.and_then(|s| s.parent_shard_meta_sequence) {
            if let Some(parent) = shard_meta.parent.as_mut() {
                parent.seq = sequence;
                parent.data_sequence = sequence;
            }
        }
        let raft_state = seed.and_then(|s| s.raft_state).unwrap_or_else(|| {
            new_raft_state_with_progress(
                spec.region_id,
                RAFT_INIT_LOG_INDEX,
                RAFT_INIT_LOG_INDEX,
                RAFT_INIT_LOG_INDEX,
            )
        });

        let mut wb = WriteBatch::new();
        write_peer_state(
            &mut wb,
            spec.region_id,
            region_state.get_region(),
            spec.state,
            None,
        );
        write_engine_meta(&mut wb, spec.region_id, &shard_meta);
        wb.set_state(
            spec.region_id,
            spec.region_id,
            spec.keyspace_id,
            &raft_state_key(spec.region_ver),
            &encode_raft_state(&raft_state),
        );
        if let Some(seed) = seed {
            for entry in &seed.raft_entries {
                wb.append_raft_log(spec.region_id, spec.region_id, spec.keyspace_id, entry);
            }
        }
        raft.write(wb).unwrap();
    }
    for &(parent_id, dependent_id) in dependents {
        raft.add_dependent(parent_id, dependent_id);
    }
    drop(raft);

    let mut manifest = Manifest::open(&manifest_dir).unwrap();
    manifest.update_store_progress(1, 1, 0);
    manifest.update_synced_target_ts(TimeStamp::new(1));
    for &ks in keyspaces {
        manifest.set_keyspace_states(ks, Bytes::new());
    }
    manifest.persist().unwrap();

    let ctx = MergedEngineContext {
        pd: Arc::new(DummyPdClient),
        fs: new_dfs_from_config(DFSConfig::default()),
        local_dir: dir.path().to_path_buf(),
        master_key: MasterKey::new(&[1u8; 32]),
        config: MergedEngineConfig {
            merged_store_id: store_id,
            force_ia: false,
            ..Default::default()
        },
        security_config: Arc::new(SecurityConfig::default()),
        encryption_key_manager: Arc::new(EncryptionKeyManager::new()),
        force_stop: Default::default(),
    };
    let engine = MergedEngine::new(ctx, None, MergedEngineInitOpts::default()).unwrap();
    TestEngineEnv { _dir: dir, engine }
}