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

use std::{
    collections::HashSet,
    future::Future,
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

use api_version::ApiV2;
use bytes::Buf;
use cloud_encryption::MasterKey;
use cloud_server::status_server::{
    ONLINE_SNAPSHOT_CAPABILITY_VERSION, OnlineSnapshotFileRefsResponse,
};
use futures::{StreamExt, TryStreamExt, stream};
use http::Request;
use hyper::Body;
use kvengine::{
    dfs::{Dfs, FileType},
    read::ONLINE_SNAPSHOT_MEM_DATA_FORMAT_VERSION,
};
use kvenginepb::ChangeSet;
use kvproto::{
    coprocessor::DelegateResponse,
    metapb::{Peer, Region, Store},
};
use pd_client::{
    PdClient,
    pd_control::{OpKind, Operator, PdControl},
};
use protobuf::Message;
use security::{HttpClient, SecurityManager};
use slog_global::{info, warn};
use tikv::storage::mvcc::Key;
use tikv_util::{box_err, box_try};
use tokio::{runtime::Runtime, sync::watch, task::JoinHandle};
use uuid::Uuid;

use crate::{
    common::{RawRegion, TableFile, send_request_to_store},
    error::{Error, Result},
    limiter::ThroughputLimiter,
    restore::RestoreConfig,
    restore_keyspace::{
        OnlineSnapshotPlacementPlan, OnlineSnapshotShard, ReportRestoreStepTrait, RestoreDeadline,
        RestoreScheduleGuard, RestoreScheduleGuardState, RestoreStep, RestoredKeyspace,
        RetainFileFilter, check_schedule_guards, restore_captured_online_snapshots, scan_regions,
    },
};

const LEASE_RENEW_INTERVAL: Duration = Duration::from_secs(20);
const SAFE_POINT_TTL: Duration = Duration::from_secs(60);
const ONLINE_SNAPSHOT_REQUEST_CONCURRENCY: usize = 32;
const SOURCE_OPERATOR_POLL_INTERVAL: Duration = Duration::from_millis(500);
const SOURCE_REGION_RETRY_INTERVAL: Duration = Duration::from_millis(200);
const SOURCE_FILE_REFS_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const SOURCE_FILE_REFS_PHASE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const SCHEDULE_GUARD_CHECK_INTERVAL: Duration = Duration::from_secs(1);

async fn wait_while_schedule_guards_healthy<F, T>(
    scheduling_guards: &[RestoreScheduleGuardState],
    future: F,
) -> Result<T>
where
    F: Future<Output = Result<T>>,
{
    tokio::pin!(future);
    loop {
        check_schedule_guards(scheduling_guards)?;
        tokio::select! {
            result = &mut future => {
                check_schedule_guards(scheduling_guards)?;
                return result;
            }
            _ = tokio::time::sleep(SCHEDULE_GUARD_CHECK_INTERVAL) => {}
        }
    }
}

struct KeyspaceSafePointLease {
    pd_client: Arc<dyn PdClient>,
    keyspace_id: u32,
    service_id: String,
    healthy: Arc<AtomicBool>,
    stop_tx: watch::Sender<bool>,
    renew_task: JoinHandle<()>,
}

impl KeyspaceSafePointLease {
    async fn acquire(
        pd_client: Arc<dyn PdClient>,
        keyspace_id: u32,
        service_id: String,
        branch_ts: u64,
    ) -> Result<Self> {
        let current = pd_client
            .update_keyspace_service_safe_point(
                keyspace_id,
                service_id.clone(),
                branch_ts.into(),
                SAFE_POINT_TTL,
            )
            .await?;
        if current > branch_ts {
            return Err(box_err!(
                "source keyspace GC safe point {} has passed branch_ts {}",
                current,
                branch_ts
            ));
        }

        let healthy = Arc::new(AtomicBool::new(true));
        let renew_healthy = healthy.clone();
        let renew_pd_client = pd_client.clone();
        let renew_service_id = service_id.clone();
        let (stop_tx, mut stop_rx) = watch::channel(false);
        let renew_task = tokio::spawn(async move {
            loop {
                tokio::select! {
                    changed = stop_rx.changed() => {
                        if changed.is_err() || *stop_rx.borrow() {
                            break;
                        }
                    }
                    _ = tokio::time::sleep(LEASE_RENEW_INTERVAL) => {
                        if let Err(err) = renew_pd_client
                            .update_keyspace_service_safe_point(
                                keyspace_id,
                                renew_service_id.clone(),
                                branch_ts.into(),
                                SAFE_POINT_TTL,
                            )
                            .await
                        {
                            renew_healthy.store(false, Ordering::Release);
                            warn!(
                                "failed to renew online snapshot GC safe point";
                                "keyspace_id" => keyspace_id,
                                "branch_ts" => branch_ts,
                                "err" => ?err,
                            );
                        }
                    }
                }
            }
        });
        Ok(Self {
            pd_client,
            keyspace_id,
            service_id,
            healthy,
            stop_tx,
            renew_task,
        })
    }

    fn check_healthy(&self) -> Result<()> {
        if self.healthy.load(Ordering::Acquire) {
            Ok(())
        } else {
            Err(box_err!("online snapshot GC safe point renewal failed"))
        }
    }

    async fn close(self) {
        let _ = self.stop_tx.send(true);
        let _ = self.renew_task.await;
        if let Err(err) = self
            .pd_client
            .remove_keyspace_service_safe_point(self.keyspace_id, self.service_id)
            .await
        {
            warn!(
                "failed to remove online snapshot GC safe point";
                "keyspace_id" => self.keyspace_id,
                "err" => ?err,
            );
        }
    }
}

async fn prepare_store(
    security_manager: &SecurityManager,
    client: &HttpClient,
    store: &Store,
    branch_ts: u64,
    timeout: Duration,
) -> Result<()> {
    let uri = box_try!(security_manager.build_uri(format!(
        "{}/kvengine/online-snapshot/prepare?branch_ts={}&wait_timeout_secs={}",
        store.get_status_address(),
        branch_ts,
        timeout.as_secs(),
    )));
    let request = Request::post(uri).body(Body::empty()).unwrap();
    send_request_to_store(request, store, client, timeout).await?;
    Ok(())
}

async fn scan_source_regions(
    pd_client: &dyn PdClient,
    keyspace_id: u32,
) -> Result<Vec<(Region, Peer)>> {
    let (start, end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
    let scanned_regions = scan_regions(pd_client, &start, &end).await?;
    if scanned_regions.is_empty() {
        return Err(Error::BackupEmptyForKeyspace(keyspace_id));
    }
    let mut region_ids = HashSet::with_capacity(scanned_regions.len());
    scanned_regions
        .into_iter()
        .map(|mut scanned_region| {
            let region = scanned_region.take_region();
            let leader = scanned_region.take_leader();
            if !region_ids.insert(region.get_id()) {
                return Err(box_err!(
                    "source Region {} appeared more than once while scanning keyspace {}",
                    region.get_id(),
                    keyspace_id
                ));
            }
            if leader.get_id() == 0 {
                return Err(box_err!("source Region {} has no leader", region.get_id()));
            }
            Ok((region, leader))
        })
        .collect()
}

#[derive(Clone, Debug, PartialEq, Eq)]
struct SourceRegionVersion {
    id: u64,
    version: u64,
    start_key: Vec<u8>,
    end_key: Vec<u8>,
}

impl From<&Region> for SourceRegionVersion {
    fn from(region: &Region) -> Self {
        Self {
            id: region.get_id(),
            version: region.get_region_epoch().get_version(),
            start_key: region.get_start_key().to_vec(),
            end_key: region.get_end_key().to_vec(),
        }
    }
}

fn source_region_versions(regions: &[(Region, Peer)]) -> Vec<SourceRegionVersion> {
    regions
        .iter()
        .map(|(region, _)| SourceRegionVersion::from(region))
        .collect()
}

fn invalid_source_region_version_transition(detail: impl std::fmt::Display) -> Error {
    box_err!(
        "source Region version did not follow a split-only transition during online snapshot branching; RestoreShard or a merge was observed: {}",
        detail
    )
}

fn validate_source_region_descendants(
    parent: &SourceRegionVersion,
    descendants: &[SourceRegionVersion],
) -> Result<()> {
    match descendants {
        [] => Err(invalid_source_region_version_transition(format!(
            "Region {}:{} has no descendants",
            parent.id, parent.version
        ))),
        [region] if region == parent => Ok(()),
        [region] => Err(invalid_source_region_version_transition(format!(
            "expected Region {}:{} [{}, {}), got Region {}:{} [{}, {})",
            parent.id,
            parent.version,
            log_wrappers::hex_encode_upper(&parent.start_key),
            log_wrappers::hex_encode_upper(&parent.end_key),
            region.id,
            region.version,
            log_wrappers::hex_encode_upper(&region.start_key),
            log_wrappers::hex_encode_upper(&region.end_key),
        ))),
        [first, .., last]
            if first.start_key != parent.start_key || last.end_key != parent.end_key =>
        {
            Err(invalid_source_region_version_transition(format!(
                "Region {}:{} boundary [{}, {}) was not preserved",
                parent.id,
                parent.version,
                log_wrappers::hex_encode_upper(&parent.start_key),
                log_wrappers::hex_encode_upper(&parent.end_key),
            )))
        }
        [_, .., last] => {
            // RFStore batch split creates N-1 new Regions, increments every
            // resulting Region version by N-1, and keeps the parent Region ID
            // on the rightmost derived Region.
            let split_increment = u64::try_from(descendants.len() - 1).map_err(|_| {
                invalid_source_region_version_transition("split child count overflowed u64")
            })?;
            let expected_version =
                parent.version.checked_add(split_increment).ok_or_else(|| {
                    invalid_source_region_version_transition("split Region version overflowed u64")
                })?;
            if last.id != parent.id {
                return Err(invalid_source_region_version_transition(format!(
                    "Region {} was not the rightmost derived Region after split",
                    parent.id
                )));
            }
            if let Some(descendant) = descendants
                .iter()
                .find(|region| region.version != expected_version)
            {
                return Err(invalid_source_region_version_transition(format!(
                    "Region {}:{} split into {} Regions, so child Region {} expected version {}, got {}",
                    parent.id,
                    parent.version,
                    descendants.len(),
                    descendant.id,
                    expected_version,
                    descendant.version
                )));
            }
            Ok(())
        }
    }
}

/// Both topologies must come from `scan_source_regions`, which guarantees an
/// ordered, contiguous keyspace cover with unique Region IDs.
fn validate_source_region_version_transition(
    baseline: &[SourceRegionVersion],
    observed: &[SourceRegionVersion],
) -> Result<()> {
    if baseline == observed {
        return Ok(());
    }

    let mut remaining = observed;
    for parent in baseline {
        let Some(last_descendant) = remaining
            .iter()
            .position(|region| region.end_key.as_slice() >= parent.end_key.as_slice())
        else {
            return Err(invalid_source_region_version_transition(format!(
                "Region {}:{} descendants do not cover its range",
                parent.id, parent.version
            )));
        };
        let (descendants, rest) = remaining.split_at(last_descendant + 1);
        validate_source_region_descendants(parent, descendants)?;
        remaining = rest;
    }

    if !remaining.is_empty() {
        return Err(invalid_source_region_version_transition(format!(
            "{} source Regions remain outside the baseline topology",
            remaining.len()
        )));
    }
    Ok(())
}

fn changes_source_topology(operator: &Operator) -> bool {
    let topology_kind_mask = OpKind::OpMerge as u32 | OpKind::OpRegion as u32;
    operator.kind_mask & topology_kind_mask != 0
}

fn source_topology_operator_region_ids(
    operators: &[Operator],
    regions: &[(Region, Peer)],
) -> HashSet<u64> {
    let source_region_ids = regions
        .iter()
        .map(|(region, _)| region.get_id())
        .collect::<HashSet<_>>();
    operators
        .iter()
        .filter(|operator| changes_source_topology(operator))
        .filter_map(|operator| {
            source_region_ids
                .contains(&operator.region_id)
                .then_some(operator.region_id)
        })
        .collect()
}

async fn scan_source_regions_after_operator_drain(
    pd_client: &dyn PdClient,
    pd_control: &PdControl,
    keyspace_id: u32,
    scheduling_guard: &RestoreScheduleGuardState,
    deadline: RestoreDeadline,
) -> Result<Vec<(Region, Peer)>> {
    deadline.remaining("source operator discovery")?;
    scheduling_guard.check_healthy()?;
    // Read the operator list before scanning Regions. Once schedule=deny is
    // installed, no new automatic move or merge operator should enter this
    // keyspace, so only this snapshot can overlap the following Region scan.
    let operators = pd_control.get_operators().await?;
    scheduling_guard.check_healthy()?;
    let regions = scan_source_regions(pd_client, keyspace_id).await?;
    let operator_region_ids = source_topology_operator_region_ids(&operators, &regions);
    if operator_region_ids.is_empty() {
        return Ok(regions);
    }

    info!(
        "waiting for existing source topology operators";
        "source_keyspace_id" => keyspace_id,
        "operators" => operator_region_ids.len(),
    );
    loop {
        deadline.remaining("source operator drain")?;
        scheduling_guard.check_healthy()?;
        let active_count = pd_control
            .get_operators()
            .await?
            .iter()
            .filter(|operator| operator_region_ids.contains(&operator.region_id))
            .filter(|operator| changes_source_topology(operator))
            .count();
        if active_count == 0 {
            break;
        }
        tokio::time::sleep(SOURCE_OPERATOR_POLL_INTERVAL).await;
    }

    deadline.remaining("source Region rescan after operator drain")?;
    scheduling_guard.check_healthy()?;
    let rescanned_regions = scan_source_regions(pd_client, keyspace_id).await?;
    info!(
        "rescanned source Regions after existing topology operators finished";
        "source_keyspace_id" => keyspace_id,
        "previous_regions" => regions.len(),
        "current_regions" => rescanned_regions.len(),
    );
    Ok(rescanned_regions)
}

async fn ensure_source_region_boundaries(
    pd_client: &dyn PdClient,
    keyspace_id: u32,
    timeout: Duration,
) -> Result<()> {
    let (start, end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
    let encoded_start = Key::from_raw(&start).into_encoded();
    let encoded_end = Key::from_raw(&end).into_encoded();
    let mut split_keys = Vec::with_capacity(2);
    for key in [&encoded_start, &encoded_end] {
        let region = pd_client.get_region_async(key).await?;
        if region.get_start_key() != key.as_slice() {
            split_keys.push(key.clone());
        }
    }
    if !split_keys.is_empty() {
        pd_client
            .split_regions_with_retry(split_keys, timeout)
            .await?;
    }
    Ok(())
}

async fn capture_region(
    pd_client: Arc<dyn PdClient>,
    prepared_stores: Arc<tokio::sync::Mutex<HashSet<u64>>>,
    security_manager: &SecurityManager,
    client: &HttpClient,
    mut region: Region,
    mut leader: Peer,
    branch_ts: u64,
    timeout: Duration,
    max_retry: usize,
    scheduling_guards: &[RestoreScheduleGuardState],
    deadline: RestoreDeadline,
) -> Result<OnlineSnapshotShard> {
    let expected_region_id = region.get_id();
    let expected_region_ver = region.get_region_epoch().get_version();
    let mut last_error = None;

    for _ in 0..max_retry.max(1) {
        deadline.remaining("source Region capture")?;
        check_schedule_guards(scheduling_guards)?;
        if region.get_id() != expected_region_id
            || region.get_region_epoch().get_version() != expected_region_ver
        {
            return Err(Error::SourceRegionEpochChanged(expected_region_id));
        }
        let store = pd_client.get_store_async(leader.get_store_id()).await?;
        check_schedule_guards(scheduling_guards)?;
        if !prepared_stores.lock().await.contains(&store.get_id()) {
            // Leader transfer to a previously unseen Store is rare. Serialize
            // this slow path so the Store-wide backup-ts barrier runs once.
            let mut prepared = prepared_stores.lock().await;
            if !prepared.contains(&store.get_id()) {
                wait_while_schedule_guards_healthy(
                    scheduling_guards,
                    prepare_store(
                        security_manager,
                        client,
                        &store,
                        branch_ts,
                        deadline.cap_timeout(timeout, "source Store prepare")?,
                    ),
                )
                .await?;
                prepared.insert(store.get_id());
            }
        }

        let uri = box_try!(security_manager.build_uri(format!(
            "{}/kvengine/online-snapshot/capture/{}?shard_ver={}&branch_ts={}",
            store.get_status_address(),
            region.get_id(),
            region.get_region_epoch().get_version(),
            branch_ts,
        )));
        let request = Request::get(uri).body(Body::empty()).unwrap();
        let request_timeout = deadline.cap_timeout(timeout, "source Region capture")?;
        check_schedule_guards(scheduling_guards)?;
        let capture_result = wait_while_schedule_guards_healthy(
            scheduling_guards,
            send_request_to_store(request, &store, client, request_timeout),
        )
        .await;
        match capture_result {
            Ok((_, body)) => {
                let mut response = DelegateResponse::default();
                box_try!(response.merge_from_bytes(&body));
                if response.has_region_error() {
                    let region_error = response.get_region_error();
                    if region_error.has_epoch_not_match() {
                        return Err(Error::SourceRegionEpochChanged(expected_region_id));
                    }
                    last_error = Some(box_err!(
                        "source Region {} leader changed: {:?}",
                        expected_region_id,
                        region_error
                    ));
                } else if !response.get_other_error().is_empty() {
                    return Err(box_err!(
                        "source Region {} capture failed: {}",
                        expected_region_id,
                        response.get_other_error()
                    ));
                } else if response.has_locked() {
                    return Err(box_err!(
                        "online snapshot capture unexpectedly rejected a lock in Region {}",
                        expected_region_id
                    ));
                } else {
                    let mem_table_data = response.take_mem_table_data();
                    if mem_table_data.len() < std::mem::size_of::<u32>()
                        || mem_table_data.as_slice().get_u32_le()
                            != ONLINE_SNAPSHOT_MEM_DATA_FORMAT_VERSION
                    {
                        return Err(box_err!(
                            "Store {} does not support full-CF online snapshot encoding",
                            store.get_id()
                        ));
                    }
                    let mut change_set = ChangeSet::default();
                    box_try!(change_set.merge_from_bytes(&response.take_snapshot()));
                    if !change_set.has_snapshot()
                        || change_set.get_shard_id() != region.get_id()
                        || change_set.get_shard_ver() != region.get_region_epoch().get_version()
                        || response.get_mem_table_sequence()
                            != change_set.get_snapshot().get_data_sequence()
                    {
                        return Err(box_err!(
                            "online snapshot response does not match Region {}:{}",
                            region.get_id(),
                            region.get_region_epoch().get_version()
                        ));
                    }
                    let raw_region = RawRegion::from(region.clone());
                    if change_set.get_snapshot().get_outer_start() != raw_region.raw_start
                        || change_set.get_snapshot().get_outer_end() != raw_region.raw_end
                    {
                        return Err(box_err!(
                            "online snapshot range does not match source Region {}",
                            region.get_id()
                        ));
                    }
                    return Ok(OnlineSnapshotShard {
                        leader,
                        region,
                        change_set,
                        mem_table_data,
                    });
                }
            }
            Err(err) => {
                last_error = Some(err);
            }
        }

        check_schedule_guards(scheduling_guards)?;
        let Some((new_region, new_leader)) = pd_client
            .get_region_leader_by_id(expected_region_id)
            .await?
        else {
            return Err(Error::SourceRegionEpochChanged(expected_region_id));
        };
        check_schedule_guards(scheduling_guards)?;
        region = new_region;
        leader = new_leader;
        tokio::time::sleep(SOURCE_REGION_RETRY_INTERVAL).await;
    }

    Err(last_error.unwrap_or_else(|| {
        box_err!(
            "source Region {} capture exhausted retries",
            expected_region_id
        )
    }))
}

fn has_same_source_region_layout(previous: &[(Region, Peer)], current: &[(Region, Peer)]) -> bool {
    previous.len() == current.len()
        && previous
            .iter()
            .zip(current)
            .all(|((previous, _), (current, _))| {
                previous.get_id() == current.get_id()
                    && previous.get_start_key() == current.get_start_key()
                    && previous.get_end_key() == current.get_end_key()
            })
}

#[allow(clippy::too_many_arguments)]
async fn capture_source_regions_allowing_splits(
    pd_client: Arc<dyn PdClient>,
    prepared_stores: Arc<tokio::sync::Mutex<HashSet<u64>>>,
    security_manager: &SecurityManager,
    client: &HttpClient,
    source_keyspace_id: u32,
    mut regions: Vec<(Region, Peer)>,
    branch_ts: u64,
    timeout: Duration,
    max_retry: usize,
    capture_concurrency: usize,
    scheduling_guards: &[RestoreScheduleGuardState],
    deadline: RestoreDeadline,
) -> Result<Vec<OnlineSnapshotShard>> {
    loop {
        check_schedule_guards(scheduling_guards)?;
        deadline.remaining("online snapshot capture")?;
        let previous_regions = regions.clone();
        let results = stream::iter(regions)
            .map(|(region, leader)| {
                capture_region(
                    pd_client.clone(),
                    prepared_stores.clone(),
                    security_manager,
                    client,
                    region,
                    leader,
                    branch_ts,
                    timeout,
                    max_retry,
                    scheduling_guards,
                    deadline,
                )
            })
            .buffer_unordered(capture_concurrency)
            .collect::<Vec<_>>()
            .await;

        let mut captured = Vec::with_capacity(results.len());
        let mut epoch_changed = false;
        let mut first_error = None;
        for result in results {
            match result {
                Ok(snapshot) => captured.push(snapshot),
                Err(Error::SourceRegionEpochChanged(_)) => epoch_changed = true,
                Err(err) if first_error.is_none() => first_error = Some(err),
                Err(_) => {}
            }
        }
        if !epoch_changed {
            return first_error.map_or_else(|| Ok(captured), Err);
        }

        check_schedule_guards(scheduling_guards)?;
        let current_regions = wait_while_schedule_guards_healthy(
            scheduling_guards,
            scan_source_regions(pd_client.as_ref(), source_keyspace_id),
        )
        .await?;
        if has_same_source_region_layout(&previous_regions, &current_regions) {
            info!(
                "source Region version changed during online snapshot capture; recapture before post-validation";
                "source_regions" => current_regions.len(),
                "branch_ts" => branch_ts,
            );
        } else {
            validate_source_region_version_transition(
                &source_region_versions(&previous_regions),
                &source_region_versions(&current_regions),
            )?;
            info!(
                "source Regions split during online snapshot capture; recapture descendants";
                "previous_regions" => previous_regions.len(),
                "current_regions" => current_regions.len(),
                "branch_ts" => branch_ts,
            );
        }
        tokio::time::sleep(SOURCE_REGION_RETRY_INTERVAL).await;
        regions = current_regions;
    }
}

async fn prepare_current_leader_stores(
    pd_client: Arc<dyn PdClient>,
    security_manager: &SecurityManager,
    client: &HttpClient,
    regions: &[(Region, Peer)],
    branch_ts: u64,
    timeout: Duration,
    scheduling_guards: &[RestoreScheduleGuardState],
    deadline: RestoreDeadline,
) -> Result<HashSet<u64>> {
    let store_ids = regions
        .iter()
        .map(|(_, leader)| leader.get_store_id())
        .collect::<HashSet<_>>();
    let prepared = stream::iter(store_ids)
        .map(|store_id| {
            let pd_client = pd_client.clone();
            async move {
                check_schedule_guards(scheduling_guards)?;
                let store = pd_client.get_store_async(store_id).await?;
                check_schedule_guards(scheduling_guards)?;
                wait_while_schedule_guards_healthy(
                    scheduling_guards,
                    prepare_store(
                        security_manager,
                        client,
                        &store,
                        branch_ts,
                        deadline.cap_timeout(timeout, "source Store prepare")?,
                    ),
                )
                .await?;
                Ok::<_, Error>(store_id)
            }
        })
        .buffer_unordered(ONLINE_SNAPSHOT_REQUEST_CONCURRENCY)
        .try_collect::<Vec<_>>()
        .await?;
    Ok(prepared.into_iter().collect())
}

struct SourceFileRefsFilter<'a> {
    pd_client: Arc<dyn PdClient>,
    safe_point: &'a KeyspaceSafePointLease,
    initial_source_versions: &'a [SourceRegionVersion],
    security_manager: Arc<SecurityManager>,
    client: HttpClient,
    runtime: &'a Runtime,
    source_keyspace_id: u32,
    timeout: Duration,
    max_retry: usize,
    deadline: RestoreDeadline,
}

impl SourceFileRefsFilter<'_> {
    async fn fetch_region_file_refs(
        &self,
        mut region: Region,
        mut leader: Peer,
        deadline: RestoreDeadline,
    ) -> Result<Vec<(FileType, u64)>> {
        let expected_region_id = region.get_id();
        let expected_region_ver = region.get_region_epoch().get_version();
        let mut last_error = None;

        for _ in 0..self.max_retry.max(1) {
            deadline.remaining("source file_refs dump")?;
            let store = self
                .pd_client
                .get_store_async(leader.get_store_id())
                .await?;
            let uri = box_try!(self.security_manager.build_uri(format!(
                "{}/kvengine/online-snapshot/file-refs/{}?shard_ver={}",
                store.get_status_address(),
                region.get_id(),
                region.get_region_epoch().get_version(),
            )));
            let request = Request::get(uri).body(Body::empty()).unwrap();
            fail::fail_point!("online_snapshot_source_file_refs_request", |_| {
                Err(box_err!("injected source file_refs request failure"))
            });
            let request_timeout = deadline.cap_timeout(self.timeout, "source file_refs dump")?;
            match tokio::time::timeout(
                request_timeout,
                send_request_to_store(request, &store, &self.client, request_timeout),
            )
            .await
            {
                Ok(Ok((_, body))) => {
                    match serde_json::from_slice::<OnlineSnapshotFileRefsResponse>(&body) {
                        Ok(response) => {
                            let raw_region = RawRegion::from(region.clone());
                            if response.version != ONLINE_SNAPSHOT_CAPABILITY_VERSION
                                || response.store_id != store.get_id()
                                || response.shard_id != region.get_id()
                                || response.shard_ver != region.get_region_epoch().get_version()
                                || response.outer_start != raw_region.raw_start
                                || response.outer_end != raw_region.raw_end
                            {
                                return Err(box_err!(
                                    "invalid file_refs response for source Region {}:{} from Store {}",
                                    region.get_id(),
                                    region.get_region_epoch().get_version(),
                                    store.get_id()
                                ));
                            }
                            return response
                                .files
                                .into_iter()
                                .map(|file| {
                                    FileType::from_u8(file.file_type)
                                        .map(|file_type| (file_type, file.file_id))
                                        .ok_or_else(|| {
                                            box_err!(
                                                "source Region {} returned unknown file type {}",
                                                region.get_id(),
                                                file.file_type
                                            )
                                        })
                                })
                                .collect();
                        }
                        Err(json_err) => {
                            let mut delegate = DelegateResponse::default();
                            if delegate.merge_from_bytes(&body).is_ok()
                                && delegate.has_region_error()
                                && delegate.get_region_error().has_epoch_not_match()
                            {
                                return Err(Error::SourceRegionEpochChanged(expected_region_id));
                            }
                            last_error = Some(box_err!(
                                "source Region {} returned an invalid file_refs response: {}",
                                expected_region_id,
                                json_err
                            ));
                        }
                    }
                }
                Ok(Err(err)) => last_error = Some(err),
                Err(_) => {
                    last_error = Some(box_err!(
                        "source Region {} file_refs request timed out after {:?}",
                        expected_region_id,
                        request_timeout
                    ));
                }
            }

            let Some((new_region, new_leader)) = self
                .pd_client
                .get_region_leader_by_id(expected_region_id)
                .await?
            else {
                return Err(box_err!(
                    "source Region {} disappeared during file_refs dump",
                    expected_region_id
                ));
            };
            if new_region.get_region_epoch().get_version() != expected_region_ver {
                return Err(Error::SourceRegionEpochChanged(expected_region_id));
            }
            if new_region.get_start_key() != region.get_start_key()
                || new_region.get_end_key() != region.get_end_key()
            {
                return Err(Error::SourceRegionEpochChanged(expected_region_id));
            }
            region = new_region;
            leader = new_leader;
            tokio::time::sleep(SOURCE_REGION_RETRY_INTERVAL).await;
        }

        Err(last_error.unwrap_or_else(|| {
            box_err!(
                "source Region {} file_refs dump exhausted retries",
                expected_region_id
            )
        }))
    }

    async fn collect_source_file_refs(&self) -> Result<HashSet<(FileType, u64)>> {
        let deadline = self.deadline.capped(SOURCE_FILE_REFS_PHASE_TIMEOUT)?;
        deadline.remaining("source file_refs dump")?;
        self.safe_point.check_healthy()?;
        let mut last_epoch_error = None;
        for _ in 0..self.max_retry.max(1) {
            let regions =
                scan_source_regions(self.pd_client.as_ref(), self.source_keyspace_id).await?;
            let scanned_versions = source_region_versions(&regions);
            let responses = stream::iter(regions)
                .map(|(region, leader)| self.fetch_region_file_refs(region, leader, deadline))
                .buffer_unordered(ONLINE_SNAPSHOT_REQUEST_CONCURRENCY)
                .try_collect::<Vec<_>>()
                .await;
            let responses = match responses {
                Ok(responses) => responses,
                Err(err @ Error::SourceRegionEpochChanged(_)) => {
                    last_epoch_error = Some(err);
                    tokio::time::sleep(SOURCE_REGION_RETRY_INTERVAL).await;
                    continue;
                }
                Err(err) => return Err(err),
            };
            let mut files = HashSet::new();
            for response_files in responses {
                files.extend(response_files);
            }
            self.safe_point.check_healthy()?;
            deadline.remaining("source file_refs dump completion")?;
            validate_source_region_version_transition(
                self.initial_source_versions,
                &scanned_versions,
            )?;
            deadline.remaining("source Region version validation completion")?;
            return Ok(files);
        }
        Err(last_epoch_error
            .unwrap_or_else(|| box_err!("source file_refs dump exhausted topology retries")))
    }
}

fn select_files_to_retain(
    target_files: Vec<TableFile>,
    source_files: &HashSet<(FileType, u64)>,
) -> Vec<TableFile> {
    let mut unique_target_files = HashSet::with_capacity(target_files.len());
    target_files
        .into_iter()
        .filter(|file| unique_target_files.insert((file.ftype, file.id)))
        .filter(|file| !source_files.contains(&(file.ftype, file.id)))
        .collect()
}

impl RetainFileFilter for SourceFileRefsFilter<'_> {
    fn filter_files(&self, target_files: Vec<TableFile>) -> Result<Vec<TableFile>> {
        self.deadline.remaining("selective file retention")?;
        let source_files = self.runtime.block_on(self.collect_source_file_refs())?;
        let target_count = target_files.len();
        let files = select_files_to_retain(target_files, &source_files);
        info!(
            "computed online snapshot selective retain set";
            "source_keyspace_id" => self.source_keyspace_id,
            "target_file_refs" => target_count,
            "source_file_refs" => source_files.len(),
            "retain_file_refs" => files.len(),
        );
        Ok(files)
    }
}

