// Copyright 2026 PingCAP, Inc.

use std::{
    fmt, io,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};

use async_trait::async_trait;
use collections::{HashMap, HashMapEntry, HashMapExt, HashSet};
use http::StatusCode;
use kvengine::dfs::Dfs;
use merged_engine::StoreProgress;
use native_br::{
    common::{
        CollectWalChunksContext, TempLocalObject, TempLocalObjectWriter, assemble_wal_chunks,
        collect_wal_chunks_with_retry,
    },
    wal::AssembledWalData,
};
use parking_lot::Mutex;
use pd_client::{PdClient, util::get_all_stores_except_tiflash_async};
use rfengine::MIN_EPOCH_ROTATE_LEN;
use tikv_util::{
    box_err, box_try, debug, info,
    time::{Instant, Limiter},
    warn,
};
use tokio::sync::{
    mpsc::{self, error::TryRecvError},
    watch,
};
use txn_types::TimeStamp;

use crate::{
    Error, ReplicationWorkerConfig, Result, WalConfig,
    error::{error_is_force_advance_eligible, error_label},
    metrics::*,
    util::{WriterWritePolicy, send_request_to_store_with_writer},
    wal::{
        FetchWalError, StoreLag, StoreWalProgresses, WalTarget, WalTargetFetcher,
        WalTargetFetcherHandle, WalTargetReceiver, wal_target_channel,
    },
    wal_spool::{create_temp_wal_spool_object, spool_dir_for_store},
};

const FETCH_WAL_TIMEOUT: Duration = Duration::from_secs(30);

#[derive(Clone, Debug)]
pub(crate) enum WalFetcherRequest {
    Fetch {
        store_progresses: Vec<StoreProgress>,
        appended_target_ts: Option<TimeStamp>,
    },
    Stop,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum WalFetchSource {
    Store,
    Dfs,
}

impl WalFetchSource {
    fn as_str(&self) -> &'static str {
        match self {
            WalFetchSource::Store => "store",
            WalFetchSource::Dfs => "dfs",
        }
    }
}

pub(crate) enum FetchedWalItem {
    Wal {
        store_id: u64,
        epoch: u32,
        start_off: u64,
        end_off: u64,
        rotate: bool,
        size: u64,
        source: WalFetchSource,
        data: FetchedWalData,
    },
    ForceAdvance {
        store_id: u64,
        epoch: u32,
        offset: u64,
    },
}

pub(crate) enum FetchedWalData {
    Spool(TempLocalObject),
    SharedSpool(SharedWal),
}

pub(crate) type SharedWal = Arc<Mutex<AssembledWalData>>;

pub(crate) struct FetchedWalsBatch {
    pub(crate) items: Vec<FetchedWalItem>,
    pub(crate) advance_to: Option<TimeStamp>,
    pub(crate) target_stores: HashSet<u64 /* store_id */>,
}

pub(crate) struct WalFetcherHandle {
    // ReplicationWorker -> WalFetcher
    request_tx: mpsc::UnboundedSender<WalFetcherRequest>,
    // ReplicationWorker <- WalFetcher
    wals_batch_rx: mpsc::UnboundedReceiver<FetchedWalsBatch>,
    join_handle: tokio::task::JoinHandle<()>,

    // ReplicationWorker -> WalFetcher/WalTargetFetcher
    stop_tx: watch::Sender<bool>,
    target_fetcher: WalTargetFetcherHandle,
}

impl WalFetcherHandle {
    // ReplicationWorker -> WalFetcher
    pub(crate) fn request_tx(&self) -> &mpsc::UnboundedSender<WalFetcherRequest> {
        &self.request_tx
    }

    // ReplicationWorker <- WalFetcher
    pub(crate) fn try_recv_batch(&mut self) -> Result<Option<FetchedWalsBatch>> {
        match self.wals_batch_rx.try_recv() {
            Ok(batch) => Ok(Some(batch)),
            Err(TryRecvError::Empty) => Ok(None),
            Err(TryRecvError::Disconnected) => Err(Error::Stopped("WalFetcher stopped".into())),
        }
    }

    pub(crate) async fn shutdown(self) {
        let _ = self.request_tx.send(WalFetcherRequest::Stop);
        let _ = self.stop_tx.send(true);

        let join_res = self.join_handle.await;

        self.target_fetcher.wait().await;

        if let Err(err) = join_res {
            if err.is_panic() {
                panic!("WalFetcher panic: {err:?}");
            }
        }
    }
}

pub(crate) fn spawn_wal_fetcher(
    pd: Arc<dyn PdClient>,
    fs: Arc<dyn Dfs>,
    http_client: Arc<security::HttpClient>,
    worker_dir: PathBuf,
    init_target_ts: TimeStamp,
    config: &ReplicationWorkerConfig,
    handle: tokio::runtime::Handle,
) -> WalFetcherHandle {
    let (request_tx, request_rx) = mpsc::unbounded_channel();
    let (wals_batch_tx, wals_batch_rx) = mpsc::unbounded_channel();

    let (target_tx, target_rx) = wal_target_channel();
    let (stop_tx, stop_rx) = watch::channel(false);
    let target_fetcher = WalTargetFetcher::spawn(
        pd.clone(),
        init_target_ts,
        config,
        handle.clone(),
        fs.clone(),
        target_tx,
        stop_rx.clone(),
    );

    let io = Arc::new(RealWalFetcherIo::new(
        pd.clone(),
        fs,
        http_client,
        worker_dir,
        &config.wal,
    ));

    let fetcher = WalFetcher::new(
        io,
        pd,
        target_rx,
        config.wal.clone(),
        wals_batch_tx,
        request_rx,
        stop_rx,
    );
    let join_handle = handle.spawn(fetcher.run());

    WalFetcherHandle {
        request_tx,
        wals_batch_rx,
        join_handle,
        stop_tx,
        target_fetcher,
    }
}

struct ActiveTarget {
    ts: TimeStamp,
    store_targets: StoreWalProgresses,
    rr_cursor: usize,
    store_status_addr: HashMap<u64, String>,
    errors: HashMap<u64, FetchWalError>,
}

