// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    borrow::Cow,
    collections::{BTreeMap, HashSet},
    fmt::{Display, Formatter},
    fs,
    fs::{File, OpenOptions, create_dir_all},
    mem,
    ops::{Deref, DerefMut},
    os::unix::fs::MetadataExt,
    path::{Path, PathBuf},
    sync::{
        Arc, Mutex, RwLock,
        atomic::{AtomicBool, AtomicI64, AtomicU32, AtomicU64, Ordering},
    },
    thread::JoinHandle,
    time::Duration,
};

use api_version::ApiV2;
use arc_swap::ArcSwap;
use bytes::{Buf, Bytes};
use chrono::Utc;
use file_system::{IoRateLimitMode, IoRateLimiter, open_direct_file};
use futures::{Future, FutureExt};
use kvengine::dfs::Dfs;
use kvproto::raft_serverpb::{RegionLocalState, StoreIdent};
use protobuf::Message;
use raft_proto::eraftpb;
use rfenginepb::{ChangeSet, KeySpaceBackupMeta, StoreBackupMeta, StoreRaftLogBackupMeta};
use tikv_util::{
    error,
    errors::Context as _,
    info,
    mpsc::{SendError, Sender},
    panic_mark_dfs_worker_file_exists, spawn_anonymous_thread_with,
    sys::thread::StdThreadBuildWrapper,
    time::Instant,
    warn,
};

use crate::{
    config::Config,
    log_batch::{RaftLogBlock, RaftLogs},
    manifest::{Manifest, manifest_path, persist_change_set},
    metrics::*,
    peers::RaftPeers,
    service_worker::{DfsWorkerStatus, ServiceTask, ServiceWorker, WalProgress},
    write_batch::{PeerBatch, WriteBatch},
    *,
};

pub const MIN_EPOCH_ROTATE_LEN: usize = 4;
pub const MAX_EPOCH_ROTATE_LEN: usize = 32;

pub const TRUNCATE_ALL_INDEX: u64 = u64::MAX;
pub const MAX_EPOCH_BACKWARD: u32 = 64;
pub const MIN_RLOG_FILE_SIZE: u64 = 16 * 1024 * 1024; // 16MB
const RAFT_INIT_LOG_INDEX: u64 = 5;

/// `RfEngine` is a persistent storage engine for multi-raft logs.
#[derive(Clone)]
pub struct RfEngine {
    core: Arc<RfEngineCore>,
}

impl Deref for RfEngine {
    type Target = RfEngineCore;

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

impl RfEngine {
    #[inline]
    pub fn db_dir(&self) -> Option<&Path> {
        Some(&self.core.dir)
    }

    #[inline]
    pub fn wal_dir(&self) -> Option<&Path> {
        Some(self.core.wal_dir())
    }

    /// Returns the directories used by synchronous WAL writers.
    pub fn sync_wal_dirs(&self) -> Vec<&Path> {
        let mut dirs = vec![self.core.wal_dir()];
        dirs.extend(self.core.wal_secondary_dir.as_deref());
        dirs
    }

    /// Returns the latest successful async WAL sync timestamp in monotonic
    /// milliseconds.
    #[inline]
    pub fn async_wal_sync_success_at_millis(&self) -> Arc<AtomicU64> {
        self.core.async_wal_sync_success_at_millis.clone()
    }
}

impl RfEngine {
    // NOTE: Pass `Some(dfs)` for usage of TiKV servers ONLY. Otherwise, it would
    // corrupt the WAL chunks in DFS.
    pub fn open(
        dir: &Path,
        cfg: &Config,
        data_dir: Option<&Path>, // for check panic mark file exists
        dfs: Option<Arc<dyn Dfs>>,
    ) -> Result<Self> {
        let core = RfEngineCore::open(dir, cfg, data_dir, dfs)?;
        Ok(Self {
            core: Arc::new(core),
        })
    }
}

pub struct RfEngineCore {
    /// The primary directory holding engine data (e.g., WAL files).
    pub dir: PathBuf,

    /// An optional directory for synchronous WAL writes. When set, WAL files
    /// are synchronously written to `wal_sync_dir` and asynchronously written
    /// to `dir`. `wal_sync_dir` is write-only, while `dir` handles reads as
    /// well (e.g. compaction and backup). This setup allows `wal_sync_dir` to
    /// use dedicated IOPS for writes, reducing tail latency.
    pub wal_sync_dir: Option<PathBuf>,

    /// An optional directory for the secondary synchronous WAL writer.
    wal_secondary_dir: Option<PathBuf>,

    /// The sync WAL writer.
    pub(crate) writer: Mutex<WalWriterExt>,

    /// A concurrent map for storing all peers' in-memory state.
    pub(crate) peers: RaftPeers,

    /// For each `region_id`, holds a set of peer IDs that depend on this
    /// region’s logs, preventing their early truncation.
    pub(crate) dependants: papaya::HashMap<u64, RwLock<HashSet<u64>>>,

    /// A channel sender used to schedule background tasks (e.g., WAL rotation,
    /// truncation).
    pub(crate) task_sender: Sender<ServiceTask>,

    /// A handle for the background worker thread.
    pub(crate) service_worker_handle: Mutex<Option<JoinHandle<()>>>,

    /// A unique id for this engine instance.
    pub(crate) engine_id: Arc<AtomicU64>,

    /// Whether lightweight backup mode is enabled at runtime.
    pub(crate) lightweight: Arc<AtomicBool>,

    pub(crate) dfs_worker_healthy: Healthy,

    pub(crate) epoch_rotate_len: usize,

    /// Tracks the latest epoch ID for WAL. Initialized during engine load and
    /// updated after wal rotation of the sync writer.
    pub(crate) current_epoch_id: Arc<AtomicU32>,

    /// Tracks the epoch ID that has been compacted.
    pub(crate) compacted_epoch: Arc<AtomicU32>,

    /// Latest successful sync of the async WAL under the primary raft data dir.
    pub(crate) async_wal_sync_success_at_millis: Arc<AtomicU64>,

    /// The number of WAL files whose raft entries can be kept in memory.
    /// Calculated from the rlog memory limit and the WAL file size.
    pub(crate) in_mem_rlog_epoch_count: u32,

    /// DFS statistics collected by DFS worker.
    pub(crate) dfs_statistics: Arc<DfsMeter>,

    /// Unix timestamp milliseconds of last rotation.
    pub(crate) last_rotate_time: Arc<AtomicI64>,
    pub(crate) force_rotation_interval: Option<Duration>,

    _lock: fslock::LockFile, // hold lock to avoid release

    #[cfg(any(test, feature = "testexport"))]
    pub(crate) compact_force_stop: Arc<AtomicBool>,
}

impl Deref for RfEngineCore {
    type Target = RaftPeers;
    fn deref(&self) -> &Self::Target {
        &self.peers
    }
}

impl RfEngineCore {
    fn open(
        dir: &Path,
        cfg: &Config,
        data_dir: Option<&Path>,
        dfs: Option<Arc<dyn Dfs>>,
    ) -> Result<Self> {
        if !cfg!(feature = "testexport")
            && (cfg.rlog_file_size.0 < MIN_RLOG_FILE_SIZE || cfg.rlog_file_size.0 > u32::MAX as u64)
        {
            return Err(Error::Other(
                "invalid config: rlog_file_size must be between [16MB, 4GB]".to_owned(),
            ));
        }

        let wal_size = cfg.target_file_size.0;
        if wal_size == 0 {
            return Err(Error::Other(
                "invalid config: wal_size must be non-zero".to_owned(),
            ));
        }
        if cfg.compaction_concurrency == 0 {
            return Err(Error::Other(
                "invalid config: compaction_concurrency must be non-zero".to_owned(),
            ));
        }
        if cfg.compaction_batch_size == 0 {
            return Err(Error::Other(
                "invalid config: compaction_batch_size must be non-zero".to_owned(),
            ));
        }
        if cfg.rlog_cache_capacity.0 > 0 && cfg.rlog_cache_size_threshold.0 == 0 {
            return Err(Error::Other(
                "invalid config: rlog_cache_size_threshold must be non-zero when rlog_cache_capacity is non-zero"
                    .to_owned(),
            ));
        }
        let in_mem_rlog_epoch_count = cfg.rlog_soft_memory_limit.0.div_ceil(wal_size);
        info!(
            "rfengine in_mem_rlog_epoch_count: {}",
            in_mem_rlog_epoch_count
        );

        let wal_sync_dir = (!cfg.wal_sync_dir.is_empty()).then(|| PathBuf::from(&cfg.wal_sync_dir));
        let wal_secondary_dir =
            (!cfg.wal_secondary_dir.is_empty()).then(|| PathBuf::from(&cfg.wal_secondary_dir));
        init_wal_files(
            dir,
            wal_sync_dir.as_ref(),
            wal_secondary_dir.as_ref(),
            cfg.epoch_rotate_len,
        )?;

        // Lock the rfengine directory to prevent concurrent opening.
        let lock = match try_acquire_lock_file(dir)? {
            Some(lock) => lock,
            None => panic!("{}", lock_conflict_message(dir)),
        };

        let engine_id = Arc::new(AtomicU64::new(0));
        let manifest = Manifest::open(dir, engine_id.clone())?;
        let (service_tx, service_rx) = tikv_util::mpsc::unbounded();
        let compacted_epoch = Arc::new(AtomicU32::new(manifest.epoch_id));
        let wal_dir = wal_sync_dir.as_deref().unwrap_or(dir);
        let writer_type = if cfg.cli_mode {
            WriterType::CliMode
        } else {
            WriterType::Sync
        };
        let writer = WalWriter::new(wal_dir, cfg, compacted_epoch.clone(), writer_type);
        let writer_ext = if cfg.wal_secondary_dir.is_empty() {
            WalWriterExt::SingleWriter(writer)
        } else {
            let wal_secondary_dir = PathBuf::from(&cfg.wal_secondary_dir);
            let secondary_writer = WalWriter::new(
                &wal_secondary_dir,
                cfg,
                compacted_epoch.clone(),
                writer_type,
            );
            WalWriterExt::DoubleWriter(DoubleWriter::new(
                writer,
                secondary_writer,
                manifest.epoch_id,
                cfg.wal_double_write_unhealthy_size.0 as usize,
            )?)
        };
        let dfs_worker_healthy = dfs_worker::Healthy::default();
        let peer_rlog_files = Arc::new(ArcSwap::from_pointee(manifest.peer_rlog_files()));
        #[cfg(any(test, feature = "testexport"))]
        let compact_force_stop = Arc::new(AtomicBool::new(false));
        let mut en = Self {
            dir: dir.to_owned(),
            wal_sync_dir,
            wal_secondary_dir,
            peers: RaftPeers::new(dir.to_path_buf(), peer_rlog_files.clone()),
            dependants: Default::default(),
            writer: Mutex::new(writer_ext),
            task_sender: service_tx,
            service_worker_handle: Mutex::new(None),
            engine_id,
            lightweight: Arc::new(AtomicBool::new(cfg.lightweight_backup)),
            dfs_worker_healthy: dfs_worker_healthy.clone(),
            epoch_rotate_len: cfg.epoch_rotate_len,
            current_epoch_id: Arc::new(AtomicU32::new(0)),
            compacted_epoch: compacted_epoch.clone(),
            async_wal_sync_success_at_millis: Arc::new(AtomicU64::new(0)),
            in_mem_rlog_epoch_count: in_mem_rlog_epoch_count as u32,
            dfs_statistics: Arc::default(),
            last_rotate_time: Arc::new(AtomicI64::new(now_unix_millis())),
            force_rotation_interval: cfg.force_rotate_interval.map(|v| v.0),
            _lock: lock,
            #[cfg(any(test, feature = "testexport"))]
            compact_force_stop: compact_force_stop.clone(),
        };
        let async_epoch_offset = en.load(&manifest)?;
        if cfg.disable_compaction {
            return Ok(en);
        }
        {
            let async_wal_writer = if let Some((async_epoch_id, async_offset)) = async_epoch_offset
            {
                debug_assert!(en.is_async_wal_enabled());
                let mut async_wal_writer =
                    WalWriter::new(dir, cfg, compacted_epoch.clone(), WriterType::Async);
                async_wal_writer.open_file(async_epoch_id, async_offset)?;
                Some(async_wal_writer)
            } else {
                None
            };

            let lightweight_backup_args: Option<(LightweightBackupConfig, Arc<dyn Dfs>, bool)> =
                if let Some(dfs) = dfs {
                    if data_dir.is_some() && panic_mark_dfs_worker_file_exists(data_dir.unwrap()) {
                        // If panic_mark_dfs_worker_file exists, skip init dfs worker thread.
                        en.dfs_worker_healthy.set_unhealthy();
                        error!("panic_mark_dfs_worker_file exists, skip init dfs worker thread");
                        None
                    } else {
                        let initially_enabled = cfg.lightweight_backup;
                        let dfs_worker_cfg = LightweightBackupConfig::new(
                            dir.to_owned(),
                            cfg.wal_chunk_target_file_size.0 as usize,
                            cfg.snapshot_rlog_multipart_upload_threshold.0 as usize,
                            cfg.snapshot_rlog_multipart_upload_part_size.0 as usize,
                            cfg.snapshot_rlog_multipart_upload_concurrency,
                            cfg.snapshot_rlog_max_object_size.0 as usize,
                            CompressionType::Lz4Compression,
                            CompressionType::Lz4Compression,
                            cfg.dfs_worker_memory_limit.as_memory_size() as usize,
                            cfg.max_wal_chunk_gap_duration,
                        );
                        Some((dfs_worker_cfg, dfs, initially_enabled))
                    }
                } else {
                    None
                };
            let compact_rate_limiter = if cfg.enable_compact_rate_limiter {
                let rate_limiter =
                    Arc::new(IoRateLimiter::new(IoRateLimitMode::WriteOnly, true, false));
                rate_limiter.set_io_rate_limit(cfg.compact_bytes_per_sec.0 as usize);
                Some(rate_limiter)
            } else {
                None
            };
            let mut service_worker = ServiceWorker::new(
                dir.to_owned(),
                cfg,
                async_wal_writer,
                service_rx,
                manifest,
                compacted_epoch.clone(),
                lightweight_backup_args,
                en.lightweight.clone(),
                dfs_worker_healthy,
                compact_rate_limiter,
                en.async_wal_sync_success_at_millis.clone(),
                #[cfg(any(test, feature = "testexport"))]
                compact_force_stop,
                peer_rlog_files,
            );
            if let Some(dfs_statistic) = service_worker.dfs_statistic() {
                en.dfs_statistics = dfs_statistic
            }
            let join_handle = spawn_anonymous_thread_with!(move || service_worker.run());
            let mut guard = en.service_worker_handle.lock().unwrap();
            *guard = Some(join_handle);
        }

        Ok(en)
    }

