// Copyright 2025 PingCAP, Inc.

use std::sync::Arc;

use bytes::Buf;
use cloud_encryption::{EncryptionState, MasterKey};
use kvengine::{GLOBAL_SHARD_END_KEY, ShardMeta, ShardTag, encryption_state_from_shard_properties};
use kvproto::{metapb, raft_cmdpb::RaftCmdRequest, raft_serverpb::MergeState};
use raft_proto::eraftpb::{Entry, HardState};
use rfengine::{RfEngine, WriteBatch};
use rfstore::{
    errors::Error as RfStoreError,
    store::{
        AffectMemtable, ApplyMsgs, Config, PeerTag, PreprocessContext, PreprocessRef,
        PreprocessScene, RAFT_INIT_LOG_TERM, get_preprocess_cmd, rlog, state::RaftState,
    },
};
use tikv_util::{box_err, codec::bytes::encode_bytes, debug, error, info, warn};

use crate::{Error, RegionProgress, Result};

pub(crate) struct Preprocessor {
    preprocessed_index: u64,
    region: metapb::Region,
    preprocessed_region: Option<metapb::Region>,
    peer: metapb::Peer,
    shard_meta: Option<ShardMeta>,
    last_committed_split_idx: u64,
    pending_truncate: Option<(u64 /* term */, u64 /* index */)>,
    raft_hard_state: HardState,
    raft_state: RaftState,
    pending_merge_state: Option<MergeState>,
    want_rollback_merge_peers: collections::HashSet<u64>,
    learner_skip_idx: u64,
    encryption_state: EncryptionState,
}

impl Preprocessor {
    pub(crate) fn from_shard_meta(
        raft: &rfengine::RfEngine,
        store_id: u64,
        shard_meta: ShardMeta,
        master_key: &MasterKey,
        encryption_key_manager: Arc<cloud_encryption::EncryptionKeyManager>,
    ) -> Self {
        let encryption_state = encryption_state_from_shard_properties(
            &shard_meta.get_properties_pb(),
            encryption_key_manager,
            master_key,
            Some(shard_meta.range.keyspace_id),
        );
        let mut region = metapb::Region::default();
        region.set_id(shard_meta.id);
        let epoch = region.mut_region_epoch();
        epoch.set_version(shard_meta.ver);
        epoch.set_conf_ver(1);
        let mut peer = metapb::Peer::default();
        peer.set_id(shard_meta.id);
        peer.set_store_id(store_id);
        region.set_peers(vec![peer.clone()].into());
        region.set_start_key(encode_bytes(&shard_meta.range.outer_start));
        if shard_meta.range.outer_end.chunk() != GLOBAL_SHARD_END_KEY {
            region.set_end_key(encode_bytes(&shard_meta.range.outer_end));
        }
        let raft_index = shard_meta.seq;
        let mut raft_state = RaftState::default();
        let mut raft_hard_state = HardState::default();
        raft_hard_state.set_vote(shard_meta.id);
        raft_hard_state.set_term(RAFT_INIT_LOG_TERM);
        raft_state.set_last_preprocessed_index(raft_index);
        raft_state.set_hard_state(&raft_hard_state);

        let mut region_state = raft
            .load_region_state(shard_meta.id, shard_meta.ver)
            .unwrap_or_else(|| {
                let states = raft.get_peer_all_states(shard_meta.id, false);
                panic!(
                    "{} failed to load region state, states: {:?}",
                    shard_meta.tag(),
                    states
                );
            });
        let merge_state = region_state
            .has_merge_state()
            .then(|| region_state.take_merge_state());

        Self {
            preprocessed_index: raft_index,
            region,
            preprocessed_region: None,
            peer,
            shard_meta: Some(shard_meta),
            last_committed_split_idx: raft_index,
            pending_truncate: None,
            raft_hard_state,
            raft_state,
            pending_merge_state: merge_state,
            want_rollback_merge_peers: collections::HashSet::default(),
            learner_skip_idx: raft_index,
            encryption_state,
        }
    }