impl fmt::Debug for ActiveTarget {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ActiveTarget")
            .field("ts", &self.ts)
            .field("store_targets", &self.store_targets)
            .field("rr_cursor", &self.rr_cursor)
            .field("errors", &self.errors)
            .finish()
    }
}

impl ActiveTarget {
    fn store_count(&self) -> usize {
        self.store_targets.len()
    }
}

#[async_trait]
trait WalFetcherIo: Send + Sync {
    async fn fetch_from_store(
        &self,
        store_id: u64,
        status_address: &str,
        epoch_id: u32,
        start_off: u64,
        end_off: u64,
    ) -> Result<(StatusCode, TempLocalObject, u64 /* read_len */)>;

    async fn fetch_from_dfs(
        &self,
        store_id: u64,
        epoch_id: u32,
        target: StoreProgress,
    ) -> Result<(
        SharedWal,
        u64,  // end_off
        bool, // rotate
    )>;
}

struct RealWalFetcherIo {
    inner: Arc<RealWalFetcherIoInner>,
}

#[derive(Default)]
struct WalCache {
    inner: HashMap<u64 /* store_id */, Option<(u32 /* epoch_id */, SharedWal)>>,
}

impl WalCache {
    fn ensure_store(&mut self, store_id: u64) {
        self.inner.entry(store_id).or_insert(None);
    }

    fn get(&self, store_id: u64, epoch: u32) -> Option<SharedWal> {
        let entry = self.inner.get(&store_id)?.as_ref()?;
        if entry.0 != epoch {
            return None;
        }
        Some(entry.1.clone())
    }

    fn put(&mut self, store_id: u64, epoch: u32, wal_data: SharedWal) {
        self.inner.insert(store_id, Some((epoch, wal_data)));
    }

    fn remove_cache(&mut self, store_id: u64) {
        if let Some(entry) = self.inner.get_mut(&store_id) {
            entry.take();
        }
    }

    fn contains_store(&self, store_id: u64) -> bool {
        self.inner.contains_key(&store_id)
    }
}

impl RealWalFetcherIo {
    fn new(
        pd: Arc<dyn PdClient>,
        fs: Arc<dyn Dfs>,
        http_client: Arc<security::HttpClient>,
        worker_dir: PathBuf,
        wal_config: &WalConfig,
    ) -> Self {
        info!("create RealWalFetcherIo"; "config" => ?wal_config);
        let store_fetch_speed_limit = wal_fetch_speed_limit(wal_config.store_fetch_max_throughput);
        Self {
            inner: Arc::new(RealWalFetcherIoInner {
                pd,
                fs,
                http_client,
                worker_dir,
                store_fetch_speed_limit,
                store_limiters: Mutex::new(HashMap::default()),
                wal_cache: Mutex::new(WalCache::default()),
            }),
        }
    }

    fn store_fetch_limiter(&self, store_id: u64) -> Option<Limiter> {
        let store_fetch_speed_limit = self.inner.store_fetch_speed_limit?;
        let mut store_limiters = self.inner.store_limiters.lock();
        Some(
            store_limiters
                .entry(store_id)
                .or_insert_with(|| build_wal_fetch_limiter(store_fetch_speed_limit))
                .clone(),
        )
    }

    fn store_cache_dir(&self, store_id: u64) -> PathBuf {
        spool_dir_for_store(&self.inner.worker_dir, store_id)
    }

    fn ensure_store_cache_dir(&self, store_id: u64) -> Result<PathBuf> {
        let cache_dir = self.store_cache_dir(store_id);
        if self.inner.wal_cache.lock().contains_store(store_id) {
            return Ok(cache_dir);
        }
        box_try!(std::fs::create_dir_all(&cache_dir));
        self.inner.wal_cache.lock().ensure_store(store_id);
        Ok(cache_dir)
    }

    fn take_or_fetch_epoch_wal_from_dfs(&self, store_id: u64, epoch_id: u32) -> Result<SharedWal> {
        if let Some(wal) = self.inner.wal_cache.lock().get(store_id, epoch_id) {
            return Ok(wal);
        }

        let cache_dir = self.ensure_store_cache_dir(store_id)?;

        let collect_ctx = CollectWalChunksContext {
            pd_client: self.inner.pd.clone(),
            dfs: self.inner.fs.clone(),
            store_id,
            complete_wal_chunks: true,
            fetch_wal_timeout: FETCH_WAL_TIMEOUT,
            cache_dir: Some(cache_dir),
            wal_chunks_cache: None,
            retry_for_wal_chunk_integrity_error: false,
        };
        let tag = format!("{}:{}", store_id, epoch_id);
        let (chunks, online_chunk, has_last_chunk) =
            collect_wal_chunks_with_retry(&tag, &collect_ctx, epoch_id, epoch_id, u64::MAX)?;
        debug_assert!(online_chunk.is_none());
        debug_assert!(has_last_chunk);

        let mut epoch_wal = assemble_wal_chunks(chunks, false)?;
        epoch_wal.freeze();
        let wal_data = Arc::new(Mutex::new(epoch_wal));
        self.inner
            .wal_cache
            .lock()
            .put(store_id, epoch_id, wal_data.clone());
        Ok(wal_data)
    }
}

struct RealWalFetcherIoInner {
    pd: Arc<dyn PdClient>,
    fs: Arc<dyn Dfs>,
    http_client: Arc<security::HttpClient>,
    worker_dir: PathBuf,
    store_fetch_speed_limit: Option<f64>,
    store_limiters: Mutex<HashMap<u64, Limiter>>,
    wal_cache: Mutex<WalCache>,
}