    pub(crate) fn wal_dir(&self) -> &Path {
        self.wal_sync_dir.as_ref().unwrap_or(&self.dir)
    }

    /// Applies and persists the write batch.
    pub fn write(&self, wb: WriteBatch) -> Result<usize> {
        self.apply(&wb);
        self.persist(wb)
    }

    /// Applies the write batch to memory without persisting it to WAL.
    pub fn apply(&self, wb: &WriteBatch) {
        let offload_epoch = calc_offload_epoch(
            self.in_mem_rlog_epoch_count,
            &self.current_epoch_id,
            &self.compacted_epoch,
        );
        let truncated_logs = self.peers.apply(wb, offload_epoch);
        if !truncated_logs.is_empty() {
            self.try_send_task(ServiceTask::Truncates(truncated_logs));
        }
    }

    /// Persists the write batch to WAL. It can be used in another thread to
    /// implement async I/O, i.e., call `apply` in the main thread and call
    /// `persist` in the I/O thread.
    /// When the epoch is rotated, the return size would be 4096 which is not
    /// accurate but ok to be used as metrics.
    pub fn persist(&self, wb: WriteBatch) -> Result<usize> {
        let timer = Instant::now();
        let wb = Arc::new(wb.into_vector());
        let mut writer = self.writer.lock().unwrap();
        let old_file_off = writer.get_file_off();
        let last_rotation = self.last_rotate_time.load(Ordering::SeqCst);
        let last_rotation_gap =
            Duration::from_millis(now_unix_millis().saturating_sub(last_rotation) as u64);
        let need_rotate = self.force_rotation_interval.is_some_and(|interval| {
            self.lightweight.load(Ordering::SeqCst) && last_rotation_gap > interval
        });
        let (epoch_id, file_off, rotated) = writer.write_batch(wb.clone(), need_rotate)?;
        if rotated {
            self.last_rotate_time
                .store(now_unix_millis(), Ordering::SeqCst);
            self.current_epoch_id.store(epoch_id, Ordering::SeqCst);
            self.try_send_task(ServiceTask::Rotate {
                epoch_id: epoch_id - 1,
                cache_wb_for_compact: true,
                require_snapshot: need_rotate,
            });
        }
        self.try_send_task(ServiceTask::Write { wb });
        ENGINE_PERSIST_DURATION_HISTOGRAM.observe(timer.saturating_elapsed_secs());
        let write_size = if rotated {
            file_off
        } else {
            file_off.saturating_sub(old_file_off)
        } as usize;
        Ok(write_size)
    }

    pub fn stop_worker(&self, force: bool) {
        #[cfg(any(test, feature = "testexport"))]
        if force {
            self.compact_force_stop.store(true, Ordering::SeqCst);
        }
        let mut handle = self.service_worker_handle.lock().unwrap();
        if let Some(h) = handle.take() {
            self.try_send_task(ServiceTask::Close { force });
            h.join().unwrap();
        }
    }

    pub fn close_writer(&self) {
        let mut _double_writer = None;
        {
            let mut writer = self.writer.lock().unwrap();
            match writer.deref_mut() {
                WalWriterExt::SingleWriter(_) => {}
                WalWriterExt::DoubleWriter(w) => {
                    _double_writer = Some(mem::take(w));
                }
            }
        }
    }

    /// After split and before the new region is initially flushed, the old
    /// region's raft log can not be truncated, otherwise, it would not be
    /// able to recover the new region. So we can call `add_dependent` after
    /// split to protect the raft log. After the new region is initially
    /// flushed or re-ingested or destroyed, call `remove_dependent` to
    /// resume truncating the raft log.
    pub fn add_dependent(&self, region_id: u64, dependent_id: u64) {
        let dependants = self.dependants.pin();
        let hs_ref = dependants.get_or_insert_with(region_id, || RwLock::new(HashSet::new()));
        let mut hs = hs_ref.write().unwrap();
        let newly_inserted = hs.insert(dependent_id);
        if !newly_inserted {
            return;
        }
        let len = hs.len();
        drop(hs);
        let tag = PeerTag::new(self.get_engine_id(), region_id);
        info!(
            "{} add dependent {}, dependents_len {}",
            tag, dependent_id, len
        );
    }

    pub fn remove_dependent(&self, region_id: u64, dependent_id: u64) -> usize {
        let dependants = self.dependants.pin();
        dependants
            .get(&region_id)
            .map(|hs| {
                let len = {
                    let mut hs = hs.write().unwrap();
                    hs.remove(&dependent_id);
                    hs.len()
                };
                let tag = PeerTag::new(self.get_engine_id(), region_id);
                info!(
                    "{} remove dependent {}, dependents_len {}",
                    tag, dependent_id, len
                );
                len
            })
            .unwrap_or_default()
    }

    pub fn with_dependents(&self, region_id: u64, f: impl FnOnce(&HashSet<u64>)) {
        let dependants = self.dependants.pin();
        if let Some(hs) = dependants.get(&region_id) {
            f(&hs.read().unwrap());
        }
    }

    pub fn has_dependents(&self, region_id: u64) -> bool {
        let dependants = self.dependants.pin();
        dependants
            .get(&region_id)
            .is_some_and(|hs| !hs.read().unwrap().is_empty())
    }

    pub fn pending_compaction_wals(&self) -> u8 {
        (self.current_epoch_id.load(Ordering::SeqCst)
            - 1
            - self.compacted_epoch.load(Ordering::SeqCst)) as u8
    }

    /// Dumps the state of the engine.
    pub fn get_engine_stats(&self) -> EngineStats {
        let mut total_mem_size = 0;
        let mut total_mem_entries = 0;
        let mut total_offloaded_entries = 0;
        let mut total_num_logs = 0;
        let peers = self.peers.peers.pin();
        let mut peers_stats = peers
            .iter()
            .map(|(peer_id, data)| {
                let peer_stats = data.read().unwrap().get_stats();
                total_mem_size += peer_stats.size;
                total_mem_entries += peer_stats.num_logs;
                total_offloaded_entries += peer_stats.num_logs_offloaded;
                // peer_stats.num_logs is the log number in memory, so we use
                // first_index and last_index to calculate the real log count.
                let truncated_index = self.get_truncated_index(*peer_id);
                let last_index = self.get_last_index(*peer_id);
                let num_logs = match (truncated_index, last_index) {
                    (Some(TRUNCATE_ALL_INDEX), Some(_)) => 0,
                    // `truncated_index` can be 0 (default/uninitialized), but
                    // the effective minimum truncated index is
                    // `RAFT_INIT_LOG_INDEX`. During startup/recovery a peer may
                    // temporarily have `last_index < RAFT_INIT_LOG_INDEX`; use
                    // `saturating_sub` to avoid underflow and treat it as 0
                    // logs.
                    (Some(t), Some(l)) => l.saturating_sub(t.max(RAFT_INIT_LOG_INDEX)),
                    // last_index is 0 means there is no raft logs,
                    _ => 0,
                };
                total_num_logs += num_logs;
                peer_stats
            })
            .collect::<Vec<PeerStats>>();
        peers_stats.sort_by(|a, b| (b.size).cmp(&a.size));
        peers_stats.truncate(10);

        let mut disk_size = 0;
        let mut num_files = 0;
        if let Ok(read_dir) = self.dir.read_dir() {
            for e in read_dir.flatten() {
                if let Ok(m) = e.metadata() {
                    num_files += 1;
                    disk_size += m.size();
                }
            }
        }
        let pending_compaction_wals = self.pending_compaction_wals();
        ENGINE_PENDING_COMPACTION_WALS_GAUGE.set(pending_compaction_wals as i64);
        // Flush metrics.
        ENGINE_RESOUCE_USAGE.memory.set(total_mem_size as i64);
        ENGINE_RESOUCE_USAGE.disk.set(disk_size as i64); // TODO: move it to CSE.store_size
        ENGINE_ENTRIES_COUNT.memory.set(total_mem_entries as i64);
        ENGINE_ENTRIES_COUNT
            .offloaded
            .set(total_offloaded_entries as i64);
        ENGINE_TOTAL_WALS_GAUGE.set(num_files as i64);

        EngineStats {
            total_mem_size,
            total_mem_entries,
            disk_size,
            num_files,
            pending_compaction_wals,
            top_10_size_peers: peers_stats,
            total_num_logs,
        }
    }