fn validate_online_snapshot_concurrency(phase: &str, concurrency: usize) -> Result<usize> {
    if concurrency == 0 {
        return Err(box_err!(
            "online snapshot {} concurrency must be greater than 0",
            phase
        ));
    }
    Ok(concurrency)
}

/// Creates a near-current branch from live source Region leader snapshots and
/// feeds it into the common keyspace restore pipeline.
#[allow(clippy::too_many_arguments)]
pub fn restore_keyspace_from_online_snapshot_with_cfg(
    config: RestoreConfig,
    keyspace_name: &str,
    target_keyspace_name: &str,
    working_path: PathBuf,
    dfs: Arc<dyn Dfs>,
    pd_client: Arc<dyn PdClient>,
    runtime: &Runtime,
    master_key: Option<MasterKey>,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    limiter: Option<Arc<ThroughputLimiter>>,
) -> Result<RestoredKeyspace> {
    let capture_concurrency = validate_online_snapshot_concurrency(
        "capture",
        config.online_snapshot_capture_concurrency,
    )?;
    validate_online_snapshot_concurrency(
        "placement",
        config.online_snapshot_placement_concurrency,
    )?;
    let mut pd_control = PdControl::new(config.pd.clone(), pd_client.get_security_mgr())?;
    pd_control.set_timeout(config.timeout_pd_control.0);
    let source = runtime.block_on(pd_control.get_keyspace_by_name(keyspace_name))?;
    let target = runtime.block_on(pd_control.get_keyspace_by_name(target_keyspace_name))?;
    if source.id == target.id {
        return Err(box_err!(
            "online snapshot branch requires different source and target keyspaces"
        ));
    }
    if ApiV2::is_default_keyspace(source.id) || ApiV2::is_default_keyspace(target.id) {
        return Err(Error::RestoreWithDefaultKeyspace);
    }
    runtime.block_on(ensure_source_region_boundaries(
        pd_client.as_ref(),
        source.id,
        config.timeout_split_regions.0,
    ))?;

    let deadline = RestoreDeadline::after(config.online_snapshot_attempt_timeout.0)?;
    // RestoreShard advances the Region version at apply. Record the source
    // topology before choosing branch_ts so the final leader-local file_refs
    // observation detects any RestoreShard reflected by that observation.
    let initial_source_versions = {
        let regions = runtime.block_on(scan_source_regions(pd_client.as_ref(), source.id))?;
        source_region_versions(&regions)
    };
    let owner = format!("online-snapshot-branch-{}", Uuid::new_v4());
    let branch_ts = runtime.block_on(pd_client.get_tso())?.into_inner();
    let safe_point = runtime.block_on(KeyspaceSafePointLease::acquire(
        pd_client.clone(),
        source.id,
        owner,
        branch_ts,
    ))?;
    let result = (|| {
        deadline.remaining("online snapshot source guard installation")?;
        let security_manager = pd_client.get_security_mgr();
        let client = box_try!(security_manager.http_client(hyper::Client::builder()));
        let (source_start, source_end) = ApiV2::get_keyspace_range_by_id(source.id);
        let source_ranges = vec![(source_start.as_slice(), source_end.as_slice())];
        let source_schedule_guard =
            RestoreScheduleGuard::install(pd_control.clone(), &source_ranges, true, runtime)?;
        let source_schedule_guard_state = source_schedule_guard.state();
        let regions = runtime.block_on(scan_source_regions_after_operator_drain(
            pd_client.as_ref(),
            &pd_control,
            source.id,
            &source_schedule_guard_state,
            deadline,
        ))?;
        let placement_plan = OnlineSnapshotPlacementPlan::new(&regions)?;
        let (target_start, target_end) = ApiV2::get_keyspace_range_by_id(target.id);
        let target_ranges = vec![(target_start.as_slice(), target_end.as_slice())];
        let target_schedule_guard =
            RestoreScheduleGuard::install(pd_control.clone(), &target_ranges, true, runtime)?;
        let target_schedule_guard_state = target_schedule_guard.state();
        let capture_schedule_guards = [
            source_schedule_guard_state.clone(),
            target_schedule_guard_state,
        ];
        // Install the target guard before the first target mutation, then run
        // one pre-capture reconciliation under both scheduling guards.
        placement_plan.prepare_target(
            target.id,
            pd_client.clone(),
            pd_control.clone(),
            &config,
            &capture_schedule_guards,
            runtime,
            deadline,
        )?;

        reporter.report_step(RestoreStep::CaptureOnlineSnapshot);
        deadline.remaining("online snapshot capture")?;
        let prepared_stores = Arc::new(tokio::sync::Mutex::new(runtime.block_on(
            prepare_current_leader_stores(
                pd_client.clone(),
                &security_manager,
                &client,
                &regions,
                branch_ts,
                config.timeout_restore_snapshot.0,
                &capture_schedule_guards,
                deadline,
            ),
        )?));
        safe_point.check_healthy()?;
        let snapshots = runtime.block_on(capture_source_regions_allowing_splits(
            pd_client.clone(),
            prepared_stores.clone(),
            &security_manager,
            &client,
            source.id,
            regions,
            branch_ts,
            config.timeout_restore_snapshot.0,
            config.max_retry,
            capture_concurrency,
            &capture_schedule_guards,
            deadline,
        ))?;
        source_schedule_guard.check_healthy()?;
        target_schedule_guard.check_healthy()?;
        drop(source_schedule_guard);
        fail::fail_point!("online_snapshot_after_source_capture");
        let prepared_store_ids = runtime.block_on(async { prepared_stores.lock().await.clone() });
        info!(
            "captured online source keyspace";
            "source_keyspace_id" => source.id,
            "target_keyspace_id" => target.id,
            "branch_ts" => branch_ts,
            "regions" => snapshots.len(),
            "prepared_stores" => ?prepared_store_ids,
        );
        safe_point.check_healthy()?;

        let retain_file_filter = SourceFileRefsFilter {
            pd_client: pd_client.clone(),
            safe_point: &safe_point,
            initial_source_versions: &initial_source_versions,
            security_manager,
            client,
            runtime,
            source_keyspace_id: source.id,
            timeout: config
                .timeout_restore_snapshot
                .0
                .min(SOURCE_FILE_REFS_REQUEST_TIMEOUT),
            max_retry: config.max_retry,
            deadline,
        };
        let restored = restore_captured_online_snapshots(
            source.id,
            target.id,
            snapshots,
            working_path,
            dfs,
            config,
            pd_client.clone(),
            pd_control,
            runtime,
            branch_ts,
            master_key,
            reporter,
            limiter,
            Some(&retain_file_filter),
            deadline,
            target_schedule_guard,
        )?;
        safe_point.check_healthy()?;
        deadline.remaining("online snapshot completion")?;
        Ok(restored)
    })();

    runtime.block_on(safe_point.close());
    result
}