#[async_trait]
impl WalFetcherIo for RealWalFetcherIo {
    async fn fetch_from_store(
        &self,
        store_id: u64,
        status_address: &str,
        epoch_id: u32,
        start_off: u64,
        end_off: u64,
    ) -> Result<(StatusCode, TempLocalObject, u64)> {
        let cache_dir = self.ensure_store_cache_dir(store_id)?;
        let store_limiter = self.store_fetch_limiter(store_id);

        let security_mgr = self.inner.pd.get_security_mgr();
        let uri = security_mgr
            .build_uri(format!(
                "{}/rfengine/wal_chunk?epoch_id={}&start_off={}&end_off={}",
                status_address, epoch_id, start_off, end_off
            ))
            .unwrap();
        let req = http::Request::get(uri).body(hyper::Body::empty()).unwrap();

        let (status, writer, read_len) = send_request_to_store_with_writer(
            req,
            self.inner.http_client.as_ref(),
            FETCH_WAL_TIMEOUT,
            WriterWritePolicy::Blocking(store_limiter),
            || {
                let obj = create_temp_wal_spool_object(&cache_dir, store_id, epoch_id, start_off)?;
                Ok(TempLocalObjectWriter::new(obj))
            },
        )
        .await?;

        let mut obj = writer.into_inner().map_err(|e| -> Error {
            box_err!("extract temp wal spool object failed: {:?}", e.into_error())
        })?;
        obj.close();
        Ok((status, obj, read_len))
    }

    async fn fetch_from_dfs(
        &self,
        store_id: u64,
        epoch_id: u32,
        target: StoreProgress,
    ) -> Result<(SharedWal, u64, bool)> {
        let inner = self.inner.clone();
        tokio::task::spawn_blocking(move || -> Result<(SharedWal, u64, bool)> {
            let io = RealWalFetcherIo { inner };

            let wal_data = io.take_or_fetch_epoch_wal_from_dfs(store_id, epoch_id)?;
            let wal_len = { wal_data.lock().len() };
            let end_off = if epoch_id == target.epoch {
                target.offset
            } else {
                wal_len
            };
            let rotate = end_off == wal_len; // `wal_data` must be a complete epoch.

            // Keep the per-store epoch cache unless we are sure this epoch is complete.
            if rotate {
                io.inner.wal_cache.lock().remove_cache(store_id);
            }

            Ok((wal_data, end_off, rotate))
        })
        .await
        .map_err(|e| -> Error { box_err!("fetch wal from dfs task join failed: {:?}", e) })?
    }
}

fn wal_fetch_speed_limit(max_throughput: tikv_util::config::ReadableSize) -> Option<f64> {
    if max_throughput.0 == 0 {
        None
    } else {
        Some(max_throughput.0 as f64)
    }
}

fn build_wal_fetch_limiter(speed_limit: f64) -> Limiter {
    <Limiter>::builder(speed_limit)
        .refill(Duration::from_millis(10))
        .build()
}

struct WalFetcher {
    io: Arc<dyn WalFetcherIo>,
    pd: Arc<dyn PdClient>,
    target_rx: WalTargetReceiver,

    config: WalConfig,
    fetch_concurrency: usize,

    wals_batch_tx: mpsc::UnboundedSender<FetchedWalsBatch>,
    request_rx: mpsc::UnboundedReceiver<WalFetcherRequest>,
    stop_rx: watch::Receiver<bool>,

    active: Option<ActiveTarget>,
    store_progresses: HashMap<u64, StoreProgress>,
}

impl WalFetcher {
    fn new(
        io: Arc<dyn WalFetcherIo>,
        pd: Arc<dyn PdClient>,
        target_rx: WalTargetReceiver,
        config: WalConfig,
        wals_batch_tx: mpsc::UnboundedSender<FetchedWalsBatch>,
        request_rx: mpsc::UnboundedReceiver<WalFetcherRequest>,
        stop_rx: watch::Receiver<bool>,
    ) -> Self {
        let fetch_concurrency = config.calc_fetch_wal_concurrency();
        Self {
            io,
            pd,
            target_rx,
            config,
            fetch_concurrency,
            wals_batch_tx,
            request_rx,
            stop_rx,
            active: None,
            store_progresses: HashMap::default(),
        }
    }

    fn tolerate_store_err(&self) -> usize {
        self.config.tolerate_store_err as usize
    }

    async fn run(mut self) {
        while let Some(req) = self.request_rx.recv().await {
            match req {
                WalFetcherRequest::Fetch {
                    store_progresses,
                    appended_target_ts,
                } => match self
                    .handle_fetch_request(store_progresses, appended_target_ts)
                    .await
                {
                    Ok(()) => {}
                    Err(Error::Stopped(msg)) => {
                        info!("WalFetcher exit"; "cause" => msg);
                        return;
                    }
                    Err(err) => {
                        warn!("WalFetcher: handle request failed: {:?}", err);
                    }
                },
                WalFetcherRequest::Stop => {
                    info!("WalFetcher exit");
                    return;
                }
            }
        }
    }

    async fn handle_fetch_request(
        &mut self,
        store_progresses: Vec<StoreProgress>,
        appended_target_ts: Option<TimeStamp>,
    ) -> Result<()> {
        let start_time = Instant::now_coarse();
        REP_FETCH_WAL_COUNTER_VEC.fetch_request.inc();

        self.merge_store_progresses(store_progresses);

        if let Some(ts) = appended_target_ts {
            if self.active.as_ref().is_some_and(|t| t.ts == ts) {
                self.active = None;
            }
        }
        self.ensure_active_target().await?;
        self.emit_next_batch().await?;

        REP_FETCH_WAL_DURATION.observe(start_time.saturating_elapsed().as_secs_f64());
        Ok(())
    }

    fn merge_store_progresses(&mut self, progresses: Vec<StoreProgress>) {
        for p in progresses {
            match self.store_progresses.entry(p.store_id) {
                HashMapEntry::Occupied(mut e) => {
                    let current = *e.get();
                    if p > current {
                        e.insert(p);
                    }
                }
                HashMapEntry::Vacant(e) => {
                    e.insert(p);
                }
            }
        }
    }

    async fn ensure_active_target(&mut self) -> Result<()> {
        const REPORT_INTERVAL: Duration = Duration::from_secs(10);

        if self.active.is_some() {
            return Ok(());
        }
        loop {
            tokio::select! {
                item = self.target_rx.recv() => {
                    return match item {
                        Some(WalTarget { ts, store_targets }) => {
                            REP_FETCH_WAL_COUNTER_VEC.recv_target.inc();
                            let store_status_addr = self.refresh_store_status_address_map().await?;
                            self.active = Some(ActiveTarget {
                                ts,
                                store_targets,
                                rr_cursor: 0,
                                store_status_addr,
                                errors: HashMap::default(),
                            });
                            Ok(())
                        }
                        None => Err(Error::Stopped("WalTargetFetcher stopped".into())),
                    };
                }
                _ = tokio::time::sleep(REPORT_INTERVAL) => {
                    REP_FETCH_WAL_COUNTER_VEC.no_target.inc();
                    warn!("WalFetcher: no target");
                }
            };
        }
    }