    pub fn dfs_stats(&self) -> &Arc<DfsMeter> {
        &self.dfs_statistics
    }

    pub fn set_engine_id(&self, engine_id: u64) {
        self.engine_id.store(engine_id, Ordering::Release)
    }

    pub fn get_engine_id(&self) -> u64 {
        self.engine_id.load(Ordering::Acquire)
    }

    // Upload latest wal chunk to object storage
    pub fn upload_wal_chunk(&self) -> impl Future<Output = ()> {
        let (tx, rx) = tokio::sync::oneshot::channel();
        self.try_send_task(ServiceTask::Upload(tx));
        rx.map(|_| ())
    }

    pub fn dump_wal_chunk(
        &self,
        epoch_id: u32,
        start_off: u64,
        end_off: u64,
        callback: Box<dyn FnOnce(Result<(Bytes, bool /* partial content */)>) + Send>,
    ) {
        self.try_send_task(ServiceTask::Dump {
            epoch_id,
            start_off,
            end_off,
            callback,
        });
    }

    pub fn get_wal_progress(&self, callback: Box<dyn FnOnce(Result<WalProgress>) + Send>) {
        self.try_send_task(ServiceTask::GetProgress { callback });
    }

    /// Returns the current runtime state of the DFS worker.
    pub async fn get_dfs_worker_status(&self) -> Result<DfsWorkerStatus> {
        let (callback, receiver) = tokio::sync::oneshot::channel();
        self.try_send_task(ServiceTask::GetDfsWorkerStatus { callback });
        receiver
            .await
            .map_err(|_| Error::Other("DFS worker dropped the status request".to_string()))?
    }

    /// Enables or disables lightweight backup for this process at runtime.
    pub async fn set_lightweight_backup_enabled(&self, enabled: bool) -> Result<()> {
        let (callback, receiver) = tokio::sync::oneshot::channel();
        self.try_send_task(ServiceTask::SetLightweightBackup { enabled, callback });
        receiver.await.map_err(|_| {
            Error::Other("DFS worker dropped the runtime configuration request".to_string())
        })?
    }

    pub fn backup(&self, mut task: BackupTask) {
        if !self.is_async_wal_enabled() {
            // Note: when async wal is enabled, `file_off` is acquired from
            // `async_wal_writer` in worker.
            let writer = self.writer.lock().unwrap();
            task.file_off = writer.get_file_off();
        }
        self.try_send_task(ServiceTask::Backup(task));
    }

    pub fn get_epoch_offset(&self) -> (u32, u64) {
        let writer = self.writer.lock().unwrap();
        (writer.get_epoch_id(), writer.get_file_off())
    }

    pub(crate) fn is_async_wal_enabled(&self) -> bool {
        self.wal_sync_dir.is_some()
    }

    pub fn is_lightweight_backup_enabled(&self) -> bool {
        self.lightweight.load(Ordering::SeqCst)
    }

    pub fn is_dfs_worker_healthy(&self) -> bool {
        self.dfs_worker_healthy.is_healthy()
    }

    pub(crate) fn try_send_task(&self, task: ServiceTask) {
        if let Err(SendError(task)) = self.task_sender.send(task) {
            warn!("send service task failed: {:?}", task);
            let err_msg = "service worker is closed".to_string();
            match task {
                ServiceTask::Dump { callback, .. } => callback(Err(Error::Other(err_msg))),
                ServiceTask::GetProgress { callback } => callback(Err(Error::Other(err_msg))),
                ServiceTask::GetDfsWorkerStatus { callback } => {
                    let _ = callback.send(Err(Error::Other(err_msg)));
                }
                ServiceTask::SetLightweightBackup { callback, .. } => {
                    let _ = callback.send(Err(Error::Other(err_msg)));
                }
                ServiceTask::Backup(task) => (task.callback)(Err(Error::Backup(err_msg))),
                _ => {}
            }
        }
    }
}

impl Drop for RfEngineCore {
    fn drop(&mut self) {
        self.stop_worker(false);
    }
}

/// Calculates the epoch up to which raft logs can be offloaded.
///
/// Keeps `in_mem_rlog_epoch_count` most recent epochs in memory. Also,
/// offloaded epoch must not exceed `compacted_epoch`.
///
/// Returns:
/// - 0 if offloading is disabled (`in_mem_rlog_epoch_count == 0`)
/// - otherwise, `min(current_epoch - N, compacted_epoch)`
pub(crate) fn calc_offload_epoch(
    in_mem_rlog_epoch_count: u32,
    current_epoch: &AtomicU32,
    compacted_epoch: &AtomicU32,
) -> u32 {
    if in_mem_rlog_epoch_count == 0 {
        return 0;
    }

    let current_epoch = current_epoch.load(Ordering::SeqCst);
    let compacted_epoch = compacted_epoch.load(Ordering::SeqCst);
    current_epoch
        .saturating_sub(in_mem_rlog_epoch_count)
        .min(compacted_epoch)
}

pub(crate) fn get_delayed_to_epoch_id(manifest: &ChangeSet) -> u32 {
    manifest.get_epoch_id() + manifest.get_delayed_epoches()
}

fn decompress_snap_rlog_file(compression_type: u32, content: &[u8]) -> Result<Cow<'_, [u8]>> {
    match CompressionType::from(compression_type) {
        CompressionType::Lz4Compression => {
            let decompressed = decompress_lz4(content).map_err(Error::from)?;
            Ok(Cow::Owned(decompressed))
        }
        CompressionType::NoCompression => Ok(Cow::Borrowed(content)),
    }
}

fn restore_all_raft_logs_with_snap_rlog_file(
    keyspace_id: Option<u32>,
    store_meta: &StoreBackupMeta,
    dir: &Path,
    rlog_data: &Bytes,
) -> Result<()> {
    let mut raft_meta = StoreRaftLogBackupMeta::default();
    let size = rlog_data.len();
    debug_assert!(size as u64 > store_meta.raft_meta_start_off);

    raft_meta
        .merge_from_bytes(&rlog_data.chunk()[store_meta.raft_meta_start_off as usize..size])
        .unwrap();
    let compression_type = raft_meta.get_header().get_compression_type();

    let handle_keyspace = |keyspace_meta: &KeySpaceBackupMeta| -> Result<()> {
        for file in keyspace_meta.get_files() {
            let path = raft_log_file_name(dir, file.peer_id, file.first_index, file.last_index);
            let content = decompress_snap_rlog_file(
                compression_type,
                &rlog_data.chunk()[file.start_off as usize..file.end_off as usize],
            )?;
            fs::write(&path, content).with_ctx(|| format!("write rlog {}", path.display()))?;
        }
        Ok(())
    };

    if let Some(keyspace_id) = keyspace_id {
        if let Some(keyspace_meta) = raft_meta.raft_logs.get(&keyspace_id) {
            handle_keyspace(keyspace_meta)?;
        }
    } else {
        for (_, keyspace_meta) in raft_meta.raft_logs {
            handle_keyspace(&keyspace_meta)?;
        }
    }
    Ok(())
}

pub fn find_latest_snapshot(
    dfs: Arc<dyn Dfs>,
    prefix: &str,
    store_id: u64,
    epoch_id: u32,
) -> Result<Option<String>> {
    let start_epoch = if epoch_id > MAX_EPOCH_BACKWARD {
        epoch_id - MAX_EPOCH_BACKWARD
    } else {
        1
    };

    // Find the latest snapshot smaller than cluster_backup epoch.
    let snapshot = match dfs.list_objects(
        &snapshot_rlog_key_suffix(start_epoch - 1),
        Some(&snapshot_rlog_key_prefix(store_id)),
        Some(MAX_EPOCH_BACKWARD),
    ) {
        Ok((objects, _)) => {
            if objects.is_empty() {
                None
            } else {
                objects.into_iter().rev().find(|obj| {
                    let snap_delayed_to_epoch =
                        parse_delayed_to_epoch_from_snapshot_key(Some(obj.key.deref())).unwrap();
                    snap_delayed_to_epoch < epoch_id
                })
            }
        }
        Err(err) => {
            error!("failed to list snapshot full backup files: {}", err);
            None
        }
    };
    Ok(snapshot.map(|snap| {
        snap.key
            .strip_prefix(&format!("{}/", prefix))
            .unwrap()
            .to_owned()
    }))
}

// For lightweight restore, find a latest snapshot full backup before
// `cluster_backup.backup_ts` and replay all WAL chunk files from snapshot epoch
// to epoch of the backup.
pub fn lightweight_restore(
    store_id: u64,
    keyspace_id: Option<u32>,
    dir: &Path,
    snap_epoch: u32,
    snap_meta: Bytes,
    snap_rlog: Bytes,
    epoch_rotate_len: usize,
) -> Result<u32> {
    let start_time = Instant::now_coarse();

    if !dir.exists() {
        create_dir_all(dir)?;
    }
    // Use the same lock file as `RfEngineCore::open` so restore fails fast if
    // the target directory is still in use by a running engine process.
    let _lock = match try_acquire_lock_file(dir)? {
        Some(lock) => lock,
        None => return Err(Error::Other(lock_conflict_message(dir))),
    };

    // Restore rewrites MANIFEST to the snapshot epoch. Delete stale WAL files
    // first so modulo-reused names cannot keep old headers from previous runs.
    clear_wal_files(dir)?;
    init_wal_files(dir, None, None, epoch_rotate_len)?;

    let mut snap_store_meta = StoreBackupMeta::default();
    if snap_epoch > 0 {
        snap_store_meta.merge_from_bytes(snap_meta.chunk()).unwrap();
        assert_eq!(snap_epoch, snap_store_meta.get_manifest().epoch_id);
        restore_all_raft_logs_with_snap_rlog_file(keyspace_id, &snap_store_meta, dir, &snap_rlog)?;
    }
    let dur_restore_rlogs = start_time.saturating_elapsed();

    info!("{} manifest file path: {:?}", store_id, manifest_path(dir); "keyspace" => ?keyspace_id);
    let manifest_file = OpenOptions::new()
        .create(true)
        .truncate(true)
        .write(true)
        .open(manifest_path(dir))
        .ctx("open manifest")?;
    if let Some(keyspace_id) = keyspace_id {
        let before = snap_store_meta.get_manifest().peers.len();
        filter_manifest_peers_for_keyspace(snap_store_meta.mut_manifest(), keyspace_id);
        let after = snap_store_meta.get_manifest().peers.len();
        info!("{} filter manifest peers: {} -> {}", store_id, before, after; "keyspace" => keyspace_id);
    }
    persist_change_set(&manifest_file, 0, snap_store_meta.get_manifest())
        .ctx("persist change set")?;
    let dur_persist_manifest = start_time.saturating_elapsed() - dur_restore_rlogs;

    info!("{} restore rfengine", store_id;
        "restore_rlogs" => ?dur_restore_rlogs, "persist_manifest" => ?dur_persist_manifest);
    Ok(snap_store_meta.get_manifest().epoch_id)
}

fn filter_manifest_peers_for_keyspace(cs: &mut rfenginepb::ChangeSet, keyspace_id: u32) {
    let peers = cs.take_peers();
    peers
        .into_iter()
        .filter(|peer| {
            peer.region_id == 0
                || peer.peer_id == 0
                || get_keyspace_id_from_peer(peer).is_some_and(|x| x == keyspace_id)
        })
        .for_each(|peer| cs.mut_peers().push(peer));
}