#[cfg(test)]
mod tests {
    use test_pd_client::TestPdClient;

    use super::*;

    #[test]
    fn test_online_snapshot_concurrency_must_be_positive() {
        assert_eq!(
            validate_online_snapshot_concurrency(
                "capture",
                crate::restore::DEFAULT_ONLINE_SNAPSHOT_CAPTURE_CONCURRENCY
            )
            .unwrap(),
            32
        );
        assert_eq!(
            validate_online_snapshot_concurrency(
                "placement",
                crate::restore::DEFAULT_ONLINE_SNAPSHOT_PLACEMENT_CONCURRENCY,
            )
            .unwrap(),
            64
        );
        validate_online_snapshot_concurrency("capture", 0).unwrap_err();
        validate_online_snapshot_concurrency("placement", 0).unwrap_err();
    }

    #[test]
    fn test_guarded_wait_stops_when_any_schedule_lease_expires() {
        let runtime = Runtime::new().unwrap();
        let source_guard = RestoreScheduleGuardState::with_valid_until_for_test(
            std::time::Instant::now() + Duration::from_secs(10),
        );
        let target_guard = RestoreScheduleGuardState::with_valid_until_for_test(
            std::time::Instant::now() + Duration::from_millis(20),
        );
        let started_at = std::time::Instant::now();

        let err = runtime
            .block_on(wait_while_schedule_guards_healthy(
                &[source_guard, target_guard],
                std::future::pending::<Result<()>>(),
            ))
            .unwrap_err();

        assert!(err.to_string().contains("lease expired"), "{err:?}");
        assert!(started_at.elapsed() < Duration::from_secs(2));
    }