    pub(crate) fn new(
        raft: &rfengine::RfEngine,
        store_id: u64,
        region_id: u64,
        master_key: &MasterKey,
        encryption_key_manager: Arc<cloud_encryption::EncryptionKeyManager>,
    ) -> Option<Self> {
        let shard_meta =
            rfstore::store::load_engine_meta(raft, store_id, region_id).or_else(|| {
                info!(
                    "{}:{} new_preprocessor: no engine meta for region",
                    store_id, region_id
                );
                None
            })?;
        Some(Self::from_shard_meta(
            raft,
            store_id,
            shard_meta,
            master_key,
            encryption_key_manager,
        ))
    }

    pub(crate) fn tag(&self) -> ShardTag {
        ShardTag::from_region(None, &self.region)
    }

    pub(crate) fn region(&self) -> &metapb::Region {
        &self.region
    }

    pub(crate) fn as_ref(&mut self) -> PreprocessRef<'_> {
        PreprocessRef {
            region: &mut self.region,
            preprocessed_index: &mut self.preprocessed_index,
            preprocessed_region: &mut self.preprocessed_region,
            peer: &mut self.peer,
            shard_meta: &mut self.shard_meta,
            last_committed_split_idx: &mut self.last_committed_split_idx,
            pending_truncate: &mut self.pending_truncate,
            raft_hard_state: self.raft_hard_state.clone(),
            raft_state: &mut self.raft_state,
            is_leader: false,
            pending_merge_state: &mut self.pending_merge_state,
            want_rollback_merge_peers: &mut self.want_rollback_merge_peers,
            learner_skip_idx: &mut self.learner_skip_idx,
            encryption_state: &mut self.encryption_state,
        }
    }

    pub(crate) fn sync_region(&mut self) -> bool /* changed */ {
        if let Some(new_region) = self.preprocessed_region.take() {
            let changed = self.region != new_region;
            self.region = new_region;
            return changed;
        }
        false
    }

    pub(crate) fn shard_meta(&self) -> Option<&ShardMeta> {
        self.shard_meta.as_ref()
    }

    pub(crate) fn mut_shard_meta(&mut self) -> Option<&mut ShardMeta> {
        self.shard_meta.as_mut()
    }

    pub(crate) fn pending_merge_state(&self) -> Option<&MergeState> {
        self.pending_merge_state.as_ref()
    }

    // Recover persisted progress & txn file locks.
    pub(crate) fn recover(
        &mut self,
        ctx: &mut PreprocessContext<'_>,
        progress: &mut RegionProgress,
        exclusive_high_idx: u64,
    ) -> Result<()> {
        const FETCH_BATCH: u64 = 1024;

        let mut preprocessor_ref = self.as_ref();
        let Some(shard_meta) = preprocessor_ref.shard_meta.as_ref() else {
            return Ok(());
        };
        let tag = shard_meta.tag();
        let region_id = shard_meta.id;
        let low_idx = shard_meta.data_sequence.saturating_add(1);
        if low_idx >= exclusive_high_idx {
            return Ok(());
        }

        let raft = ctx.raft;
        let base_version = shard_meta.base_version;
        let shard_ver = shard_meta.ver;
        let peer_tag = tag.into();
        let mut entry_buf: Vec<Entry> =
            Vec::with_capacity((exclusive_high_idx - low_idx).min(FETCH_BATCH) as usize);
        let mut txn_file_log_idxes = Vec::new();
        let mut fetch_low = low_idx;
        while fetch_low < exclusive_high_idx {
            let fetch_high = fetch_low
                .saturating_add(FETCH_BATCH)
                .min(exclusive_high_idx);
            entry_buf.clear();
            raft.fetch_raft_entries_to(
                region_id,
                fetch_low,
                fetch_high,
                None,
                &mut entry_buf,
                None,
            )
            .map_err(|err| -> Error {
                let stats = raft.get_peer_stats(region_id);
                let truncated_state = rfstore::store::load_raft_truncated_state(raft, region_id);
                error!(
                    "{} Preprocessor.recover: fetch raft entries failed", tag;
                    "low" => fetch_low,
                    "high" => fetch_high,
                    "stats" => ?stats,
                    "truncated_state" => ?truncated_state,
                    "err" => ?err,
                );
                box_err!(err)
            })?;
            for entry in &entry_buf {
                let log_index = entry.index;
                let (affect_memtable, cmd) =
                    get_affect_memtable_for_recover_entry(entry, peer_tag, base_version, shard_ver);
                progress.persist_progress.update(log_index, affect_memtable);

                if let Some(cmd) = cmd {
                    if cmd.has_custom_request() {
                        let custom_req = cmd.get_custom_request();
                        if !rlog::is_txn_file_ref(custom_req.get_data()) {
                            continue;
                        }
                        if let Err(err) =
                            preprocessor_ref.preprocess_txn_file_ref(ctx, entry, custom_req)
                        {
                            warn!(
                                "{} Preprocessor.recover: preprocess txn file ref failed", tag;
                                "log_idx" => entry.index,
                                "err" => ?err,
                            );
                            if !matches!(err, RfStoreError::EpochNotMatch(..)) {
                                return Err(box_err!(err));
                            }
                        } else {
                            txn_file_log_idxes.push(entry.index);
                        }
                    }
                }
            }
            fetch_low = fetch_high;
        }
        debug!("{} Preprocessor.recover: [{}, {})", tag, low_idx, exclusive_high_idx;
            "progress" => ?progress, "txn_file_log_idxes" => ?txn_file_log_idxes);
        Ok(())
    }
}