fn clear_wal_files(dir: &Path) -> Result<()> {
    if !dir.exists() {
        return Ok(());
    }

    for entry in fs::read_dir(dir).with_ctx(|| format!("read dir {}", dir.display()))? {
        let entry = entry.with_ctx(|| format!("read dir entry in {}", dir.display()))?;
        if !entry
            .file_type()
            .with_ctx(|| format!("read file type {}", entry.path().display()))?
            .is_file()
        {
            continue;
        }

        let path = entry.path();
        // `lightweight_restore` is an offline procedure and should be run
        // against a directory that is not concurrently used by an engine
        // process. We delete only rfengine's numeric WAL slot files (e.g.
        // `0.wal`, `1.wal`) to avoid surprising users by deleting unrelated
        // `.wal` files that happen to live in the same directory.
        if path.extension().is_none_or(|ext| ext != "wal") {
            continue;
        }
        let Some(stem) = path.file_stem().and_then(|stem| stem.to_str()) else {
            continue;
        };
        if stem.is_empty() || !stem.as_bytes().iter().all(|b| b.is_ascii_digit()) {
            continue;
        }
        fs::remove_file(&path).with_ctx(|| format!("remove wal file {}", path.display()))?;
    }
    Ok(())
}

pub(crate) fn init_wal_files(
    dir: &Path,
    wal_sync_dir: Option<&PathBuf>,
    wal_secondary_dir: Option<&PathBuf>,
    epoch_rotate_len: usize,
) -> Result<()> {
    if !dir.exists() {
        create_dir_all(dir)?;
    }
    open_wal_files(dir, epoch_rotate_len)?;
    if let Some(wal_secondary_dir) = wal_secondary_dir {
        if !wal_secondary_dir.exists() {
            create_dir_all(wal_secondary_dir)?;
        }
    }
    if wal_sync_dir.is_none() {
        return Ok(());
    }
    let wal_sync_dir = wal_sync_dir.unwrap();
    if !wal_sync_dir.exists() {
        create_dir_all(wal_sync_dir)?;
    }
    // upgrade_mark file is used to make the upgrade procedure idempotent.
    // In case the upgrade process is interrupted, we can resume it later.
    let upgrade_mark_file = upgrade_mark_file_path(dir);
    if !upgrade_mark_file.exists() {
        if all_wal_files_exists(wal_sync_dir.as_path(), epoch_rotate_len) {
            return Ok(()); // already upgraded.
        }
        File::create(upgrade_mark_file.as_path())?;
    }
    copy_wal_files(dir, wal_sync_dir.as_path(), epoch_rotate_len)?;
    fs::remove_file(upgrade_mark_file.as_path())?;
    file_system::sync_dir(dir)?;
    Ok(())
}

fn upgrade_mark_file_path(dir: &Path) -> PathBuf {
    dir.join("upgrade_mark")
}

fn open_wal_files(dir: &Path, epoch_rotate_len: usize) -> Result<()> {
    // create epoch_rotate_len wal files and always reuse them, so we never need to
    // sync dir on writer thread.
    for i in 0..epoch_rotate_len {
        let file_path = wal_file_path(dir, i);
        let _ = open_direct_file(&file_path, true)?;
    }
    file_system::sync_dir(dir)?;
    Ok(())
}

pub(crate) fn wal_file_path(dir: &Path, idx: usize) -> PathBuf {
    dir.join(format!("{}.wal", idx))
}

fn all_wal_files_exists(dir: &Path, epoch_rotate_len: usize) -> bool {
    for i in 0..epoch_rotate_len {
        if !wal_file_path(dir, i).exists() {
            return false;
        }
    }
    true
}

fn copy_wal_files(dir: &Path, wal_sync_dir: &Path, epoch_rotate_len: usize) -> Result<()> {
    for i in 0..epoch_rotate_len {
        let src_file_path = wal_file_path(dir, i);
        let dst_file_path = wal_file_path(wal_sync_dir, i);
        fs::copy(src_file_path, dst_file_path)?;
    }
    file_system::sync_dir(wal_sync_dir)?;
    Ok(())
}

pub fn load_store_ident(rf: &RfEngine) -> Option<StoreIdent> {
    let val = rf.get_state(0, STORE_IDENT_KEY);
    val.as_ref()?;
    let mut ident = StoreIdent::new();
    ident.merge_from_bytes(val.unwrap().chunk()).unwrap();
    Some(ident)
}

pub fn save_store_ident(rf: &RfEngine, store_ident: &StoreIdent) {
    let val = store_ident.write_to_bytes().unwrap();
    let mut wb = WriteBatch::new();
    wb.set_state(0, 0, 0, STORE_IDENT_KEY, &val);
    rf.write(wb).unwrap();
}

#[derive(Debug, Clone, Default, PartialEq)]
pub(crate) struct PeerMeta {
    pub(crate) region_id: u64,
    pub(crate) truncated_idx: u64,
    pub(crate) keyspace_id: u32,
    pub(crate) has_region_meta: bool,
    pub(crate) states: BTreeMap<Bytes, Bytes>,
    pub(crate) states_encoded_len: usize,
}

const ENTRY_BASE_LEN: usize = 2 /* key len */ + 4 /* value len */;

impl PeerMeta {
    pub(crate) fn new(region_id: u64, keyspace_id: u32) -> Self {
        Self {
            region_id,
            keyspace_id,
            ..Default::default()
        }
    }

    pub fn set_state(&mut self, key: &[u8], val: &[u8]) {
        self.set_state_bytes(Bytes::copy_from_slice(key), Bytes::copy_from_slice(val))
    }

    pub fn set_state_bytes(&mut self, key: Bytes, val: Bytes) {
        if key[0] == REGION_META_KEY_BYTE && !val.is_empty() {
            self.has_region_meta = true;
        }
        let key_len = key.len();
        self.states_encoded_len += key_len + val.len() + ENTRY_BASE_LEN;
        let old = self.states.insert(key, val);
        if let Some(old) = old {
            self.states_encoded_len -= key_len + old.len() + ENTRY_BASE_LEN;
        }
    }

    pub fn remove_state(&mut self, key: &[u8]) {
        if let Some(old) = self.states.remove(key) {
            self.states_encoded_len -= key.len() + old.len() + ENTRY_BASE_LEN;
            if self.has_region_meta && key[0] == REGION_META_KEY_BYTE {
                self.has_region_meta = self.get_latest_state(REGION_META_KEY_PREFIX).is_some()
            }
        }
    }

    pub fn get_state(&self, key: &[u8]) -> Option<&[u8]> {
        self.states.get(key).map(|v| v.chunk())
    }

    pub fn get_state_bytes(&self, key: &[u8]) -> Option<Bytes> {
        self.states.get(key).cloned()
    }

    pub fn get_latest_state(&self, key_prefix: &[u8]) -> Option<&[u8]> {
        self.states
            .iter()
            .rev()
            .find(|(k, _)| k.starts_with(key_prefix))
            .map(|(_, v)| v.chunk())
    }

    pub fn get_latest_peer_state(&self) -> Option<RegionLocalState> {
        let bin = self.get_latest_state(REGION_META_KEY_PREFIX)?;
        let mut region_local_state = RegionLocalState::new();
        region_local_state.merge_from_bytes(bin).unwrap();
        Some(region_local_state)
    }

    pub fn get_keyspace_id(&self) -> Option<u32> {
        if self.keyspace_id > 0 {
            return Some(self.keyspace_id);
        }
        if !self.has_region_meta {
            return None;
        }
        self.get_latest_peer_state().map(|local_state| {
            ApiV2::get_u32_keyspace_id_by_key(local_state.get_region().get_start_key())
                .unwrap_or_default()
        })
    }

    pub(crate) fn merge(&mut self, other: &PeerMeta, keep_empty: bool) {
        assert_eq!(self.region_id, other.region_id);
        if self.truncated_idx < other.truncated_idx {
            self.truncated_idx = other.truncated_idx;
        }
        if let Some(keyspace_id) = other.get_keyspace_id() {
            self.keyspace_id = keyspace_id;
        }
        for (key, val) in &other.states {
            if keep_empty || !val.is_empty() {
                self.set_state_bytes(key.clone(), val.clone());
            } else {
                self.remove_state(key)
            };
        }
    }
}

/// `PeerData` contains region data and state in memory.
#[derive(Clone, Default)]
pub(crate) struct PeerData {
    pub(crate) peer_id: u64,
    pub(crate) meta: PeerMeta,
    pub(crate) raft_logs: RaftLogs,
}

impl Deref for PeerData {
    type Target = PeerMeta;

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

impl DerefMut for PeerData {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.meta
    }
}

impl PeerData {
    pub(crate) fn new(peer_id: u64, region_id: u64, keyspace_id: u32) -> Self {
        Self {
            peer_id,
            meta: PeerMeta::new(region_id, keyspace_id),
            ..Default::default()
        }
    }

    pub(crate) fn get(&self, index: u64) -> Option<eraftpb::Entry> {
        self.raft_logs.get(index)
    }

    pub(crate) fn term(&self, index: u64) -> Option<u64> {
        self.raft_logs.get(index).map(|e| e.term)
    }

    pub(crate) fn get_state(&self, key: &[u8]) -> Option<&Bytes> {
        self.states.get(key)
    }

    pub(crate) fn apply(&mut self, batch: &PeerBatch) -> Vec<RaftLogBlock> {
        debug_assert_eq!(self.peer_id, batch.peer_id);
        let mut truncated_blocks = vec![];
        for op in &batch.raft_logs {
            let truncated = self.raft_logs.append(self.peer_id, op.clone());
            if !truncated.is_empty() {
                truncated_blocks.extend(truncated);
            }
        }
        let truncated_index = batch.truncated_idx;
        if self.truncated_idx < truncated_index {
            self.truncated_idx = truncated_index;
            truncated_blocks.extend(self.raft_logs.truncate(truncated_index));
        }
        if self.truncated_idx == TRUNCATE_ALL_INDEX
            && truncated_index > 0
            && truncated_index != TRUNCATE_ALL_INDEX
        {
            warn!(
                "region: {} peer:{} restore truncate all index to index {}",
                self.region_id, self.peer_id, truncated_index,
            );
            self.truncated_idx = truncated_index;
        }
        self.meta.merge(&batch.meta, false);
        truncated_blocks
    }

    pub(crate) fn get_stats(&self) -> PeerStats {
        let size = self.raft_logs.size();
        let first_idx = self.raft_logs.first_index();
        let last_idx = self.raft_logs.last_index();
        let num_logs = if last_idx != 0 {
            (last_idx - first_idx + 1) as usize
        } else {
            0
        };
        let num_logs_offloaded = self.get_num_logs_offloaded() as usize;
        PeerStats {
            peer_id: self.peer_id,
            region_id: self.meta.region_id,
            keyspace_id: self.meta.keyspace_id,
            size,
            num_logs,
            num_logs_offloaded,
            num_states: self.meta.states.len(),
            first_idx,
            last_idx,
            truncated_idx: self.meta.truncated_idx,
        }
    }

    // Returns the number of raft log entries that have been offloaded.
    //
    // Caveat: if all raft logs have been offloaded (which can only happen if
    // TiKV restarts and the peer has no new writes), the returned value will be
    // 0, which is inaccurate.
    fn get_num_logs_offloaded(&self) -> u64 {
        self.raft_logs.first_index().saturating_sub(
            self.meta
                .truncated_idx
                .max(RAFT_INIT_LOG_INDEX)
                .saturating_add(1),
        )
    }