    async fn refresh_store_status_address_map(&self) -> Result<HashMap<u64, String>> {
        const MAX_BACKOFF: Duration = Duration::from_secs(5);

        let mut stop_rx = self.stop_rx.clone();
        let mut retry_backoff = Duration::from_millis(200);
        loop {
            if *stop_rx.borrow() {
                return Err(Error::Stopped("WalFetcher stopped".into()));
            }

            let stores = tokio::select! {
                res = get_all_stores_except_tiflash_async(self.pd.as_ref(), true) => res,
                res = stop_rx.changed() => {
                    if res.is_err() || *stop_rx.borrow() {
                        return Err(Error::Stopped("WalFetcher stopped".into()));
                    }
                    continue;
                }
            };

            match stores {
                Ok(stores) => {
                    let mut map = HashMap::with_capacity(stores.len());
                    for store in stores {
                        map.insert(store.id, store.status_address);
                    }
                    REP_FETCH_WAL_COUNTER_VEC.refresh_store.inc();
                    return Ok(map);
                }
                Err(err) => {
                    REP_FETCH_WAL_COUNTER_VEC.refresh_store_err.inc();
                    warn!(
                        "WalFetcher: refresh store status address map failed: {:?}",
                        err
                    );
                }
            }

            tokio::select! {
                _ = tokio::time::sleep(retry_backoff) => {}
                res = stop_rx.changed() => {
                    if res.is_err() || *stop_rx.borrow() {
                        return Err(Error::Stopped("WalFetcher stopped".into()));
                    }
                }
            }
            retry_backoff = retry_backoff.saturating_mul(2).min(MAX_BACKOFF);
        }
    }

    fn effective_concurrency(&self, store_count: usize) -> usize {
        self.fetch_concurrency.min(store_count).max(1)
    }

    // Return error only on stopped.
    async fn emit_next_batch(&mut self) -> Result<()> {
        let Some(mut active) = self.active.take() else {
            // `ensure_active_target` must have set the active target.
            debug_assert!(false);
            return Ok(());
        };

        let (mut items, decision) = self.build_batch(&mut active).await?;
        let target_stores = active.store_targets.keys().copied().collect();
        let advance_to = match decision {
            TargetDecision::Advance(target_ts) => {
                // Encode force-advances for errors with explicit targets.
                for err in active.errors.values() {
                    if let Some(target) = err.target {
                        if err.is_force_advance_eligible() {
                            items.push(FetchedWalItem::ForceAdvance {
                                store_id: target.store_id,
                                epoch: target.epoch,
                                offset: target.offset,
                            });
                        }
                    }
                }
                REP_FETCH_WAL_ADVANCE_TARGET_TS.set(target_ts.physical() as i64);
                debug!("WalFetcher: target finished"; "target" => ?active);
                Some(target_ts)
            }
            TargetDecision::Drop => {
                warn!("WalFetcher: drop target"; "target" => ?active);
                None
            }
            TargetDecision::Keep => {
                debug!("WalFetcher: keep target"; "target" => ?active);
                self.active = Some(active);
                None
            }
        };

        if let Err(err) = self.wals_batch_tx.send(FetchedWalsBatch {
            items,
            advance_to,
            target_stores,
        }) {
            return Err(Error::Stopped(format!("send WALs batch failed: {err}")));
        }

        REP_FETCH_WAL_COUNTER_VEC.emit_batch.inc();
        Ok(())
    }