fn get_affect_memtable_for_recover_entry(
    entry: &Entry,
    peer_tag: PeerTag,
    base_version: u64,
    shard_ver: u64,
) -> (AffectMemtable, Option<RaftCmdRequest>) {
    if entry.get_entry_type() != raft_proto::eraftpb::EntryType::EntryNormal
        || entry.get_data().is_empty()
    {
        return (AffectMemtable::None, None);
    }
    let Some(cmd) = get_preprocess_cmd(entry) else {
        // Normal writes (no preprocess).
        return (AffectMemtable::Write, None);
    };
    let Some(custom) = rlog::get_custom_log(&cmd) else {
        return (AffectMemtable::None, Some(cmd));
    };
    let affect_memtable = custom.is_affect_memtable(peer_tag, Some(base_version), Some(shard_ver));
    (affect_memtable, Some(cmd))
}

pub(crate) struct RecoverTxnFileLocksContext<'a> {
    store_id: u64,
    raft: &'a RfEngine,
    raft_wb: WriteBatch,
    remove_dependents: Vec<(u64 /* parent_id */, u64 /* dependent_id */)>,
    apply_msgs: ApplyMsgs,
    cfg: Config,
    destroying: std::collections::HashSet<u64>,
}

impl<'a> RecoverTxnFileLocksContext<'a> {
    pub(crate) fn new(store_id: u64, raft: &'a RfEngine) -> Self {
        Self {
            store_id,
            raft,
            raft_wb: WriteBatch::new(),
            remove_dependents: Vec::new(),
            apply_msgs: ApplyMsgs::default(),
            cfg: Config::default(),
            destroying: std::collections::HashSet::default(),
        }
    }

    pub(crate) fn as_ref(&mut self) -> PreprocessContext<'_> {
        PreprocessContext {
            store_id: self.store_id,
            raft: self.raft,
            raft_wb: &mut self.raft_wb,
            remove_dependents: &mut self.remove_dependents,
            apply_msgs: &mut self.apply_msgs,
            cfg: &self.cfg,
            destroying: &mut self.destroying,
            scene: PreprocessScene::MergedEngine { router: None },
        }
    }
}