    /// Returns the index such that truncating logs up to this index will reduce
    /// its total raft log usage (including in-memory and offloaded entries) to
    /// below the specified size limit.
    pub(crate) fn index_to_truncate_to_size(&self, size_limit: usize) -> u64 {
        let in_mem_bytes = self.raft_logs.size();
        // If in‑memory logs alone exceed the limit, no need to check offloaded
        // entries.
        if in_mem_bytes >= size_limit {
            return self.raft_logs.index_to_truncate_to_size(size_limit);
        }

        // Check if we need to truncate the offloaded raft logs. We estimate the
        // offloaded raft log size based on the average raft entry size in
        // memory.
        let in_mem_cnt = self.raft_logs.len();
        let offloaded_cnt = self.get_num_logs_offloaded();
        if in_mem_cnt == 0 || offloaded_cnt == 0 || in_mem_bytes == 0 {
            return 0;
        }
        let avg_entry_size = in_mem_bytes.div_ceil(in_mem_cnt);
        let offloaded_cnt_allowed = (size_limit - in_mem_bytes) / avg_entry_size;
        let to_truncate = offloaded_cnt.saturating_sub(offloaded_cnt_allowed as u64);
        if to_truncate == 0 {
            return 0;
        }

        info!(
            "peer {}: truncating offloaded entries, \
            in_mem_bytes {}, in_mem_cnt {}, offloaded_cnt {}, to_truncate {}",
            self.peer_id, in_mem_bytes, in_mem_cnt, offloaded_cnt, to_truncate
        );
        self.meta.truncated_idx.max(RAFT_INIT_LOG_INDEX) + to_truncate
    }
}

#[derive(Default, Serialize, Deserialize, Debug)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct EngineStats {
    pub total_mem_size: usize,
    pub total_mem_entries: usize,
    pub num_files: usize,
    pub disk_size: u64,
    pub pending_compaction_wals: u8,
    pub top_10_size_peers: Vec<PeerStats>,
    pub total_num_logs: u64,
}

#[derive(Default, Serialize, Deserialize, Debug)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct EngineDfsStats {
    pub uploaded_bytes: u64,
    pub requests: u64,
}

#[derive(Default, Serialize, Deserialize, Debug, PartialEq)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct PeerStats {
    pub peer_id: u64,
    pub region_id: u64,
    pub keyspace_id: u32,
    pub size: usize,
    pub num_logs: usize,
    pub num_logs_offloaded: usize,
    pub num_states: usize,
    pub first_idx: u64,
    pub last_idx: u64,
    pub truncated_idx: u64,
}

pub struct PeerTag {
    pub engine_id: u64,
    pub region_id: u64,
}

impl PeerTag {
    pub fn new(engine_id: u64, region_id: u64) -> Self {
        Self {
            engine_id,
            region_id,
        }
    }
}

impl Display for PeerTag {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}:{}", self.engine_id, self.region_id)
    }
}

fn lock_conflict_message(dir: &Path) -> String {
    format!(
        "rfengine lock failed for {}, maybe already used by another process",
        dir.display()
    )
}

fn try_acquire_lock_file(dir: &Path) -> Result<Option<fslock::LockFile>> {
    let lock_path = dir.join("LOCK");
    let mut lock = fslock::LockFile::open(&lock_path)?;
    if !lock.try_lock()? {
        return Ok(None);
    }
    Ok(Some(lock))
}