    fn source_region(id: u64, start: &[u8], end: &[u8]) -> (Region, Peer) {
        let mut peer = Peer::default();
        peer.set_id(id);
        peer.set_store_id(1);
        let mut region = Region::default();
        region.set_id(id);
        region.set_start_key(start.to_vec());
        region.set_end_key(end.to_vec());
        region.mut_peers().push(peer.clone());
        (region, peer)
    }

    fn source_region_with_version(
        id: u64,
        version: u64,
        start: &[u8],
        end: &[u8],
    ) -> (Region, Peer) {
        let (mut region, peer) = source_region(id, start, end);
        region.mut_region_epoch().set_version(version);
        (region, peer)
    }

    fn versions(regions: &[(Region, Peer)]) -> Vec<SourceRegionVersion> {
        regions
            .iter()
            .map(|(region, _)| SourceRegionVersion::from(region))
            .collect()
    }

    #[test]
    fn test_source_operator_filter_waits_only_for_overlapping_topology_changes() {
        let operator = |region_id, kind| Operator {
            region_id,
            kind_mask: kind as u32,
            ..Default::default()
        };
        let regions = vec![
            source_region(1, b"a", b"f"),
            source_region(2, b"f", b"m"),
            source_region(3, b"m", b"t"),
            source_region(4, b"t", b"z"),
        ];
        let operators = vec![
            operator(1, OpKind::OpMerge),
            operator(2, OpKind::OpRegion),
            operator(3, OpKind::OpSplit),
            operator(4, OpKind::OpLeader),
            operator(99, OpKind::OpRegion),
        ];

        assert_eq!(
            source_topology_operator_region_ids(&operators, &regions),
            HashSet::from_iter([1, 2])
        );
    }