impl Drop for RecoverTxnFileLocksContext<'_> {
    fn drop(&mut self) {
        debug_assert!(self.raft_wb.is_empty());
        debug_assert!(self.remove_dependents.is_empty());
        debug_assert!(self.apply_msgs.is_empty());
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use bytes::Bytes;
    use cloud_encryption::{EncryptionKeyManager, MasterKey};
    use kvengine::ShardMeta;
    use kvenginepb::{ChangeSet, Snapshot, TxnFileRef};
    use kvproto::{
        metapb,
        raft_cmdpb::{AdminCmdType, AdminRequest, RaftCmdRequest},
        raft_serverpb::{PeerState, RegionLocalState},
    };
    use raft_proto::eraftpb::{Entry, EntryType, HardState};
    use rfengine::{RfEngine, WriteBatch, raft_state_key};
    use rfstore::store::{
        AffectMemtable, CustomBuilder, PeerTag, RAFT_INIT_LOG_INDEX, RegionIdVer, state::RaftState,
        write_engine_meta, write_peer_state,
    };
    use tikv_util::codec::bytes::encode_bytes;

    use super::{Preprocessor, RecoverTxnFileLocksContext, get_affect_memtable_for_recover_entry};
    use crate::{
        RegionProgress,
        test_util::{
            build_preprocess_entry, build_txn_file_ref_entry, encode_raft_state, make_lock_prefix,
        },
    };

    const TEST_STORE_ID: u64 = 1024;
    const TEST_KEYSPACE_ID: u32 = 1;

    fn new_region_state(region_id: u64, region_ver: u64) -> RegionLocalState {
        let mut region = metapb::Region::default();
        region.set_id(region_id);
        region.set_start_key(encode_bytes(b"a"));
        region.set_end_key(encode_bytes(b"z"));
        region.mut_region_epoch().set_version(region_ver);
        region.mut_region_epoch().set_conf_ver(1);
        let mut peer = metapb::Peer::default();
        peer.set_id(region_id);
        peer.set_store_id(TEST_STORE_ID);
        region.set_peers(vec![peer].into());

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

    fn new_shard_meta(region_id: u64, region_ver: u64) -> ShardMeta {
        let mut snap = Snapshot::default();
        snap.set_outer_start(b"a".to_vec());
        snap.set_outer_end(b"z".to_vec());

        let mut cs = 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 = ShardMeta::new(TEST_STORE_ID, &cs);
        meta.seq = RAFT_INIT_LOG_INDEX;
        meta.data_sequence = RAFT_INIT_LOG_INDEX;
        meta
    }

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

    struct RecoverTestEnv {
        _dir: tempfile::TempDir,
        raft: RfEngine,
        preprocessor: Preprocessor,
        progress: RegionProgress,
    }

    fn seed_recover_test_peer(
        region_id: u64,
        region_ver: u64,
        shard_meta: ShardMeta,
        raft_state: RaftState,
        entries: &[Entry],
    ) -> RecoverTestEnv {
        let dir = tempfile::Builder::new()
            .prefix("merged_engine_preprocessor_recover_")
            .tempdir()
            .unwrap();
        let raft =
            RfEngine::open(dir.path(), &rfengine::RfEngineConfig::default(), None, None).unwrap();
        raft.set_engine_id(TEST_STORE_ID);

        let region_state = new_region_state(region_id, region_ver);
        let mut wb = WriteBatch::new();
        write_peer_state(
            &mut wb,
            region_id,
            region_state.get_region(),
            PeerState::Normal,
            None,
        );
        write_engine_meta(&mut wb, region_id, &shard_meta);
        wb.set_state(
            region_id,
            region_id,
            TEST_KEYSPACE_ID,
            &raft_state_key(region_ver),
            &encode_raft_state(&raft_state),
        );
        for entry in entries {
            wb.append_raft_log(region_id, region_id, TEST_KEYSPACE_ID, entry);
        }
        raft.write(wb).unwrap();

        let preprocessor = Preprocessor::from_shard_meta(
            &raft,
            TEST_STORE_ID,
            shard_meta,
            &MasterKey::new(&[1u8; 32]),
            Arc::new(EncryptionKeyManager::new()),
        );

        RecoverTestEnv {
            _dir: dir,
            raft,
            preprocessor,
            progress: RegionProgress::new(TEST_KEYSPACE_ID, region_id),
        }
    }

    #[test]
    fn test_recovery_classification_matches_sync_path() {
        let merged_store_id = 101;
        let peer_tag = PeerTag::new(merged_store_id, RegionIdVer::new(7, 0));
        let base_version = 1000;
        let shard_ver = 7;

        let mut cases = Vec::new();

        let mut conf_change = Entry::default();
        conf_change.set_entry_type(EntryType::EntryConfChange);
        conf_change.set_data(Bytes::from_static(b"conf-change"));
        cases.push(("conf-change", conf_change));

        let mut empty_normal = Entry::default();
        empty_normal.set_entry_type(EntryType::EntryNormal);
        cases.push(("empty-normal", empty_normal));

        let mut plain_write = Entry::default();
        plain_write.set_entry_type(EntryType::EntryNormal);
        plain_write.set_data(Bytes::from_static(b"plain-write"));
        cases.push(("plain-write", plain_write));

        let mut admin_cmd = RaftCmdRequest::default();
        let mut admin_req = AdminRequest::default();
        admin_req.set_cmd_type(AdminCmdType::TransferLeader);
        admin_cmd.set_admin_request(admin_req);
        cases.push(("admin-request", build_preprocess_entry(&admin_cmd)));

        let mut switch_memtable_builder = CustomBuilder::new();
        switch_memtable_builder.set_switch_mem_table(128);
        let mut switch_memtable_cmd = RaftCmdRequest::default();
        switch_memtable_cmd.set_custom_request(switch_memtable_builder.build());
        cases.push((
            "custom-switch-memtable",
            build_preprocess_entry(&switch_memtable_cmd),
        ));

        let mut txn_file_builder = CustomBuilder::new();
        txn_file_builder.set_txn_file(&TxnFileRef::new());
        let mut txn_file_cmd = RaftCmdRequest::default();
        txn_file_cmd.set_custom_request(txn_file_builder.build());
        cases.push(("custom-txn-file-ref", build_preprocess_entry(&txn_file_cmd)));

        let mut change_set = ChangeSet::new();
        change_set.set_shard_ver(shard_ver);
        change_set.mut_flush().set_version(base_version + 17);
        let mut engine_meta_builder = CustomBuilder::new();
        engine_meta_builder.set_change_set(&change_set);
        let mut engine_meta_cmd = RaftCmdRequest::default();
        engine_meta_cmd.set_custom_request(engine_meta_builder.build());
        cases.push((
            "custom-engine-meta-flush",
            build_preprocess_entry(&engine_meta_cmd),
        ));

        for (name, entry) in cases {
            let (recover_affect_memtable, _) =
                get_affect_memtable_for_recover_entry(&entry, peer_tag, base_version, shard_ver);

            let mut sync_entry = entry.clone();
            let preprocess_cmd = crate::update_entry(&mut sync_entry, merged_store_id);
            let sync_affect_memtable = if sync_entry.get_data().is_empty() {
                AffectMemtable::None
            } else {
                crate::get_affect_memtable(peer_tag, preprocess_cmd.as_ref(), Some(base_version))
            };

            let same = match (&recover_affect_memtable, &sync_affect_memtable) {
                (AffectMemtable::None, AffectMemtable::None)
                | (AffectMemtable::Write, AffectMemtable::Write) => true,
                (
                    AffectMemtable::Persist { data_seq: left_seq },
                    AffectMemtable::Persist {
                        data_seq: right_seq,
                    },
                ) => left_seq == right_seq,
                _ => false,
            };
            assert!(
                same,
                "classification mismatch for case: {}, recover: {:?}, sync: {:?}",
                name, recover_affect_memtable, sync_affect_memtable,
            );
        }

        let mut mismatch_cs = ChangeSet::new();
        mismatch_cs.set_shard_ver(shard_ver + 1);
        mismatch_cs.mut_flush().set_version(base_version + 17);
        let mut mismatch_builder = CustomBuilder::new();
        mismatch_builder.set_change_set(&mismatch_cs);
        let mut mismatch_cmd = RaftCmdRequest::default();
        mismatch_cmd.set_custom_request(mismatch_builder.build());
        let mismatch_entry = build_preprocess_entry(&mismatch_cmd);

        let (affect_memtable, _) = get_affect_memtable_for_recover_entry(
            &mismatch_entry,
            peer_tag,
            base_version,
            shard_ver,
        );
        assert!(
            matches!(affect_memtable, AffectMemtable::None),
            "mismatched shard version should be skipped in recover, got: {:?}",
            affect_memtable
        );
    }

    #[test]
    fn test_recover_skips_mismatched_txn_file_ref_version() {
        let region_id = 13;
        let region_ver = 1;
        let existing_start_ts = 301;
        let existing_lock = make_lock_prefix(b"existing-primary".to_vec(), existing_start_ts);
        let entry = build_txn_file_ref_entry(
            RAFT_INIT_LOG_INDEX + 1,
            region_ver + 1,
            303,
            make_lock_prefix(b"primary-a".to_vec(), 303),
            vec![],
        );
        let raft_state = new_raft_state(region_id, entry.index, entry.index);
        let mut shard_meta = new_shard_meta(region_id, region_ver);
        let mut existing_ref = TxnFileRef::new();
        existing_ref.set_start_ts(existing_start_ts);
        existing_ref.set_lock_val_prefix(existing_lock.clone());
        shard_meta.merge_txn_file_ref(&existing_ref, RAFT_INIT_LOG_INDEX);
        let mut env =
            seed_recover_test_peer(region_id, region_ver, shard_meta, raft_state, &[entry]);
        let mut ctx = RecoverTxnFileLocksContext::new(TEST_STORE_ID, &env.raft);

        env.preprocessor
            .recover(
                &mut ctx.as_ref(),
                &mut env.progress,
                RAFT_INIT_LOG_INDEX + 2,
            )
            .unwrap();

        let shard_meta = env.preprocessor.shard_meta().unwrap();
        assert!(shard_meta.has_txn_file_locks());
        assert_eq!(shard_meta.txn_file_locks().len(), 1);
        assert_eq!(shard_meta.txn_file_locks().seq(), RAFT_INIT_LOG_INDEX);
        assert_eq!(
            shard_meta
                .txn_file_locks()
                .get_locks()
                .get(&existing_start_ts)
                .unwrap()
                .as_ref(),
            existing_lock.as_slice()
        );
    }

    #[test]
    fn test_recover_boundary() {
        let region_id = 20;
        let region_ver = 1;
        let first_start_ts = 404;
        let second_start_ts = 405;
        let first = build_txn_file_ref_entry(
            RAFT_INIT_LOG_INDEX + 1,
            region_ver,
            first_start_ts,
            make_lock_prefix(b"parent-primary-a".to_vec(), first_start_ts),
            vec![],
        );
        let second = build_txn_file_ref_entry(
            RAFT_INIT_LOG_INDEX + 2,
            region_ver,
            second_start_ts,
            make_lock_prefix(b"parent-primary-b".to_vec(), second_start_ts),
            vec![],
        );
        let raft_state = new_raft_state(region_id, second.index, second.index);
        let mut shard_meta = new_shard_meta(region_id, region_ver);
        shard_meta.data_sequence = RAFT_INIT_LOG_INDEX + 2;
        shard_meta.parent = Some(Box::new(new_shard_meta(region_id, region_ver)));
        let parent_shard_meta = shard_meta.parent.as_deref().unwrap().clone();
        let mut env = seed_recover_test_peer(
            region_id,
            region_ver,
            shard_meta.clone(),
            raft_state,
            &[first, second],
        );
        env.preprocessor = Preprocessor::from_shard_meta(
            &env.raft,
            TEST_STORE_ID,
            parent_shard_meta,
            &MasterKey::new(&[1u8; 32]),
            Arc::new(EncryptionKeyManager::new()),
        );
        let mut ctx = RecoverTxnFileLocksContext::new(TEST_STORE_ID, &env.raft);

        env.preprocessor
            .recover(
                &mut ctx.as_ref(),
                &mut env.progress,
                shard_meta.data_sequence,
            )
            .unwrap();

        let locks = env.preprocessor.shard_meta().unwrap().txn_file_locks();
        assert_eq!(locks.len(), 1);
        assert_eq!(locks.seq(), RAFT_INIT_LOG_INDEX + 1);
        let lock_info = locks
            .get_lock_infos_by_start_tss(b"target-key", [first_start_ts])
            .unwrap()
            .pop()
            .unwrap();
        assert_eq!(lock_info.get_primary_lock(), b"parent-primary-a");
        assert!(lock_info.get_is_txn_file());
        assert!(!locks.get_locks().contains_key(&second_start_ts));
    }

    #[test]
    fn test_recover_txn_file_locks_across_parent_and_child() {
        let region_id = 81;
        let parent_region_ver = 1;
        let child_region_ver = parent_region_ver + 1;
        let parent_log_index = RAFT_INIT_LOG_INDEX + 1;
        let merge_log_index = RAFT_INIT_LOG_INDEX + 2;
        let child_log_index = RAFT_INIT_LOG_INDEX + 3;
        let parent_start_ts = 800;
        let child_start_ts = 810;
        let expected_parent_lock = make_lock_prefix(b"parent-primary".to_vec(), parent_start_ts);
        let expected_child_lock = make_lock_prefix(b"child-primary".to_vec(), child_start_ts);

        let dir = tempfile::Builder::new()
            .prefix("merged_engine_preprocessor_startup_")
            .tempdir()
            .unwrap();
        let merged_dir = dir.path().join("merged");
        let raft = RfEngine::open(
            merged_dir.as_path(),
            &rfengine::RfEngineConfig::default(),
            None,
            None,
        )
        .unwrap();
        raft.set_engine_id(TEST_STORE_ID);

        // Seed the original parent shard and recover its txn-file entry.
        let parent_shard_meta = new_shard_meta(region_id, parent_region_ver);
        let parent_entry = build_txn_file_ref_entry(
            parent_log_index,
            parent_region_ver,
            parent_start_ts,
            expected_parent_lock.clone(),
            vec![],
        );
        let parent_raft_state = new_raft_state(region_id, parent_log_index, parent_log_index);
        let mut wb = WriteBatch::new();
        let parent_region_state = new_region_state(region_id, parent_region_ver);
        write_peer_state(
            &mut wb,
            region_id,
            parent_region_state.get_region(),
            PeerState::Normal,
            None,
        );
        write_engine_meta(&mut wb, region_id, &parent_shard_meta);
        wb.set_state(
            region_id,
            region_id,
            TEST_KEYSPACE_ID,
            &raft_state_key(parent_region_ver),
            &encode_raft_state(&parent_raft_state),
        );
        wb.append_raft_log(region_id, region_id, TEST_KEYSPACE_ID, &parent_entry);
        raft.write(wb).unwrap();

        let mut progress = RegionProgress::new(TEST_KEYSPACE_ID, region_id);
        let mut preprocessor = Preprocessor::from_shard_meta(
            &raft,
            TEST_STORE_ID,
            parent_shard_meta,
            &MasterKey::new(&[1u8; 32]),
            Arc::new(EncryptionKeyManager::new()),
        );
        let mut recover_ctx = RecoverTxnFileLocksContext::new(TEST_STORE_ID, &raft);
        preprocessor
            .recover(&mut recover_ctx.as_ref(), &mut progress, merge_log_index)
            .unwrap();

        let mut shard_meta = preprocessor.shard_meta().unwrap().clone();
        assert!(shard_meta.has_txn_file_locks());
        assert_eq!(shard_meta.txn_file_locks().len(), 1);
        assert_eq!(shard_meta.txn_file_locks().seq(), parent_log_index);
        assert_eq!(
            shard_meta
                .txn_file_locks()
                .get_locks()
                .get(&parent_start_ts)
                .unwrap()
                .as_ref(),
            expected_parent_lock.as_slice()
        );
        let lock_info = shard_meta
            .txn_file_locks()
            .get_lock_infos_by_start_tss(b"target-key", [parent_start_ts])
            .unwrap()
            .pop()
            .unwrap();
        assert_eq!(lock_info.get_lock_version(), parent_start_ts);
        assert_eq!(lock_info.get_primary_lock(), b"parent-primary");
        assert!(lock_info.get_is_txn_file());

        // Simulate the merge version bump, then the initial flush that persists
        // the parent lock and clears the parent pointer.
        shard_meta.prepare_merge(merge_log_index);
        assert_eq!(shard_meta.ver, child_region_ver);
        assert!(shard_meta.parent.is_some());

        let mut initial_flush = ChangeSet::new();
        initial_flush.set_shard_id(region_id);
        initial_flush.set_shard_ver(child_region_ver);
        initial_flush.set_sequence(merge_log_index);
        initial_flush
            .mut_initial_flush()
            .set_outer_start(b"a".to_vec());
        initial_flush
            .mut_initial_flush()
            .set_outer_end(b"z".to_vec());
        initial_flush
            .mut_initial_flush()
            .set_inner_key_off(shard_meta.inner_key_off as u32);
        initial_flush
            .mut_initial_flush()
            .set_base_version(shard_meta.base_version);
        initial_flush
            .mut_initial_flush()
            .set_data_sequence(merge_log_index);
        shard_meta.apply_initial_flush(&initial_flush);
        assert!(shard_meta.parent.is_none());

        // Persist the post-flush shard meta, then recover the child txn-file
        // entry from the new shard version on top of the carried parent lock.
        let child_entry = build_txn_file_ref_entry(
            child_log_index,
            child_region_ver,
            child_start_ts,
            expected_child_lock.clone(),
            vec![],
        );
        let child_raft_state = new_raft_state(region_id, child_log_index, child_log_index);
        let mut wb = WriteBatch::new();
        let child_region_state = new_region_state(region_id, child_region_ver);
        write_peer_state(
            &mut wb,
            region_id,
            child_region_state.get_region(),
            PeerState::Normal,
            None,
        );
        write_engine_meta(&mut wb, region_id, &shard_meta);
        wb.set_state(
            region_id,
            region_id,
            TEST_KEYSPACE_ID,
            &raft_state_key(child_region_ver),
            &encode_raft_state(&child_raft_state),
        );
        wb.append_raft_log(region_id, region_id, TEST_KEYSPACE_ID, &child_entry);
        raft.write(wb).unwrap();
        drop(recover_ctx);

        let mut progress = RegionProgress::new(TEST_KEYSPACE_ID, region_id);
        let mut recover_ctx = RecoverTxnFileLocksContext::new(TEST_STORE_ID, &raft);
        let mut preprocessor = Preprocessor::from_shard_meta(
            &raft,
            TEST_STORE_ID,
            shard_meta,
            &MasterKey::new(&[1u8; 32]),
            Arc::new(EncryptionKeyManager::new()),
        );
        preprocessor
            .recover(
                &mut recover_ctx.as_ref(),
                &mut progress,
                child_log_index.saturating_add(1),
            )
            .unwrap();

        let shard_meta = preprocessor.shard_meta().unwrap();
        assert!(shard_meta.parent.is_none());
        assert_eq!(shard_meta.ver, child_region_ver);

        let locks = shard_meta.txn_file_locks();
        assert_eq!(locks.len(), 2);
        assert_eq!(locks.seq(), child_log_index);
        assert_eq!(
            locks.get_locks().get(&parent_start_ts).unwrap().as_ref(),
            expected_parent_lock.as_slice()
        );
        assert_eq!(
            locks.get_locks().get(&child_start_ts).unwrap().as_ref(),
            expected_child_lock.as_slice()
        );

        let parent_lock_info = locks
            .get_lock_infos_by_start_tss(b"target-key", [parent_start_ts])
            .unwrap()
            .pop()
            .unwrap();
        assert_eq!(parent_lock_info.get_lock_version(), parent_start_ts);
        assert_eq!(parent_lock_info.get_primary_lock(), b"parent-primary");
        assert!(parent_lock_info.get_is_txn_file());

        let child_lock_info = locks
            .get_lock_infos_by_start_tss(b"target-key", [child_start_ts])
            .unwrap()
            .pop()
            .unwrap();
        assert_eq!(child_lock_info.get_lock_version(), child_start_ts);
        assert_eq!(child_lock_info.get_primary_lock(), b"child-primary");
        assert!(child_lock_info.get_is_txn_file());
    }
}