    // Return error only on stopped.
    async fn build_batch(
        &mut self,
        active: &mut ActiveTarget,
    ) -> Result<(Vec<FetchedWalItem>, TargetDecision)> {
        if active.store_targets.is_empty() {
            // All stores are filter out.
            return Ok((Vec::new(), TargetDecision::Advance(active.ts)));
        }

        let effective_concurrency = self.effective_concurrency(active.store_count());
        let targets_and_lags =
            get_decreasing_stores_lag(&active.store_targets, &self.store_progresses);
        let selected = select_round_robin_stores(
            &targets_and_lags,
            effective_concurrency,
            &self.store_progresses,
            &active.store_targets,
            &mut active.rr_cursor,
        );

        let mut predicted = HashMap::default();
        predicted.extend(self.store_progresses.iter().map(|(&k, &v)| (k, v)));

        // Pre-fill target-none errors.
        for (&store_id, target) in &active.store_targets {
            if target.is_none() {
                active.errors.insert(
                    store_id,
                    FetchWalError {
                        err: Error::StoreUnhealthy { store_id },
                        target: None,
                    },
                );
            }
        }

        // Target-level drop: store-unhealthy beyond tolerance cannot be recovered for
        // this target.
        let unhealthy_cnt = active
            .store_targets
            .values()
            .filter(|t| t.is_none())
            .count();
        if unhealthy_cnt > self.tolerate_store_err() {
            // WalTargetFetcher should not return such target.
            debug_assert!(false);
            return Ok((Vec::new(), TargetDecision::Drop));
        }

        let start_time = Instant::now_coarse();
        let mut join_set = tokio::task::JoinSet::new();
        for store_id in selected {
            let Some(target) = active.store_targets.get(&store_id).copied().flatten() else {
                continue;
            };

            let current = predicted.get(&store_id).copied().unwrap_or(StoreProgress {
                store_id,
                epoch: 1,
                offset: 0,
            });
            if current >= target {
                active.errors.remove(&store_id);
                continue;
            }

            let store_status_addr = active.store_status_addr.get(&store_id).cloned();
            let io = self.io.clone();
            join_set.spawn(async move {
                fetch_one_store(io, store_id, store_status_addr, current, target).await
            });
        }

        let mut items = Vec::new();
        while let Some(task) = join_set.join_next().await {
            match task {
                Ok(Ok(item)) => {
                    // Update predicted progress for finish decision.
                    if let FetchedWalItem::Wal {
                        store_id,
                        epoch,
                        end_off,
                        rotate,
                        ..
                    } = &item
                    {
                        let next = if *rotate {
                            StoreProgress {
                                store_id: *store_id,
                                epoch: epoch + 1,
                                offset: 0,
                            }
                        } else {
                            StoreProgress {
                                store_id: *store_id,
                                epoch: *epoch,
                                offset: *end_off,
                            }
                        };
                        predicted.insert(*store_id, next);
                        active.errors.remove(store_id);
                    }
                    items.push(item);
                }
                Ok(Err((store_id, target, err))) => {
                    warn!("WalFetcher: fetch failed: {:?}", err; "store" => store_id);
                    let eligible = error_is_force_advance_eligible(&err);
                    active.errors.insert(
                        store_id,
                        FetchWalError {
                            err,
                            target: eligible.then_some(target),
                        },
                    );
                }
                Err(join_err) if join_err.is_panic() => {
                    panic!("WalFetcher: fetch task panicked: {:?}", join_err);
                }
                Err(_join_err) => return Err(Error::Stopped("task cancelled".into())),
            }
        }
        REP_FETCH_STORE_WAL_DURATION
            .with_label_values(&["all"])
            .observe(start_time.saturating_elapsed().as_secs_f64());

        // Decide whether the target can be advanced after applying this batch.
        let decision = decide_advance_to(
            active.ts,
            &active.store_targets,
            &predicted,
            &mut active.errors,
            self.tolerate_store_err(),
        );
        Ok((items, decision))
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum TargetDecision {
    Advance(TimeStamp),
    Keep,
    Drop,
}

async fn fetch_one_store(
    io: Arc<dyn WalFetcherIo>,
    store_id: u64,
    store_status_addr: Option<String>,
    current: StoreProgress,
    target: StoreProgress,
) -> std::result::Result<FetchedWalItem, (u64 /* store_id */, StoreProgress, Error)> {
    REP_FETCH_WAL_INFLIGHT.inc();
    tikv_util::defer!({
        REP_FETCH_WAL_INFLIGHT.dec();
    });
    let start_time = tikv_util::time::Instant::now_coarse();

    let store_id_str = store_id.to_string();
    let record_store_metrics = |result: &'static str| {
        REP_FETCH_STORE_WAL_COUNTER
            .with_label_values(&[&store_id_str, result])
            .inc();
        REP_FETCH_STORE_WAL_DURATION
            .with_label_values(&[&store_id_str])
            .observe(start_time.saturating_elapsed().as_secs_f64());
    };

    let (epoch_id, start_off) = (current.epoch, current.offset);
    let end_off = if epoch_id < target.epoch {
        0
    } else {
        target.offset
    };

    // Use DFS when store is missing or WAL is near overwritten.
    if store_status_addr.is_none()
        || epoch_id <= near_overwritten_epoch(target.epoch, MIN_EPOCH_ROTATE_LEN)
    {
        return match io.fetch_from_dfs(store_id, epoch_id, target).await {
            Ok((wal_data, effective_end_off, rotate)) => {
                record_store_metrics("ok");
                let size = effective_end_off.saturating_sub(start_off);
                REP_FETCH_STORE_WAL_BYTES
                    .with_label_values(&[&store_id_str, WalFetchSource::Dfs.as_str()])
                    .inc_by(size);

                Ok(FetchedWalItem::Wal {
                    store_id,
                    epoch: epoch_id,
                    start_off,
                    end_off: effective_end_off,
                    rotate,
                    size,
                    source: WalFetchSource::Dfs,
                    data: FetchedWalData::SharedSpool(wal_data),
                })
            }
            Err(err) => {
                record_store_metrics(error_label(&err));
                Err((store_id, target, err))
            }
        };
    }

    let status_addr = store_status_addr.unwrap();
    let (status, spool, read_len) = match io
        .fetch_from_store(store_id, &status_addr, epoch_id, start_off, end_off)
        .await
    {
        Ok(res) => res,
        Err(err) => {
            record_store_metrics(error_label(&err));
            return Err((store_id, target, err));
        }
    };

    if status == StatusCode::GONE {
        drop(spool);
        let (wal_data, effective_end_off, rotate) =
            match io.fetch_from_dfs(store_id, epoch_id, target).await {
                Ok(res) => res,
                Err(err) => {
                    record_store_metrics(error_label(&err));
                    return Err((store_id, target, err));
                }
            };
        record_store_metrics("ok");
        let size = effective_end_off.saturating_sub(start_off);
        REP_FETCH_STORE_WAL_BYTES
            .with_label_values(&[&store_id_str, WalFetchSource::Dfs.as_str()])
            .inc_by(size);
        return Ok(FetchedWalItem::Wal {
            store_id,
            epoch: epoch_id,
            start_off,
            end_off: effective_end_off,
            rotate,
            size,
            source: WalFetchSource::Dfs,
            data: FetchedWalData::SharedSpool(wal_data),
        });
    }

    if !status.is_success() {
        // Keep allocation bounded for unexpected large error bodies.
        let err_prefix = read_file_prefix(spool.path.as_ref(), 8 * 1024)
            .unwrap_or_else(|_| "<failed to read error body>".to_string());
        drop(spool);
        record_store_metrics("error");
        return Err((
            store_id,
            target,
            box_err!("store {} dump wal_chunk failed: {}", store_id, err_prefix),
        ));
    }

    let effective_end_off = start_off + read_len;
    // `/rfengine/wal_chunk` status semantics (see
    // `cloud_server::status_server::rfengine_wal_chunk`
    // and `rfengine::service_worker::handle_dump`):
    //
    // - `StatusCode::PARTIAL_CONTENT` (206) means `partial_content=true`:
    //   - when `end_off==0`, it means there are newer epochs after `epoch_id` (i.e.
    //     this is NOT the latest epoch on the store);
    //   - when `end_off>0`, it means this epoch has more bytes after the returned
    //     `effective_end_off` (including the case where `epoch_id` is the current
    //     epoch and can still grow).
    // - `StatusCode::OK` (200) means `partial_content=false`, which only happens
    //   when the returned range reaches the end of a *closed* epoch file.
    //
    // `rotate=true` indicates that the returned data completes the current epoch,
    // so the apply side should call `MergedEngine::rotate_wal(store_id,
    // epoch_id, effective_end_off)` after a successful `update_wal`.
    let rotate = if epoch_id < target.epoch {
        debug_assert_eq!(status, StatusCode::PARTIAL_CONTENT);
        true
    } else {
        status == StatusCode::OK
    };
    record_store_metrics("ok");
    REP_FETCH_STORE_WAL_BYTES
        .with_label_values(&[&store_id_str, WalFetchSource::Store.as_str()])
        .inc_by(read_len);
    Ok(FetchedWalItem::Wal {
        store_id,
        epoch: epoch_id,
        start_off,
        end_off: effective_end_off,
        rotate,
        size: read_len,
        source: WalFetchSource::Store,
        data: FetchedWalData::Spool(spool),
    })
}

fn read_file_prefix(path: &Path, max_len: usize) -> io::Result<String> {
    use std::io::Read;
    let mut file = std::fs::File::open(path)?;
    let mut buf = vec![0u8; max_len];
    let n = file.read(&mut buf)?;
    buf.truncate(n);
    Ok(String::from_utf8_lossy(&buf).to_string())
}

// Ref: ObjectStorageWorker::near_overwritten_epoch
fn near_overwritten_epoch(current_epoch: u32, epoch_rotate_len: usize) -> u32 {
    current_epoch.saturating_sub(epoch_rotate_len as u32 - 2)
}

fn get_decreasing_stores_lag(
    targets: &StoreWalProgresses,
    store_progresses: &HashMap<u64, StoreProgress>,
) -> Vec<(u64, StoreLag)> {
    let mut lags = Vec::with_capacity(targets.len());
    for (&store_id, target) in targets {
        let target_and_lag = if let Some(target) = target {
            let current = store_progresses
                .get(&store_id)
                .copied()
                .unwrap_or(StoreProgress {
                    store_id,
                    epoch: 1,
                    offset: 0,
                });
            let epoch_lag = target.epoch.saturating_sub(current.epoch);
            let offset_lag = target.offset as i64 - current.offset as i64;
            StoreLag {
                epoch_lag,
                offset_lag,
            }
        } else {
            StoreLag {
                epoch_lag: 0,
                offset_lag: 0,
            }
        };
        lags.push((store_id, target_and_lag));
    }
    lags.sort_by(|m, n| n.1.compare(&m.1)); // Decreasing sort.

    // Report lags.
    for (store_id, lag) in &lags {
        if targets.get(store_id).is_some_and(|t| t.is_some()) {
            let store_id_str = store_id.to_string();
            REP_FETCH_STORE_WAL_EPOCH_LAG
                .with_label_values(&[&store_id_str])
                .set(lag.epoch_lag as i64);
        }
    }

    lags
}

fn select_round_robin_stores(
    targets_and_lags: &[(u64, StoreLag)],
    max_count: usize,
    store_progresses: &HashMap<u64, StoreProgress>,
    targets: &StoreWalProgresses,
    rr_cursor: &mut usize,
) -> Vec<u64> {
    // Select up to `max_count` stores using a round-robin cursor over a list that
    // is already sorted by decreasing lag (`targets_and_lags`).
    //
    // Summary:
    // - Start from `rr_cursor` and walk the list once with wrap-around.
    // - Skip unhealthy stores (`target=None`) and stores already caught up
    //   (`current >= target`).
    // - Stop early once we have enough stores.
    // - Update `rr_cursor` so the next call continues from the next store, which
    //   avoids repeatedly picking the same "most lagging" store and improves
    //   fairness under a strict per-batch budget.
    if max_count == 0 || targets_and_lags.is_empty() {
        return Vec::new();
    }

    let mut selected = Vec::with_capacity(max_count);
    let mut idx = *rr_cursor % targets_and_lags.len();
    let mut visited = 0usize;
    while selected.len() < max_count && visited < targets_and_lags.len() {
        let store_id = targets_and_lags[idx].0;
        idx = (idx + 1) % targets_and_lags.len();
        visited += 1;

        let Some(target) = targets.get(&store_id).copied().flatten() else {
            continue;
        };
        let current = store_progresses
            .get(&store_id)
            .copied()
            .unwrap_or(StoreProgress {
                store_id,
                epoch: 1,
                offset: 0,
            });
        if current < target {
            selected.push(store_id);
        }
    }
    *rr_cursor = idx;
    selected
}

fn decide_advance_to(
    target_ts: TimeStamp,
    targets: &StoreWalProgresses,
    predicted: &HashMap<u64, StoreProgress>,
    errors: &mut HashMap<u64, FetchWalError>,
    tolerate_store_err: usize,
) -> TargetDecision {
    // Decide what to do with the current target after emitting a fetch batch.
    //
    // Summary of the decision logic:
    // - Iterate all stores in `targets` and classify each store as:
    //   - `CaughtUp`: `predicted[store] >= target` (remove any stale error entry).
    //   - `Behind`: `predicted[store] < target` and no recorded error for this
    //     store.
    //   - `Error`:  `predicted[store] < target` and there is a recorded error for
    //     this store.
    //   - `Unhealthy`: `target=None` (store progress unavailable for this target).
    // - If any store is `Behind`, keep the target active (`Keep`) so we can fetch
    //   more WAL.
    // - Otherwise (all remaining stores are either `CaughtUp`, `Error`, or
    //   `Unhealthy`):
    //   - Advance the target (`Advance`) if the number of `Error` + `Unhealthy`
    //     stores is within `tolerate_store_err`.
    //   - If beyond tolerance, drop the target (`Drop`) only when there exists at
    //     least one `Unhealthy` store (`target=None`), since that condition is
    //     treated as non-retryable for the current target. Otherwise keep it
    //     (`Keep`) and retry later.
    //
    // Assumptions / correctness boundary:
    // - Advancing `target_ts` is allowed even when up to `tolerate_store_err`
    //   stores have not reached their target progress for this target (either
    //   because `target=None` or due to a recorded fetch error). This relies on
    //   merged_engine's quorum-based commit inference (typically 3 replicas,
    //   quorum=2) and on the operational expectation that ignoring one store does
    //   not violate CDC ordering guarantees.
    // - `predicted` must reflect the store progresses after applying the batch
    //   being emitted (including `rotate` decisions). Incorrect `predicted` can
    //   cause premature `Advance`.
    // - `tolerate_store_err` is expected to be small (0/1). Increasing it changes
    //   correctness and should be revisited alongside merged_engine quorum
    //   assumptions.
    let mut error_count = 0usize;
    let mut behind_cont = 0usize;
    let mut has_store_unhealthy = false;

    for (&store_id, target) in targets {
        let Some(target) = target else {
            has_store_unhealthy = true;
            error_count += 1;
            continue;
        };

        let current = predicted.get(&store_id).copied().unwrap_or(StoreProgress {
            store_id,
            epoch: 1,
            offset: 0,
        });
        if current >= *target {
            errors.remove(&store_id);
            continue;
        }

        if errors.contains_key(&store_id) {
            error_count += 1;
        } else {
            behind_cont += 1;
        }
    }

    if behind_cont > 0 {
        return TargetDecision::Keep;
    }

    if error_count <= tolerate_store_err {
        return TargetDecision::Advance(target_ts);
    }

    // Beyond tolerance:
    // - StoreUnhealthy is not retryable for this target: drop target and leaving
    //   `active` cleared by the caller.
    // - Other errors are retryable: keep the current target active.
    if has_store_unhealthy {
        return TargetDecision::Drop;
    }
    TargetDecision::Keep
}

#[cfg(test)]
mod tests {
    use std::io::Write;