    #[test]
    fn test_source_region_version_transition_detects_restore_shard() {
        let initial = vec![
            source_region_with_version(1, 7, b"a", b"m"),
            source_region_with_version(2, 11, b"m", b"z"),
        ];
        let current = vec![
            source_region_with_version(1, 8, b"a", b"m"),
            source_region_with_version(2, 11, b"m", b"z"),
        ];

        let err =
            validate_source_region_version_transition(&versions(&initial), &versions(&current))
                .unwrap_err();
        assert!(err.to_string().contains("RestoreShard"), "{err:?}");
        assert!(err.to_string().contains("Region 1:7"), "{err:?}");
        assert!(err.to_string().contains("Region 1:8"), "{err:?}");
    }

    #[test]
    fn test_source_region_version_transition_ignores_leader_change() {
        let initial = vec![source_region_with_version(1, 7, b"a", b"z")];
        let mut current = initial.clone();
        current[0].1.set_id(99);
        current[0].1.set_store_id(9);

        validate_source_region_version_transition(&versions(&initial), &versions(&current))
            .unwrap();
    }

    #[test]
    fn test_source_region_version_transition_accepts_batch_split() {
        let initial = vec![
            source_region_with_version(1, 5, b"a", b"g"),
            source_region_with_version(2, 7, b"g", b"t"),
            source_region_with_version(3, 11, b"t", b"z"),
        ];
        let current = vec![
            source_region_with_version(1, 5, b"a", b"g"),
            source_region_with_version(4, 9, b"g", b"m"),
            source_region_with_version(5, 9, b"m", b"q"),
            source_region_with_version(2, 9, b"q", b"t"),
            source_region_with_version(3, 11, b"t", b"z"),
        ];

        validate_source_region_version_transition(&versions(&initial), &versions(&current))
            .unwrap();
    }