fn now_unix_millis() -> i64 {
    fail::fail_point!("rfengine::now_unix_millis", |v| {
        v.expect("rfengine::now_unix_millis requires a return value")
            .parse::<i64>()
            .expect("rfengine::now_unix_millis return value must be an integer")
    });
    Utc::now().naive_utc().timestamp_millis()
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, fs::OpenOptions, os::unix::prelude::FileExt, time::Duration};

    use ::test_util::eventually;
    use eraftpb::{Entry, EntryType};
    use matches::assert_matches;
    use protobuf::Message;
    use tikv_util::config::ReadableSize;

    use super::*;
    use crate::{
        log_batch::{RAFT_LOG_BLOCK_CAP, RaftLogOp},
        test_util::{
            init_logger, make_log_data, make_region_state, make_state_kv, new_raft_entry, try_wait,
        },
    };

    #[test]
    fn test_rfengine() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let cfg = Config::new(128 * 1024_usize);
        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
        assert_eq!(engine.is_empty(), true);

        let init_stats = engine.get_engine_stats();
        assert_eq!(init_stats.total_mem_size, 0);
        assert_eq!(init_stats.total_mem_entries, 0);
        assert!(init_stats.num_files > 0);
        assert!(init_stats.disk_size > 0);
        assert_eq!(init_stats.pending_compaction_wals, 0);
        assert_eq!(init_stats.top_10_size_peers.len(), 0);

        let mut wb = WriteBatch::new();
        for peer_id in 1..=10_u64 {
            let (key, val) = make_state_kv(3, 1);
            let region_id = peer_id + 1;
            wb.set_state(peer_id, region_id, 0, key.chunk(), val.chunk());
        }
        engine.write(wb).unwrap();
        assert_eq!(engine.is_empty(), false);

        let mut truncated_regions = vec![];
        for idx in 1..=1050_u64 {
            let mut wb = WriteBatch::new();
            for peer_id in 1..=10_u64 {
                let region_id = peer_id + 1;
                let keyspace_id = 1;
                if peer_id == 1 {
                    if idx > 100 && idx < 900 {
                        continue;
                    } else if idx == 900 {
                        wb.truncate_raft_log(peer_id, region_id, keyspace_id, 899);
                    }
                }
                wb.append_raft_log(peer_id, region_id, keyspace_id, &make_log_data(idx, 128));
                let (key, val) = make_state_kv(1, idx);
                wb.set_state(peer_id, region_id, keyspace_id, key.chunk(), val.chunk());
                if idx % 100 == 0 && peer_id != 1 {
                    truncated_regions.push((peer_id, region_id, keyspace_id, idx - 100));
                    wb.truncate_raft_log(peer_id, region_id, keyspace_id, idx - 100);
                }
            }
            engine.write(wb).unwrap();
        }
        assert_eq!(engine.peers.peers.len(), 10);
        let wal_cnt = engine
            .dir
            .read_dir()
            .unwrap()
            .filter(|p| {
                p.as_ref()
                    .unwrap()
                    .path()
                    .extension()
                    .is_some_and(|e| e == "wal")
            })
            .count();
        assert_eq!(wal_cnt, 4);

        let mut old_entries_map = HashMap::new();
        let peers = engine.peers.peers.pin();
        for (peer_id, peer_ref) in peers.iter() {
            let peer_data = peer_ref.read().unwrap();
            assert_eq!(*peer_id, peer_data.peer_id);
            old_entries_map.insert(peer_data.peer_id, peer_data.clone());
        }
        drop(peers);
        assert_eq!(old_entries_map.len(), 10);
        engine.stop_worker(true);

        for _ in 0..2 {
            let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
            let mut wb = WriteBatch::new();
            for &(peer_id, region_id, keyspace_id, truncated_idx) in truncated_regions.iter() {
                wb.truncate_raft_log(peer_id, region_id, keyspace_id, truncated_idx);
            }
            engine.apply(&wb);
            engine.iterate_all_states(false, |peer_id, _, key, _| {
                let old_region_data = old_entries_map.get(&peer_id).unwrap();
                assert!(old_region_data.get_state(key).is_some());
                true
            });
            assert_eq!(engine.peers.peers.len(), 10);
            let peers = engine.peers.peers.pin();
            for (peer_id, new_data_ref) in peers.iter() {
                let new_data = new_data_ref.read().unwrap();
                let old_data = old_entries_map.get(peer_id).unwrap();
                assert_eq!(
                    old_data.raft_logs.first_index(),
                    engine.get_truncated_index(*peer_id).unwrap() + 1,
                    "peer: {}",
                    *peer_id,
                );
                assert_eq!(
                    old_data.raft_logs.last_index(),
                    engine.get_last_index(*peer_id).unwrap(),
                );
                for i in new_data.raft_logs.first_index()..=new_data.raft_logs.last_index() {
                    let entry = new_data.raft_logs.get(i).unwrap();
                    assert_eq!(
                        old_data.get(entry.index).unwrap().data.chunk(),
                        entry.data.chunk()
                    );
                }
            }
            engine.stop_worker(true);
        }
    }

    #[test]
    fn test_open_validates_rlog_cache_configuration() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::new(128 * 1024_usize);
        cfg.rlog_cache_size_threshold = ReadableSize(0);

        match RfEngine::open(tmp_dir.path(), &cfg, None, None) {
            Err(Error::Other(msg)) => assert_eq!(
                msg,
                "invalid config: rlog_cache_size_threshold must be non-zero when rlog_cache_capacity is non-zero"
            ),
            Ok(_) => panic!("expected open to reject zero rlog_cache_size_threshold"),
            Err(err) => panic!("unexpected error: {:?}", err),
        }

        let tmp_dir = tempfile::tempdir().unwrap();
        cfg.rlog_cache_capacity = ReadableSize(0);
        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None)
            .expect("zero capacity should disable the rlog cache");
        engine.stop_worker(true);
    }

    #[test]
    fn test_open_rejects_zero_compaction_knobs() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();

        let mut cfg = Config::new(128 * 1024_usize);
        cfg.compaction_concurrency = 0;
        match RfEngine::open(tmp_dir.path(), &cfg, None, None) {
            Err(Error::Other(msg)) => {
                assert_eq!(
                    msg,
                    "invalid config: compaction_concurrency must be non-zero"
                )
            }
            Ok(_) => panic!("expected open to reject zero compaction_concurrency"),
            Err(err) => panic!("unexpected error: {:?}", err),
        }

        let mut cfg = Config::new(128 * 1024_usize);
        cfg.compaction_batch_size = 0;
        match RfEngine::open(tmp_dir.path(), &cfg, None, None) {
            Err(Error::Other(msg)) => {
                assert_eq!(
                    msg,
                    "invalid config: compaction_batch_size must be non-zero"
                )
            }
            Ok(_) => panic!("expected open to reject zero compaction_batch_size"),
            Err(err) => panic!("unexpected error: {:?}", err),
        }
    }

    #[test]
    fn test_region_data() {
        init_logger();
        let mut region_data = PeerData::new(1, 2, 1);

        let mut region_batch = PeerBatch::new(1, 2, 1);
        for i in 1..=5 {
            region_batch.append_raft_log(RaftLogOp::new(&new_raft_entry(
                EntryType::EntryNormal,
                i,
                i,
                b"data",
                0,
            )));
        }
        assert!(region_data.apply(&region_batch).is_empty());
        for i in 1..=5 {
            assert_eq!(
                region_data.get(i).unwrap(),
                region_batch.raft_logs[(i - 1) as usize].to_entry()
            );
            assert_eq!(region_data.term(i).unwrap(), i);
        }
        assert!(region_data.get(6).is_none());
        let region_stats = region_data.get_stats();
        assert_eq!(
            region_stats,
            PeerStats {
                peer_id: 1,
                region_id: 2,
                keyspace_id: 1,
                size: 20,
                num_logs: 5,
                num_logs_offloaded: 0,
                num_states: 0,
                first_idx: 1,
                last_idx: 5,
                truncated_idx: 0,
            }
        );

        region_batch = PeerBatch::new(1, 2, 1);
        region_batch.truncate(5);
        region_batch.set_state(b"k1", b"v1");
        region_batch.set_state(b"k2", b"v2");
        let truncated = region_data.apply(&region_batch);
        assert_eq!(truncated.len(), 1);
        assert_eq!(truncated[0].first_index(), 1);
        assert_eq!(truncated[0].last_index(), 5);
        for i in 1..=5 {
            assert!(region_data.get(i).is_none());
        }
        assert_eq!(region_data.get_state(b"k1"), Some(&b"v1".to_vec().into()));
        assert_eq!(region_data.get_state(b"k2"), Some(&b"v2".to_vec().into()));
        let region_stats = region_data.get_stats();
        assert_eq!(
            region_stats,
            PeerStats {
                peer_id: 1,
                region_id: 2,
                keyspace_id: 1,
                size: 0,
                num_logs: 0,
                num_logs_offloaded: 0,
                num_states: 2,
                first_idx: 0,
                last_idx: 0,
                truncated_idx: 5,
            }
        );

        region_batch = PeerBatch::new(1, 2, 1);
        region_batch.truncate(5);
        region_batch.set_state(b"k1", b"");
        assert!(region_data.apply(&region_batch).is_empty());
        assert!(region_data.get_state(b"k1").is_none());
        assert_eq!(region_data.get_state(b"k2"), Some(&b"v2".to_vec().into()));

        region_batch = PeerBatch::new(1, 2, 1);
        region_batch.truncate(100);
        assert!(region_data.apply(&region_batch).is_empty());
        assert_eq!(region_data.truncated_idx, 100);
    }

    #[test]
    fn test_rfengine_basic() {
        init_logger();
        const STATE_PREFIX: u8 = b'p';

        let dir = tempfile::tempdir().unwrap();
        let cfg = Config::new(128 * 1024);
        let engine = RfEngine::open(dir.path(), &cfg, None, None).unwrap();

        // Write 10 logs and states to 2 region.
        let mut data_map = HashMap::new();
        let mut wb = WriteBatch::new();
        for peer_id in 1..=2 {
            let region_id = peer_id + 1;
            let keyspace_id = 1;
            for i in 1..=10 {
                let entry = new_raft_entry(EntryType::EntryNormal, peer_id, i, b"data", 0);
                let (state_key, state_val) = (&[STATE_PREFIX, i as u8], &[i as u8]);
                wb.append_raft_log(peer_id, region_id, keyspace_id, &entry);
                wb.set_state(peer_id, region_id, keyspace_id, state_key, state_val);

                let (entries, states) = data_map
                    .entry(peer_id)
                    .or_insert_with(|| (vec![], BTreeMap::new()));
                entries.push(entry);
                states.insert(state_key.to_vec(), state_val.to_vec());
            }
        }
        engine.write(wb).unwrap();

        assert_eq!(engine.get_term(1, 1), Some(1));
        assert_eq!(engine.get_term(1, 11), None);
        assert_eq!(engine.get_term(3, 1), None);
        assert_eq!(engine.get_last_index(1), Some(10));
        assert_eq!(engine.get_last_index(3), None);

        // Test `get_entry` and `get_state`.
        for (&peer_id, (entries, states)) in &data_map {
            entries.iter().for_each(|entry| {
                assert_eq!(entry, &engine.get_raft_entry(peer_id, entry.index).unwrap(),);
            });
            states
                .iter()
                .for_each(|(key, val)| assert_eq!(val, &engine.get_state(peer_id, key).unwrap()));
        }
        assert!(engine.get_raft_entry(1, 11).is_none());
        assert!(engine.get_raft_entry(3, 1).is_none());
        assert!(engine.get_state(1, b"k").is_none());

        // Test `fetch_entries_to`.
        let mut buf = vec![];
        for peer_id in 1..=2 {
            for low in 1..=10 {
                for high in low + 1..=11 {
                    assert_eq!(
                        engine
                            .fetch_raft_entries_to(peer_id, low, high, None, &mut buf, None)
                            .unwrap(),
                        (high - low) as usize
                    );
                    assert_eq!(
                        data_map.get(&peer_id).unwrap().0[(low - 1) as usize..(high - 1) as usize],
                        buf
                    );
                    buf.clear();
                }
            }
        }
        // Test `fetch_entries_to` should push logs to the buf.
        let peer1_entries = &data_map.get(&1).unwrap().0;
        for i in 1..=10 {
            assert_eq!(
                engine
                    .fetch_raft_entries_to(1, i, i + 1, None, &mut buf, None)
                    .unwrap(),
                1
            );
            assert_eq!(buf, peer1_entries[..i as usize]);
        }
        assert_matches!(
            engine.fetch_raft_entries_to(1, 11, 12, None, &mut buf, None),
            Err(Error::EntriesUnavailable)
        );
        // Test `fetch_entries_to` limits size.
        let mut max_size = 0;
        for (i, entry) in peer1_entries.iter().enumerate() {
            buf.clear();
            max_size += entry.compute_size();
            assert_eq!(
                engine
                    .fetch_raft_entries_to(1, 1, 11, Some(max_size as usize), &mut buf, None)
                    .unwrap(),
                i + 1
            );
            assert_eq!(buf, peer1_entries[..=i]);
        }

        // Test fetch empty logs.
        buf.clear();
        assert_eq!(
            engine
                .fetch_raft_entries_to(1, 1, 1, None, &mut buf, None)
                .unwrap(),
            0
        );
        assert!(buf.is_empty());

        // Test `get_last_state_with_prefix`
        assert_eq!(
            engine
                .get_last_state_with_prefix(1, &[STATE_PREFIX])
                .unwrap(),
            [10_u8].as_slice()
        );
        assert!(
            engine
                .get_last_state_with_prefix(1, &[STATE_PREFIX + 1])
                .is_none()
        );

        // Test `iterate_region_states`
        for desc in [false, true] {
            let mut expect_index = if desc { 10 } else { 1 };
            engine.iterate_peer_states(1, desc, |k, v| {
                assert_eq!(k.chunk(), &[STATE_PREFIX, expect_index]);
                assert_eq!(v.chunk(), &[expect_index]);
                if desc {
                    expect_index -= 1;
                } else {
                    expect_index += 1;
                }
                true
            });
            assert_eq!(expect_index, if desc { 0 } else { 11 });
        }

        // Test `iterate_all_states`
        for desc in [false, true] {
            let mut count = 0;
            engine.iterate_all_states(desc, |id, _, k, v| {
                assert_eq!(v, data_map.get(&id).unwrap().1.get(k).unwrap());
                count += 1;
                true
            });
            assert_eq!(count, 20);
        }
        // Test `iterate_all_states` breaks.
        let mut count = 0;
        engine.iterate_all_states(false, |_, _, _, _| {
            count += 1;
            false
        });
        assert_eq!(count, 2);

        // Test `add_dependent` and `remove_dependent`.
        engine.add_dependent(1, 1);
        let dependants = engine.dependants.pin();
        assert!(dependants.get(&1).unwrap().read().unwrap().contains(&1));
        engine.remove_dependent(1, 1);
        assert!(!dependants.get(&1).unwrap().read().unwrap().contains(&1));
    }

    #[test]
    fn test_rfengine_wal() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let wal_size = 128 * 1024_usize;
        let dir_path = tmp_dir.path();
        let cfg = Config::new(wal_size);
        let mut check_cfg = cfg.clone();
        check_cfg.disable_compaction = true;
        let engine = RfEngine::open(dir_path, &cfg, None, None).unwrap();
        let mut wb = WriteBatch::new();
        for peer_id in 1..=10_u64 {
            let (key, val) = make_state_kv(3, 1);
            let region_id = peer_id + 1;
            wb.set_state(peer_id, region_id, 0, key.chunk(), val.chunk());
        }
        engine.write(wb).unwrap();
        for idx in 1..=1050_u64 {
            let mut wb = WriteBatch::new();
            for peer_id in 1..=10_u64 {
                let region_id = peer_id + 1;
                let keyspace_id = 1;
                wb.append_raft_log(peer_id, region_id, keyspace_id, &make_log_data(idx, 128));
                let (key, val) = make_state_kv(1, idx);
                wb.set_state(peer_id, region_id, keyspace_id, key.chunk(), val.chunk());
            }
            engine.write(wb).unwrap();
        }
        assert_eq!(engine.peers.peers.len(), 10);
        engine.stop_worker(true);
        for _ in 0..2 {
            let engine = RfEngine::open(dir_path, &check_cfg, None, None).unwrap();
            assert_eq!(engine.peers.peers.len(), 10);
            engine.stop_worker(true);
        }
        let compacted_epoch = engine.compacted_epoch.load(Ordering::Relaxed);
        let current_epoch = {
            let writer = engine.writer.lock().unwrap();
            writer.get_epoch_id()
        };
        {
            let mut it =
                WalIterator::new(dir_path, current_epoch + 1, cfg.epoch_rotate_len).unwrap();
            let Error::Corruption {
                msg: _,
                epoch_id: _,
                offset,
                data: _,
            } = it.check_wal_header().unwrap_err()
            else {
                panic!("expected corruption error");
            };
            // header epoch mismatch error offset should be 0
            assert_eq!(offset, 0);
        }
        for ep in compacted_epoch + 1..=current_epoch {
            let filename = wal_file_name(dir_path, ep, cfg.epoch_rotate_len);
            let mut it = WalIterator::new(dir_path, ep, cfg.epoch_rotate_len).unwrap();
            it.check_wal_header().unwrap();
            let mut offsets = vec![it.offset];
            loop {
                match it.read_batch() {
                    Err(err) => {
                        if let Error::Eof = err {
                            break;
                        }
                        panic!("{:?}", err);
                    }
                    Ok(_data) => offsets.push(it.offset),
                }
            }
            offsets.pop().unwrap();
            for (idx, offset) in offsets.iter().enumerate() {
                if idx == 0 || idx == offsets.len() / 2 || idx == offsets.len() - 1 {
                    for pos in &[0, 4, 8, 12] {
                        let fd = OpenOptions::new()
                            .read(true)
                            .write(true)
                            .open(filename.as_path())
                            .unwrap();
                        let mut buf = [0u8; 4096];
                        fd.read_exact_at(&mut buf, *offset).unwrap();
                        if buf[*pos] == 255 {
                            continue;
                        }
                        buf[*pos] += 1;
                        fd.write_all_at(buf.as_ref(), *offset).unwrap();
                        fd.sync_data().unwrap();
                        let open_engine = RfEngine::open(dir_path, &check_cfg, None, None);
                        // RfEngine can auto recover from corruption for the last epoch wal
                        // corruption.
                        assert!(if ep == current_epoch {
                            open_engine.is_ok()
                        } else {
                            open_engine.is_err()
                        });
                        buf[*pos] -= 1;
                        fd.write_all_at(buf.as_ref(), *offset).unwrap();
                        fd.sync_data().unwrap();
                    }
                }
            }
        }
    }

    #[test]
    fn test_truncate_all_logs() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let wal_size = 4096 * 10;
        let cfg = Config::new(wal_size);
        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
        {
            let mut wb = WriteBatch::new();
            let (key, val) = make_region_state(10, 42);
            wb.set_state(1, 2, 1, &key, &val);
            engine.write(wb).unwrap();
        }
        for i in 1..=50 {
            let mut wb = WriteBatch::new();
            wb.append_raft_log(1, 2, 1, &make_log_data(i, 128));
            engine.write(wb).unwrap();
        }

        // Truncate all index.
        let mut wb = WriteBatch::new();
        wb.truncate_raft_log(1, 2, 1, TRUNCATE_ALL_INDEX);
        engine.write(wb).unwrap();

        // Write more batch to trigger WAL compaction.
        {
            let mut wb = WriteBatch::new();
            let (key, val) = make_region_state(11, 43);
            wb.set_state(2, 3, 0, &key, &val);
            engine.write(wb).unwrap();
        }
        for i in 1..=10 {
            let mut wb = WriteBatch::new();
            wb.append_raft_log(2, 3, 0, &make_log_data(i, wal_size));
            engine.write(wb).unwrap();
        }

        // Check no file of peer 1 left.
        wait_for_rlogs_truncated(&engine, 1, 10);
    }

    #[test]
    fn test_init_wal_files() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let epoch_rotate_len = 4;
        init_wal_files(tmp_dir.path(), None, None, epoch_rotate_len).unwrap();
        let check_file_exists = |path: &Path| {
            for idx in 0..epoch_rotate_len {
                assert!(wal_file_path(path, idx).exists());
            }
        };
        check_file_exists(tmp_dir.path());

        let file_contents: Vec<String> = (0..epoch_rotate_len)
            .map(|i| format!("wal {}", i))
            .collect();
        let write_files = |path: &Path| {
            for idx in 0..epoch_rotate_len {
                let wal_file_path = wal_file_path(path, idx);
                fs::write(wal_file_path.as_path(), file_contents[idx].as_bytes()).unwrap();
            }
        };
        write_files(tmp_dir.path());
        File::create(manifest_path(tmp_dir.path())).unwrap();

        // upgrade to use wal_sync_dir
        let wal_sync_dir = tmp_dir.path().join("wal_sync_dir");
        init_wal_files(tmp_dir.path(), Some(&wal_sync_dir), None, epoch_rotate_len).unwrap();
        assert!(!upgrade_mark_file_path(tmp_dir.path()).exists());
        let check_files = || {
            for idx in 0..epoch_rotate_len {
                let async_wal_file_path = wal_file_path(tmp_dir.path(), idx);
                assert!(async_wal_file_path.exists());
                let sync_wal_file_path = wal_file_path(&wal_sync_dir, idx);
                assert!(sync_wal_file_path.exists());
                let data = fs::read_to_string(sync_wal_file_path.as_path()).unwrap();
                assert_eq!(data, file_contents[idx]);
            }
        };
        check_files();

        // simulate upgrade interrupted.
        write_files(tmp_dir.path());
        fs::remove_file(wal_file_path(wal_sync_dir.as_path(), 3)).unwrap();
        File::create(upgrade_mark_file_path(tmp_dir.path())).unwrap();

        // init_wal_files again should recover from the interrupted upgrade.
        init_wal_files(tmp_dir.path(), Some(&wal_sync_dir), None, epoch_rotate_len).unwrap();
        check_files();

        // wal_secondary_dir is created no matter if wal_sync_dir is provided.
        let wal_secondary_dir = tmp_dir.path().join("wal_secondary_dir");
        init_wal_files(
            tmp_dir.path(),
            Some(&wal_sync_dir),
            Some(&wal_secondary_dir),
            epoch_rotate_len,
        )
        .unwrap();
        assert!(wal_secondary_dir.exists());
        fs::remove_dir(wal_secondary_dir.as_path()).unwrap();
        init_wal_files(
            tmp_dir.path(),
            None,
            Some(&wal_secondary_dir),
            epoch_rotate_len,
        )
        .unwrap();
        assert!(wal_secondary_dir.exists());
    }

    #[test]
    fn test_lightweight_restore_clears_stale_wal_headers() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::new(16 * 1024);
        cfg.disable_compaction = true;

        // Create a stale WAL header for epoch 1.
        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
        engine.stop_worker(true);
        drop(engine);

        // Restore to snapshot epoch 8 so startup expects WAL epoch 9, which maps
        // to the same WAL file slot as epoch 1 when epoch_rotate_len is 4.
        let mut snap_store_meta = StoreBackupMeta::default();
        snap_store_meta.mut_manifest().set_epoch_id(8);
        snap_store_meta.raft_meta_start_off = 0;
        let snap_meta = Bytes::from(snap_store_meta.write_to_bytes().unwrap());

        let mut snap_rlog_meta = StoreRaftLogBackupMeta::default();
        snap_rlog_meta
            .mut_header()
            .set_compression_type(CompressionType::Lz4Compression.to());
        let snap_rlog = Bytes::from(snap_rlog_meta.write_to_bytes().unwrap());

        lightweight_restore(
            1,
            None,
            tmp_dir.path(),
            8,
            snap_meta,
            snap_rlog,
            cfg.epoch_rotate_len,
        )
        .unwrap();

        let wal_cnt = fs::read_dir(tmp_dir.path())
            .unwrap()
            .flatten()
            .filter(|entry| {
                let path = entry.path();
                if path.extension().is_none_or(|ext| ext != "wal") {
                    return false;
                }
                let Some(stem) = path.file_stem().and_then(|stem| stem.to_str()) else {
                    return false;
                };
                !stem.is_empty() && stem.as_bytes().iter().all(|b| b.is_ascii_digit())
            })
            .count();
        assert_eq!(wal_cnt, cfg.epoch_rotate_len);

        let reopened = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
        reopened.stop_worker(true);
    }

    #[rstest::rstest]
    #[case(false)]
    #[case::wal_sync(true)]
    fn test_rfengine_compact_force_restart(#[case] with_wal_sync: bool) {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let wal_sync_dir = tmp_dir.path().join("wal_sync");
        let mut cfg = Config::new(16 * 1024);
        if with_wal_sync {
            cfg.wal_sync_dir = wal_sync_dir.to_str().unwrap().to_owned();
        }
        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();

        let peer_id = 1;
        let region_id = 2;
        let keyspace_id = 1;
        let (key1, val1) = make_state_kv(b'a', 1);
        let key1_vec = key1.clone().freeze().to_vec();
        let val1_vec = val1.clone().freeze().to_vec();
        let mut wb = WriteBatch::new();
        wb.set_state(peer_id, region_id, keyspace_id, key1.chunk(), val1.chunk());
        engine.write(wb).unwrap();

        // Force stop before WAL rotate for stable interruption.
        engine.compact_force_stop.store(true, Ordering::SeqCst);

        let mut log_index = 1;
        let entry_size = (cfg.target_file_size.0 as usize / 2).max(1024);
        let start_epoch = engine.writer.lock().unwrap().get_epoch_id();
        while engine.writer.lock().unwrap().get_epoch_id() == start_epoch {
            let mut wb = WriteBatch::new();
            wb.append_raft_log(
                peer_id,
                region_id,
                keyspace_id,
                &make_log_data(log_index, entry_size),
            );
            engine.write(wb).unwrap();
            log_index += 1;
        }

        let (key2, val2) = make_state_kv(b'b', 2);
        let key2_vec = key2.clone().freeze().to_vec();
        let val2_vec = val2.clone().freeze().to_vec();
        let mut wb = WriteBatch::new();
        wb.set_state(peer_id, region_id, keyspace_id, key2.chunk(), val2.chunk());
        engine.write(wb).unwrap();

        engine.stop_worker(true);
        drop(engine);

        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
        let epoch_after_restart = engine.writer.lock().unwrap().get_epoch_id();
        assert_eq!(epoch_after_restart, start_epoch + 1);

        while engine.writer.lock().unwrap().get_epoch_id() == epoch_after_restart {
            let mut wb = WriteBatch::new();
            wb.append_raft_log(
                peer_id,
                region_id,
                keyspace_id,
                &make_log_data(log_index, entry_size),
            );
            engine.write(wb).unwrap();
            log_index += 1;
        }

        let (key3, val3) = make_state_kv(b'c', 3);
        let key3_vec = key3.clone().freeze().to_vec();
        let val3_vec = val3.clone().freeze().to_vec();
        let mut wb = WriteBatch::new();
        wb.set_state(peer_id, region_id, keyspace_id, key3.chunk(), val3.chunk());
        engine.write(wb).unwrap();

        assert!(
            try_wait(
                || engine.compacted_epoch.load(Ordering::SeqCst) >= epoch_after_restart,
                10
            ),
            "compact epoch {} not finished",
            epoch_after_restart
        );

        engine.stop_worker(false);
        drop(engine);

        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();
        assert_eq!(
            val1_vec.as_slice(),
            engine.get_state(peer_id, &key1_vec).unwrap().as_ref()
        );
        assert_eq!(
            val2_vec.as_slice(),
            engine.get_state(peer_id, &key2_vec).unwrap().as_ref()
        );
        assert_eq!(
            val3_vec.as_slice(),
            engine.get_state(peer_id, &key3_vec).unwrap().as_ref()
        );
        engine.stop_worker(true);
    }

    fn wait_for_rlogs_truncated(en: &RfEngine, peer_id: u64, seconds: usize) {
        let mut ok = false;
        let peer_id_str = format!("{:016x}", peer_id);

        let start_time = Instant::now_coarse();
        let timeout = Duration::from_secs(seconds as u64);
        while start_time.saturating_elapsed() < timeout {
            let read_dir = en.dir.read_dir().unwrap();
            let found = read_dir.into_iter().any(|entry| {
                let filename = entry.unwrap().file_name();
                let filename = filename.to_string_lossy();
                let parts: Vec<_> = filename.as_ref().split('_').collect();
                parts.len() == 3 && parts[0] == peer_id_str
            });
            if !found {
                ok = true;
                break;
            }
            std::thread::sleep(Duration::from_secs(1));
        }

        assert!(ok);
    }

    #[test]
    fn test_get_region_peer_map() {
        let dir = tempfile::tempdir().unwrap();
        let cfg = Config::new(128 * 1024);
        let engine = RfEngine::open(dir.path(), &cfg, None, None).unwrap();

        // ---------------------
        // Region 10: two peers (peer 1 -> older, peer 2 -> newer)
        // ---------------------
        // 1) Write something for peer 1, region 10:
        let mut wb = WriteBatch::new();
        wb.append_raft_log(
            1,
            10,
            1,
            &new_raft_entry(EntryType::EntryNormal, 1, 1, b"data", 0),
        );
        engine.write(wb).unwrap();

        // 2) Write something for peer 2, region 10 (the "newer" peer).
        let mut wb = WriteBatch::new();
        wb.append_raft_log(
            2,
            10,
            1,
            &new_raft_entry(EntryType::EntryNormal, 2, 1, b"data", 0),
        );
        engine.write(wb).unwrap();

        // At this point, region 10 should map to peer_id = 2 (newer).
        let map = engine.get_region_peer_map();
        assert_eq!(map.get(&10), Some(&2), "peer 2 should override peer 1");

        // 3) Truncate the newer peer to TRUNCATE_ALL_INDEX, which means that region
        // is destroyed.
        let mut wb = WriteBatch::new();
        wb.truncate_raft_log(2, 10, 1, TRUNCATE_ALL_INDEX);
        engine.write(wb).unwrap();

        // Because the newest peer is truncated, region 10 should be removed entirely.
        let map = engine.get_region_peer_map();
        assert!(!map.contains_key(&10), "region 10 should be removed");

        // ---------------------
        // Region 20: multiple peers, confirm the highest ID remains.
        // ---------------------
        // Add peer 5 (older) and peer 7 (newer) to the same region 20.
        for pid in [5u64, 7u64] {
            let mut wb = WriteBatch::new();
            wb.append_raft_log(
                pid,
                20,
                1,
                &new_raft_entry(EntryType::EntryNormal, pid, 1, b"data", 0),
            );
            engine.write(wb).unwrap();
        }
        // get_region_peer_map should pick peer 7 for region 20.
        let map = engine.get_region_peer_map();
        assert_eq!(map.get(&20), Some(&7), "peer 7 should override peer 5");
    }

    #[test]
    fn test_load_store_region_states() {
        init_logger();
        let tmp_dir = tempfile::tempdir().unwrap();
        let cfg = Config::new(128 * 1024_usize);
        let engine = RfEngine::open(tmp_dir.path(), &cfg, None, None).unwrap();

        let cluster_id = 1;
        let store_id = 2;
        let engine_id = 3;

        engine.set_engine_id(engine_id);
        assert_eq!(engine.get_engine_id(), engine_id);

        let mut ident = StoreIdent::default();
        ident.set_cluster_id(cluster_id);
        ident.set_store_id(store_id);
        let bin = ident.write_to_bytes().unwrap();
        let mut wb = WriteBatch::new();
        wb.set_state(0, 0, 1, STORE_IDENT_KEY, bin.as_slice());
        engine.write(wb).unwrap();

        let loaded_ident = load_store_ident(&engine).unwrap();
        assert_eq!(loaded_ident.cluster_id, cluster_id);
        assert_eq!(loaded_ident.store_id, store_id);

        let peer_id = 5;
        let region_id = 6;
        let region_epoch = 11;
        let mut wb = WriteBatch::new();
        let (key, val) = make_region_state(region_epoch, 12 /* keyspace_id */);
        wb.set_state(peer_id, region_id, 1, &key, &val);
        engine.write(wb).unwrap();

        let state = engine.load_region_state(peer_id, region_epoch);
        assert!(state.is_some(), "the region state should not be None");
    }

    #[test]
    fn test_calc_offload_epoch() {
        let current_epoch = Arc::new(AtomicU32::new(10)); // unchanged
        let compacted_epoch = Arc::new(AtomicU32::new(0));
        for (count, compacted, expected) in [
            (0, 10, 0),
            (0, 9, 0),
            (5, 4, 4),
            (5, 5, 5),
            (5, 6, 5),
            (10, 5, 0),
            (10, 10, 0),
            (20, 10, 0),
        ] {
            compacted_epoch.store(compacted, Ordering::SeqCst);
            let offload_epoch = calc_offload_epoch(count, &current_epoch, &compacted_epoch);
            assert_eq!(
                offload_epoch, expected,
                "current_epoch=10, compacted={}, in_mem_rlog_epoch_count={}, expected offload_epoch={}, got {}",
                compacted, count, expected, offload_epoch,
            );
        }
    }

    #[test]
    fn test_rfengine_fetch_offloaded() {
        init_logger();
        const STATE_PREFIX: u8 = b'p';

        let dir = tempfile::tempdir().unwrap();
        let mut cfg = Config::new(8 * 1024); // WAL size
        cfg.rlog_soft_memory_limit = ReadableSize(8 * 1024); // Offload immediately after compaction.

        let engine = RfEngine::open(dir.path(), &cfg, None, None).unwrap();
        // Record the entries and states inserted into rfengine. Will be used
        // for comparison later.
        let mut data_map = HashMap::new();
        // Track the next index that should be inserted for each peer.
        let mut next_index = [1u64; 2];
        // The keyspace id of PeerBatch exists only in memory and is not persisted to
        // disk, so it is lost after restarting rfengine. This test checks whether peer
        // meta remains consistent before and after restarting rfengine.  To avoid the
        // keyspace id affecting the result, the keyspace id is set to zero.
        let keyspace_id = 0;
        for peer_id in [1, 2, 1, 2] {
            let mut wb = WriteBatch::new();
            for i in 1..=1000 {
                let entry_index = next_index[peer_id - 1];
                next_index[peer_id - 1] += 1;

                let peer_id = peer_id as u64;
                let region_id = peer_id + 1000;
                let entry =
                    new_raft_entry(EntryType::EntryNormal, peer_id, entry_index, b"data", 0);
                let (state_key, state_val) = (&[STATE_PREFIX, i as u8], &[i as u8]);
                wb.append_raft_log(peer_id, region_id, keyspace_id, &entry);
                wb.set_state(peer_id, region_id, keyspace_id, state_key, state_val);

                let (entries, states) = data_map
                    .entry(peer_id)
                    .or_insert_with(|| (vec![], BTreeMap::new()));
                entries.push(entry);
                states.insert(state_key.to_vec(), state_val.to_vec());
            }
            engine.write(wb).unwrap();
            // Wait for the compaction to finish.
            eventually(Duration::from_millis(100), Duration::from_secs(1), || {
                engine.current_epoch_id.load(Ordering::SeqCst)
                    == engine.compacted_epoch.load(Ordering::SeqCst) + 1
            });
        }
        assert_eq!(engine.get_last_index(1), Some(2000));
        assert_eq!(engine.get_last_index(2), Some(2000));

        // Overwrite one entry at index 500 on peer 2.
        let mut wb = WriteBatch::new();
        wb.append_raft_log(
            2,
            1002,
            keyspace_id,
            &new_raft_entry(EntryType::EntryNormal, 10_u64, 500, b"data", 0),
        );
        engine.write(wb).unwrap();
        assert_eq!(engine.get_last_index(2), Some(500));

        let peer_meta_1 = extract_peer_meta(&engine, 1, |raft_logs| {
            // Entries [1,1000] are offloaded, [1001,2000] are eligible for
            // offloading but still in memory.
            assert_eq!(raft_logs.first_index(), 1001);
        });
        let peer_meta_2 = extract_peer_meta(&engine, 2, |_| {});
        check_engine_data(&engine, &data_map);

        // Restart rfengine.
        engine.stop_worker(true);
        let engine = RfEngine::open(dir.path(), &cfg, None, None).unwrap();

        // Check rfengine again after the restart.
        check_engine_data(&engine, &data_map);
        assert_eq!(
            peer_meta_1,
            extract_peer_meta(&engine, 1, |raft_logs| {
                // all entries are offloaded after the restart.
                assert_eq!(raft_logs.first_index(), 0);
            })
        );
        assert_eq!(peer_meta_2, extract_peer_meta(&engine, 2, |_| {}));
    }

    // Extracts and returns the peer meta after applying a custom check on its
    // raft logs.
    fn extract_peer_meta<F>(engine: &RfEngine, peer_id: u64, raft_logs_check: F) -> PeerMeta
    where
        F: FnOnce(&RaftLogs),
    {
        let peers = engine.peers.peers.pin();
        let peer_ref = peers.get(&peer_id).unwrap();
        let peer = peer_ref.read().unwrap();
        raft_logs_check(&peer.raft_logs);
        peer.meta.clone()
    }

    // Tests raft log fetching across in-memory and offloaded entries.
    fn check_engine_data(
        engine: &RfEngine,
        data_map: &HashMap<u64, (Vec<Entry>, BTreeMap<Vec<u8>, Vec<u8>>)>,
    ) {
        assert_eq!(engine.get_term(1, 100), Some(1));
        assert_eq!(engine.get_term(1, 1000), Some(1)); // offloaded entry
        assert_eq!(engine.get_term(1, 1001), Some(1));
        assert!(engine.get_raft_entry(1, 1000).is_some());
        assert!(engine.get_raft_entry(1, 2000).is_some());

        // Test `fetch_entries_to`.
        let mut buf = vec![];

        for (low, high) in [(100, 200), (900, 1100), (1200, 1300), (1, 2001)] {
            assert_eq!(
                engine
                    .fetch_raft_entries_to(1, low, high, None, &mut buf, None)
                    .unwrap(),
                (high - low) as usize
            );
            assert_eq!(
                data_map.get(&1).unwrap().0[(low - 1) as usize..(high - 1) as usize],
                buf
            );
            buf.clear();
        }

        // Test fetch unavailable entries.
        assert!(matches!(
            engine.fetch_raft_entries_to(1, 2001, 2002, None, &mut buf, None),
            Err(Error::EntriesUnavailable),
        ));

        // Test fetch empty logs.
        buf.clear();
        assert_eq!(
            engine
                .fetch_raft_entries_to(1, 1, 1, None, &mut buf, None)
                .unwrap(),
            0
        );
        assert!(buf.is_empty());

        // Test `fetch_entries_to` should push logs to the buf.
        let peer1_entries = &data_map.get(&1).unwrap().0;
        for i in 1..=10 {
            assert_eq!(
                engine
                    .fetch_raft_entries_to(1, i, i + 1, None, &mut buf, None)
                    .unwrap(),
                1
            );
            assert_eq!(buf, peer1_entries[..i as usize]);
        }

        // Test `fetch_entries_to` should respect size limit.
        let mut max_size = 0;
        for (i, entry) in peer1_entries.iter().enumerate() {
            if i == 10 {
                break;
            }
            buf.clear();
            max_size += entry.compute_size();
            assert_eq!(
                engine
                    .fetch_raft_entries_to(1, 1, 11, Some(max_size as usize), &mut buf, None)
                    .unwrap(),
                i + 1
            );
            assert_eq!(buf, peer1_entries[..=i]);
        }

        // Peer 2: entries [1, 499] are offloaded, entry at index 500 is
        // in-memory.
        assert_eq!(engine.get_term(2, 100), Some(2));
        assert_eq!(engine.get_term(2, 499), Some(2));
        assert_eq!(engine.get_term(2, 500), Some(10));
        assert_eq!(engine.get_term(2, 501), None);
        assert!(matches!(
            engine.fetch_raft_entries_to(2, 501, 502, None, &mut buf, None),
            Err(Error::EntriesUnavailable),
        ));

        buf.clear();
        assert_eq!(
            engine
                .fetch_raft_entries_to(2, 1, 501, None, &mut buf, None)
                .unwrap(),
            500
        );

        let last_idx = buf.len() - 1;
        let second_last_idx = last_idx - 1;
        assert_eq!(buf[last_idx].term, 10);
        assert_eq!(buf[second_last_idx].term, 2);
    }

    #[test]
    fn test_peer_index_to_truncate_to_size() {
        let mut peer = PeerData::default();
        peer.peer_id = 1;

        // Insert 4 blocks of logs into the peer.
        let data = b"data";
        let total_cnt = RAFT_LOG_BLOCK_CAP * 4;
        let total_size = total_cnt * data.len();
        for i in 1..=total_cnt {
            let log = RaftLogOp::new(&new_raft_entry(
                EntryType::EntryNormal,
                1,
                i as u64,
                data,
                0,
            ));
            peer.raft_logs.append(peer.peer_id, log);
        }

        fn check_truncate_idx(peer: &PeerData, total_cnt: usize, total_size: usize) {
            let q = total_cnt as u64 / 4;
            assert_eq!(peer.index_to_truncate_to_size(total_size + 1), 0);
            assert_eq!(peer.index_to_truncate_to_size(total_size / 4 * 3 + 1), q);
            assert_eq!(peer.index_to_truncate_to_size(total_size / 2 + 1), q * 2);
            assert_eq!(peer.index_to_truncate_to_size(total_size / 4 + 1), q * 3);
        }

        for i in 0..=3 {
            // Simulate progressive offloading by truncating more in-memory logs.
            peer.raft_logs
                .truncate(i as u64 * RAFT_LOG_BLOCK_CAP as u64);
            check_truncate_idx(&peer, total_cnt, total_size);
        }
    }
}