    use test_pd_client::TestPdClient;
    use tikv_util::config::ReadableSize;

    use super::*;
    use crate::error::BrError;

    #[test]
    fn test_round_robin_selection() {
        let progress = |store_id: u64, epoch: u32, offset: u64| StoreProgress {
            store_id,
            epoch,
            offset,
        };

        // Empty input / zero count: should not change the cursor.
        {
            let targets = StoreWalProgresses::default();
            let progresses = HashMap::default();

            let mut rr = 123;
            let picked = select_round_robin_stores(&[], 1, &progresses, &targets, &mut rr);
            assert!(picked.is_empty());
            assert_eq!(rr, 123);

            let targets_and_lags = [(
                1,
                StoreLag {
                    epoch_lag: 0,
                    offset_lag: 0,
                },
            )];
            let picked =
                select_round_robin_stores(&targets_and_lags, 0, &progresses, &targets, &mut rr);
            assert!(picked.is_empty());
            assert_eq!(rr, 123);
        }

        // Skip stores that are:
        // - Unhealthy (`target=None`)
        // - Already caught up (`current >= target`)
        // - Missing progress (defaults to epoch=1, offset=0)
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(1, Some(progress(1, 1, 10))); // behind (default progress)
            targets.insert(2, Some(progress(2, 1, 10))); // caught up
            targets.insert(3, None); // unhealthy

            let mut progresses = HashMap::default();
            progresses.insert(2, progress(2, 1, 10));

            let targets_and_lags = vec![
                (
                    1,
                    StoreLag {
                        epoch_lag: 0,
                        offset_lag: 0,
                    },
                ),
                (
                    2,
                    StoreLag {
                        epoch_lag: 0,
                        offset_lag: 0,
                    },
                ),
                (
                    3,
                    StoreLag {
                        epoch_lag: 0,
                        offset_lag: 0,
                    },
                ),
            ];
            let mut rr = 5; // Starts from store 3 after modulo.
            let picked =
                select_round_robin_stores(&targets_and_lags, 2, &progresses, &targets, &mut rr);
            assert_eq!(picked, vec![1]);
            assert_eq!(rr, 2);
        }