    #[test]
    fn test_source_region_version_transition_rejects_merge() {
        let initial = vec![
            source_region_with_version(1, 7, b"a", b"m"),
            source_region_with_version(2, 11, b"m", b"z"),
        ];
        let current = vec![source_region_with_version(2, 12, b"a", b"z")];

        let err =
            validate_source_region_version_transition(&versions(&initial), &versions(&current))
                .unwrap_err();
        assert!(
            err.to_string().contains("RestoreShard or a merge"),
            "{err:?}"
        );
    }

    #[test]
    fn test_source_region_version_transition_rejects_second_split_generation() {
        let initial = vec![source_region_with_version(1, 7, b"a", b"z")];
        let first_split = vec![
            source_region_with_version(2, 8, b"a", b"m"),
            source_region_with_version(1, 8, b"m", b"z"),
        ];
        let second_split = vec![
            source_region_with_version(3, 9, b"a", b"g"),
            source_region_with_version(2, 9, b"g", b"m"),
            source_region_with_version(1, 8, b"m", b"z"),
        ];

        validate_source_region_version_transition(&versions(&initial), &versions(&first_split))
            .unwrap();
        let err = validate_source_region_version_transition(
            &versions(&initial),
            &versions(&second_split),
        )
        .unwrap_err();
        assert!(err.to_string().contains("expected version 9"), "{err:?}");
        assert!(err.to_string().contains("got 8"), "{err:?}");
    }