        // Cursor advances with early exit (selected reaches max_count).
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(10, Some(progress(10, 1, 10)));
            targets.insert(11, Some(progress(11, 1, 10)));
            targets.insert(12, Some(progress(12, 1, 10)));
            let progresses = HashMap::default(); // all stores behind by default

            let targets_and_lags = vec![
                (
                    10,
                    StoreLag {
                        epoch_lag: 0,
                        offset_lag: 0,
                    },
                ),
                (
                    11,
                    StoreLag {
                        epoch_lag: 0,
                        offset_lag: 0,
                    },
                ),
                (
                    12,
                    StoreLag {
                        epoch_lag: 0,
                        offset_lag: 0,
                    },
                ),
            ];

            let mut rr = 0;
            let picked =
                select_round_robin_stores(&targets_and_lags, 2, &progresses, &targets, &mut rr);
            assert_eq!(picked, vec![10, 11]);
            assert_eq!(rr, 2);

            let picked =
                select_round_robin_stores(&targets_and_lags, 2, &progresses, &targets, &mut rr);
            assert_eq!(picked, vec![12, 10]);
            assert_eq!(rr, 1);

            let picked =
                select_round_robin_stores(&targets_and_lags, 2, &progresses, &targets, &mut rr);
            assert_eq!(picked, vec![11, 12]);
            assert_eq!(rr, 0);
        }
    }

    #[test]
    fn test_decide_advance_to() {
        let target_ts = TimeStamp::new(42);
        let progress = |store_id: u64, epoch: u32, offset: u64| StoreProgress {
            store_id,
            epoch,
            offset,
        };

        // A store behind without error should keep the target.
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(1, Some(progress(1, 1, 10)));

            let mut predicted = HashMap::default();
            predicted.insert(1, progress(1, 1, 0));

            let mut errors = HashMap::default();
            let decision = decide_advance_to(target_ts, &targets, &predicted, &mut errors, 0);
            assert_eq!(decision, TargetDecision::Keep);
            assert!(errors.is_empty());
        }

        // All stores caught up should advance the target, and clear stale errors.
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(1, Some(progress(1, 1, 10)));

            let mut predicted = HashMap::default();
            predicted.insert(1, progress(1, 1, 10));

            let mut errors = HashMap::default();
            errors.insert(
                1,
                FetchWalError {
                    err: Error::StoreTimeout("timeout".to_string()),
                    target: Some(progress(1, 1, 10)),
                },
            );
            let decision = decide_advance_to(target_ts, &targets, &predicted, &mut errors, 0);
            assert_eq!(decision, TargetDecision::Advance(target_ts));
            assert!(errors.is_empty());
        }

        // Stores behind with errors can still be advanced when within tolerance.
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(1, Some(progress(1, 2, 10)));
            targets.insert(2, Some(progress(2, 2, 10)));

            let mut predicted = HashMap::default();
            predicted.insert(1, progress(1, 2, 10));
            predicted.insert(2, progress(2, 2, 0));

            let mut errors = HashMap::default();
            errors.insert(
                2,
                FetchWalError {
                    err: Error::StoreTimeout("timeout".to_string()),
                    target: Some(progress(2, 2, 10)),
                },
            );
            let decision = decide_advance_to(target_ts, &targets, &predicted, &mut errors, 1);
            assert_eq!(decision, TargetDecision::Advance(target_ts));
            assert!(errors.contains_key(&2));
        }

        // Unhealthy stores beyond tolerance drop the target.
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(1, None);
            targets.insert(2, Some(progress(2, 2, 10)));

            let mut predicted = HashMap::default();
            predicted.insert(2, progress(2, 2, 10));

            let mut errors = HashMap::default();
            errors.insert(
                1,
                FetchWalError {
                    err: Error::StoreUnhealthy { store_id: 1 },
                    target: None,
                },
            );

            let decision = decide_advance_to(target_ts, &targets, &predicted, &mut errors, 0);
            assert_eq!(decision, TargetDecision::Drop);
        }

        // Retryable errors beyond tolerance keep the target (no unhealthy store).
        {
            let mut targets = StoreWalProgresses::default();
            targets.insert(1, Some(progress(1, 2, 10)));

            let mut predicted = HashMap::default();
            predicted.insert(1, progress(1, 2, 0));

            let mut errors = HashMap::default();
            errors.insert(
                1,
                FetchWalError {
                    err: Error::StoreTimeout("timeout".to_string()),
                    target: Some(progress(1, 2, 10)),
                },
            );

            let decision = decide_advance_to(target_ts, &targets, &predicted, &mut errors, 0);
            assert_eq!(decision, TargetDecision::Keep);
        }
    }

    #[test]
    fn test_force_advance_eligible_error() {
        let err = FetchWalError {
            err: Error::BrError(BrError::WalChunkIntegrityError("oops".to_string())),
            target: Some(StoreProgress {
                store_id: 7,
                epoch: 3,
                offset: 123,
            }),
        };
        assert!(err.is_force_advance_eligible());

        let err = FetchWalError {
            err: Error::StoreTimeout("timeout".to_string()),
            target: Some(StoreProgress {
                store_id: 7,
                epoch: 3,
                offset: 123,
            }),
        };
        assert!(!err.is_force_advance_eligible());
    }

    struct FakeIo {
        worker_dir: PathBuf,
    }

    #[async_trait]
    impl WalFetcherIo for FakeIo {
        async fn fetch_from_store(
            &self,
            store_id: u64,
            _status_address: &str,
            epoch_id: u32,
            start_off: u64,
            end_off: u64,
        ) -> Result<(StatusCode, TempLocalObject, u64)> {
            let read_len = end_off.saturating_sub(start_off);
            let spool_dir = spool_dir_for_store(&self.worker_dir, store_id);
            box_try!(std::fs::create_dir_all(&spool_dir));
            let mut obj = create_temp_wal_spool_object(&spool_dir, store_id, epoch_id, start_off)?;
            let mut remaining = read_len;
            let buf = [0u8; 4096];
            while remaining > 0 {
                let n = (buf.len() as u64).min(remaining) as usize;
                obj.write_all(&buf[..n])
                    .map_err(|e| -> Error { box_err!("write fake spool failed: {:?}", e) })?;
                remaining -= n as u64;
            }
            obj.flush()
                .map_err(|e| -> Error { box_err!("flush fake spool failed: {:?}", e) })?;
            obj.close();
            Ok((StatusCode::PARTIAL_CONTENT, obj, read_len))
        }

        async fn fetch_from_dfs(
            &self,
            _store_id: u64,
            _epoch_id: u32,
            _target: StoreProgress,
        ) -> Result<(SharedWal, u64, bool)> {
            Err(box_err!("unexpected dfs fetch in test"))
        }
    }

    #[test]
    fn test_merge_store_progresses_never_regresses() {
        let dir = test_util::temp_dir(Some("wal_fetcher_merge_progress"), false);
        let (_wal_targets_tx, wal_targets_rx) = wal_target_channel();
        let (apply_batch_tx, _apply_batch_rx) = mpsc::unbounded_channel();
        let (_request_tx, request_rx) = mpsc::unbounded_channel();

        let (_stop_tx, stop_rx) = watch::channel(false);
        let wal_config = WalConfig {
            tolerate_store_err: false,
            tikv_rfengine_wal_size: ReadableSize(1),
            update_stores_wal_size_limit: 1.into(),
            max_fetch_concurrency: 1,
            ..Default::default()
        };
        let mut fetcher = WalFetcher::new(
            Arc::new(FakeIo {
                worker_dir: dir.path().to_path_buf(),
            }),
            Arc::new(TestPdClient::new(1, false)),
            wal_targets_rx,
            wal_config,
            apply_batch_tx,
            request_rx,
            stop_rx,
        );

        fetcher.merge_store_progresses(vec![StoreProgress {
            store_id: 1,
            epoch: 1,
            offset: 100,
        }]);
        fetcher.merge_store_progresses(vec![StoreProgress {
            store_id: 1,
            epoch: 1,
            offset: 50,
        }]);
        assert_eq!(fetcher.store_progresses.get(&1).unwrap().offset, 100);

        fetcher.merge_store_progresses(vec![StoreProgress {
            store_id: 1,
            epoch: 2,
            offset: 0,
        }]);
        fetcher.merge_store_progresses(vec![StoreProgress {
            store_id: 1,
            epoch: 1,
            offset: 200,
        }]);
        assert_eq!(fetcher.store_progresses.get(&1).unwrap().epoch, 2);
    }
}