    #[test]
    fn test_source_region_version_transition_rejects_split_with_extra_increment() {
        let initial = vec![source_region_with_version(1, 7, b"a", b"z")];
        let current = vec![
            source_region_with_version(2, 8, b"a", b"m"),
            source_region_with_version(1, 9, b"m", b"z"),
        ];

        let err =
            validate_source_region_version_transition(&versions(&initial), &versions(&current))
                .unwrap_err();
        assert!(err.to_string().contains("expected version 8"), "{err:?}");
        assert!(err.to_string().contains("got 9"), "{err:?}");
    }

    #[test]
    fn test_same_source_region_layout_ignores_version_change() {
        let previous = vec![source_region_with_version(1, 7, b"a", b"z")];
        let current = vec![source_region_with_version(1, 8, b"a", b"z")];

        assert!(has_same_source_region_layout(&previous, &current));
    }

    #[tokio::test]
    async fn test_scan_source_regions_paginates_pd_requests() {
        let keyspace_id = 42;
        let region_count = crate::restore_keyspace::SCAN_REGION_BATCH_SIZE + 1;
        let (start, end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
        let mut boundaries = Vec::with_capacity(region_count + 1);
        boundaries.push(Key::from_raw(&start).into_encoded());
        for index in 1..region_count {
            let mut key = start.clone();
            key.extend_from_slice(&(index as u16).to_be_bytes());
            boundaries.push(Key::from_raw(&key).into_encoded());
        }
        boundaries.push(Key::from_raw(&end).into_encoded());

        let pd_client = TestPdClient::new(1, false);
        pd_client.set_bootstrap(true);
        for index in 0..region_count {
            let mut leader = Peer::default();
            leader.set_id(index as u64 + 1);
            leader.set_store_id(1);
            let mut region = Region::default();
            region.set_id(index as u64 + 1);
            region.set_start_key(boundaries[index].clone());
            region.set_end_key(boundaries[index + 1].clone());
            region.mut_peers().push(leader.clone());
            pd_client.add_region(&region);
            pd_client
                .region_heartbeat(
                    1,
                    region,
                    leader,
                    pd_client::RegionStat::default(),
                    None,
                    None,
                )
                .await
                .unwrap();
        }

        let regions = scan_source_regions(&pd_client, keyspace_id).await.unwrap();
        assert_eq!(regions.len(), region_count);
        assert_eq!(regions.first().unwrap().0.get_start_key(), boundaries[0]);
        assert_eq!(
            regions.last().unwrap().0.get_end_key(),
            boundaries[region_count]
        );
    }

    #[test]
    fn test_select_files_to_retain_is_typed_and_deduplicated() {
        let file = |id, ftype| TableFile {
            id,
            ftype,
            shard_id: 1,
            size: 10,
        };
        let target = vec![
            file(1, FileType::Sst),
            file(1, FileType::Sst),
            file(1, FileType::Blob),
            file(2, FileType::TxnChunk),
        ];
        let source = HashSet::from([(FileType::Sst, 1), (FileType::TxnChunk, 2)]);

        let retained = select_files_to_retain(target, &source);
        assert_eq!(retained.len(), 1);
        assert_eq!(retained[0].id, 1);
        assert_eq!(retained[0].ftype, FileType::Blob);
    }
}