// Copyright 2026 PingCAP, Inc.

use std::{
    collections::{HashMap, VecDeque},
    panic::AssertUnwindSafe,
    sync::{
        Arc, Condvar, Mutex,
        atomic::{AtomicBool, AtomicU64, Ordering},
    },
    thread::{self, JoinHandle},
    time::Duration,
};

use etcd_client::{
    Client as EtcdClient, EventType, GetOptions, WatchOptions, WatchResponse, WatchStream, Watcher,
};
use futures::FutureExt;
use kvengine::dfs::new_dfs_from_config;
use pd_client::PdClient;
use rfenginepb::ClusterBackupMeta;
use slog_global::{info, warn};
use tikv_util::{retry::try_wait_result_async, sys::thread::StdThreadBuildWrapper};
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};

use crate::{
    backup::{BackupConfig, Result},
    common::{generate_etcd_connect_opt, get_latest_backup_meta},
    error::{Error, PdMetaWatchCacheError},
    metrics::{
        NATIVE_BR_PD_META_BACKUP_SNAPSHOT_TOTAL, NATIVE_BR_PD_META_CACHE_KVS,
        NATIVE_BR_PD_META_CACHE_REVISION, NATIVE_BR_PD_META_CACHE_STATE,
        NATIVE_BR_PD_META_FALLBACK_TOTAL, NATIVE_BR_PD_META_WATCH_EVENTS_TOTAL,
        NATIVE_BR_PD_META_WATCH_KEYS_TOTAL, NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL,
    },
};

const MAX_BATCH_GET_CNT: i64 = 1024;
const RETRY_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const RETRY_INTERVAL: Duration = Duration::from_secs(1);

// `PD_KEY_SPACE_META_PATH` is an array of tuples each containing:
// 1. A string representing a keyspace meta path.
// 2. A boolean indicating whether the path is prefixed by 'pd/$cluster_id/'.
//
// Keyspace meta paths in PD are:
// 1. keyspace meta: "/pd/$cluster_id/keyspaces/"
// 2. keyspace region label: "/pd/$cluster_id/region_label/keyspaces/"
// 3. keyspace placement rules: "/pd/$cluster_id/rules/"
// 4. rule group: "/pd/$cluster_id/rule_group/"
// 5. keyspace group membership:
//    "/pd/$cluster_id/tso/keyspace_groups/membership/"
// 6. resource group information: "resource_group/" (Note: no leading slash)
// 7. tidb worker keys: "/tidb/remote/worker/"
// 8. pd config: "/pd/$cluster_id/config"
// 9. meta-service groups statuses "/pd/$cluster_id/meta_service_groups/"
const PD_KEY_SPACE_META_PATH: [(&str, bool); 9] = [
    ("keyspaces/", true),
    ("region_label/keyspaces/", true),
    ("rules/", true),
    ("rule_group/", true),
    ("tso/keyspace_groups/membership/", true),
    ("resource_group/", false),
    ("/tidb/remote/worker/", false),
    ("config", true),
    ("meta_service_groups/", true),
];

const CACHE_STATES: [&str; 5] = ["initializing", "ready", "stale", "resyncing", "stopped"];

/// A complete raw PD metadata image at one etcd revision.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PdMetaSnapshot {
    pub revision: i64,
    pub kvs: HashMap<Vec<u8>, Vec<u8>>,
}

/// Supplies a complete PD metadata snapshot for the backup path.
pub trait PdMetaProvider: Send + Sync {
    fn snapshot(&self, timeout: Duration) -> Result<PdMetaSnapshot>;
}

/// Long-lived watch cache for raw PD metadata used by native BR backups.
pub struct PdMetaWatchCache {
    inner: Arc<CacheInner>,
    command_tx: UnboundedSender<WatchCommand>,
    join_handle: Mutex<Option<JoinHandle<()>>>,
}

impl PdMetaWatchCache {
    pub fn start(config: BackupConfig, pd_client: Arc<dyn PdClient>) -> Result<Arc<Self>> {
        let (command_tx, command_rx) = unbounded_channel();
        let inner = Arc::new(CacheInner::new());
        let worker_inner = inner.clone();
        let join_handle = thread::Builder::new()
            .name("pd-meta-watch-cache".to_owned())
            .spawn_wrapper(move || run_cache_worker(config, pd_client, worker_inner, command_rx))
            .map_err(|err| {
                Error::BackupError(format!("failed to start pd meta watch cache: {err:?}"))
            })?;
        Ok(Arc::new(Self {
            inner,
            command_tx,
            join_handle: Mutex::new(Some(join_handle)),
        }))
    }
}

impl PdMetaProvider for PdMetaWatchCache {
    fn snapshot(&self, timeout: Duration) -> Result<PdMetaSnapshot> {
        let request_id = self.inner.next_progress_id.fetch_add(1, Ordering::SeqCst);
        {
            let state = self.inner.data.lock().unwrap();
            if state.status != CacheStatus::Ready {
                return Err(cache_error(state.status));
            }
        }

        self.command_tx
            .send(WatchCommand::RequestProgress(request_id))
            .map_err(|_| PdMetaWatchCacheError::Stopped)?;

        let mut state = self.inner.data.lock().unwrap();
        let wait_result = self
            .inner
            .condvar
            .wait_timeout_while(state, timeout, |state| {
                state.status == CacheStatus::Ready && state.progress_ack_id < request_id
            })
            .unwrap();
        state = wait_result.0;

        if state.status != CacheStatus::Ready {
            return Err(cache_error(state.status));
        }
        if state.progress_ack_id < request_id {
            return Err(PdMetaWatchCacheError::ProgressTimeout.into());
        }

        Ok(PdMetaSnapshot {
            revision: state.applied_revision,
            kvs: state.kvs.clone(),
        })
    }
}

impl Drop for PdMetaWatchCache {
    fn drop(&mut self) {
        self.inner.stopped.store(true, Ordering::SeqCst);
        let _ = self.command_tx.send(WatchCommand::Stop);
        if let Some(join_handle) = self.join_handle.lock().unwrap().take() {
            if join_handle.thread().id() != thread::current().id() {
                let _ = join_handle.join();
            }
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum CacheStatus {
    Initializing,
    Ready,
    Stale,
    Resyncing,
    Stopped,
}

impl CacheStatus {
    fn label(self) -> &'static str {
        match self {
            CacheStatus::Initializing => "initializing",
            CacheStatus::Ready => "ready",
            CacheStatus::Stale => "stale",
            CacheStatus::Resyncing => "resyncing",
            CacheStatus::Stopped => "stopped",
        }
    }
}

struct CacheData {
    status: CacheStatus,
    kvs: HashMap<Vec<u8>, Vec<u8>>,
    /// Highest revision whose watched events were either applied or proven
    /// empty by progress.
    applied_revision: i64,
    progress_ack_id: u64,
}

struct CacheInner {
    data: Mutex<CacheData>,
    condvar: Condvar,
    stopped: AtomicBool,
    next_progress_id: AtomicU64,
}

impl CacheInner {
    fn new() -> Self {
        set_state_metric(CacheStatus::Initializing);
        Self {
            data: Mutex::new(CacheData {
                status: CacheStatus::Initializing,
                kvs: HashMap::new(),
                applied_revision: 0,
                progress_ack_id: 0,
            }),
            condvar: Condvar::new(),
            stopped: AtomicBool::new(false),
            next_progress_id: AtomicU64::new(1),
        }
    }

    fn set_status(&self, status: CacheStatus) {
        {
            let mut data = self.data.lock().unwrap();
            data.status = status;
        }
        set_state_metric(status);
        self.condvar.notify_all();
    }

    fn replace_snapshot(&self, snapshot: PdMetaSnapshot) {
        let kv_count = snapshot.kvs.len();
        let revision = snapshot.revision;
        {
            let mut data = self.data.lock().unwrap();
            data.kvs = snapshot.kvs;
            data.applied_revision = revision;
            data.status = CacheStatus::Ready;
        }
        set_state_metric(CacheStatus::Ready);
        NATIVE_BR_PD_META_CACHE_KVS.set(kv_count as i64);
        NATIVE_BR_PD_META_CACHE_REVISION.set(revision);
        self.condvar.notify_all();
    }

    fn ack_progress(&self, request_id: u64) {
        {
            let mut data = self.data.lock().unwrap();
            data.progress_ack_id = data.progress_ack_id.max(request_id);
        }
        self.condvar.notify_all();
    }

    fn apply_watch_response(&self, prefixes: &[Vec<u8>], resp: &WatchResponse) {
        let mut data = self.data.lock().unwrap();
        let mut event_revision: Option<i64> = None;
        for event in resp.events() {
            let Some(kv) = event.kv() else {
                continue;
            };
            let mod_revision = kv.mod_revision();
            event_revision =
                Some(event_revision.map_or(mod_revision, |revision| revision.max(mod_revision)));
            match event.event_type() {
                EventType::Put => {
                    apply_raw_event(
                        prefixes,
                        &mut data.kvs,
                        RawPdMetaEvent::Put {
                            key: kv.key(),
                            value: kv.value(),
                        },
                    );
                }
                EventType::Delete => {
                    apply_raw_event(
                        prefixes,
                        &mut data.kvs,
                        RawPdMetaEvent::Delete { key: kv.key() },
                    );
                }
            }
        }
        // Event response headers can be ahead of the last delivered event, so
        // only empty progress responses can safely advance to the header.
        let safe_revision = if resp.events().is_empty() {
            resp.header().map(|header| header.revision())
        } else {
            event_revision
        };
        if let Some(revision) = safe_revision {
            data.applied_revision = data.applied_revision.max(revision);
        }
        NATIVE_BR_PD_META_CACHE_KVS.set(data.kvs.len() as i64);
        NATIVE_BR_PD_META_CACHE_REVISION.set(data.applied_revision);
        drop(data);
        self.condvar.notify_all();
    }

    fn local_seed(&self) -> Option<PdMetaSnapshot> {
        let data = self.data.lock().unwrap();
        (data.applied_revision > 0).then(|| PdMetaSnapshot {
            revision: data.applied_revision,
            kvs: data.kvs.clone(),
        })
    }

    fn is_stopped(&self) -> bool {
        self.stopped.load(Ordering::SeqCst)
    }
}

enum WatchCommand {
    RequestProgress(u64),
    Stop,
}

#[derive(Clone)]
enum BootstrapPlan {
    LatestBackup,
    FullSnapshot,
    LocalSeed(PdMetaSnapshot),
    IncrementalScan(PdMetaSnapshot),
}

struct BootstrapSnapshot {
    source: &'static str,
    snapshot: PdMetaSnapshot,
}

enum WatchLoopExit {
    Reconnect(&'static str),
    IncrementalScan(&'static str),
    Stopped,
}

enum WatchStartError {
    Retry(&'static str, Error),
    FullResync(&'static str),
}

fn run_cache_worker(
    config: BackupConfig,
    pd_client: Arc<dyn PdClient>,
    inner: Arc<CacheInner>,
    command_rx: UnboundedReceiver<WatchCommand>,
) {
    let runtime = tokio::runtime::Builder::new_current_thread()
        .thread_name("pd-meta-watch-cache")
        .enable_all()
        .build()
        .unwrap();
    runtime.block_on(run_cache_worker_async(config, pd_client, inner, command_rx));
}

async fn run_cache_worker_async(
    config: BackupConfig,
    pd_client: Arc<dyn PdClient>,
    inner: Arc<CacheInner>,
    mut command_rx: UnboundedReceiver<WatchCommand>,
) {
    let mut initialized = false;
    let mut plan = BootstrapPlan::LatestBackup;

    while !inner.is_stopped() {
        inner.set_status(if initialized {
            CacheStatus::Resyncing
        } else {
            CacheStatus::Initializing
        });

        let cluster_id = match pd_client.get_cluster_id() {
            Ok(cluster_id) => cluster_id,
            Err(err) => {
                warn!("pd meta watch cache: get cluster id failed"; "err" => ?err);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&["cluster_id_error"])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                if wait_retry_or_stop(&mut command_rx).await {
                    break;
                }
                continue;
            }
        };

        let prefixes = pd_keyspace_meta_prefixes(cluster_id);
        let mut etcd_client = match connect_etcd_with_retry(&config).await {
            Ok(client) => client,
            Err(err) => {
                warn!("pd meta watch cache: connect etcd failed"; "err" => ?err);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&["connect_error"])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                if wait_retry_or_stop(&mut command_rx).await {
                    break;
                }
                continue;
            }
        };

        let bootstrap = match bootstrap_snapshot(
            &config,
            cluster_id,
            &prefixes,
            &mut etcd_client,
            plan.clone(),
        )
        .await
        {
            Ok(bootstrap) => bootstrap,
            Err(err) => {
                warn!("pd meta watch cache: bootstrap failed"; "err" => ?err);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&["bootstrap_error"])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                plan = BootstrapPlan::FullSnapshot;
                if wait_retry_or_stop(&mut command_rx).await {
                    break;
                }
                continue;
            }
        };

        let snapshot = bootstrap.snapshot;
        let start_revision = snapshot.revision + 1;
        let (watcher, stream) = match start_watch(&mut etcd_client, start_revision).await {
            Ok(watch) => watch,
            Err(WatchStartError::FullResync(reason)) => {
                warn!("pd meta watch cache: watch start needs full resync"; "reason" => reason);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&[reason])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                plan = BootstrapPlan::FullSnapshot;
                continue;
            }
            Err(WatchStartError::Retry(reason, err)) => {
                warn!("pd meta watch cache: watch start failed"; "reason" => reason, "err" => ?err);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&[reason])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                if wait_retry_or_stop(&mut command_rx).await {
                    break;
                }
                continue;
            }
        };

        info!("pd meta watch cache: bootstrap finished";
            "source" => bootstrap.source,
            "revision" => snapshot.revision,
            "kvs" => snapshot.kvs.len());
        inner.replace_snapshot(snapshot);
        initialized = true;

        match watch_loop(inner.clone(), &prefixes, watcher, stream, &mut command_rx).await {
            WatchLoopExit::Stopped => break,
            WatchLoopExit::Reconnect(reason) => {
                warn!("pd meta watch cache: reconnect watch"; "reason" => reason);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&[reason])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                plan = inner
                    .local_seed()
                    .map(BootstrapPlan::LocalSeed)
                    .unwrap_or(BootstrapPlan::FullSnapshot);
            }
            WatchLoopExit::IncrementalScan(reason) => {
                warn!("pd meta watch cache: incremental scan resync"; "reason" => reason);
                NATIVE_BR_PD_META_WATCH_RESYNC_TOTAL
                    .with_label_values(&[reason])
                    .inc();
                inner.set_status(CacheStatus::Stale);
                plan = inner
                    .local_seed()
                    .map(BootstrapPlan::IncrementalScan)
                    .unwrap_or(BootstrapPlan::FullSnapshot);
            }
        }
    }

    inner.set_status(CacheStatus::Stopped);
}

async fn wait_retry_or_stop(command_rx: &mut UnboundedReceiver<WatchCommand>) -> bool {
    loop {
        tokio::select! {
            _ = tokio::time::sleep(RETRY_INTERVAL) => return false,
            command = command_rx.recv() => {
                match command {
                    Some(WatchCommand::Stop) | None => return true,
                    Some(WatchCommand::RequestProgress(_)) => {}
                }
            }
        }
    }
}

async fn bootstrap_snapshot(
    config: &BackupConfig,
    cluster_id: u64,
    prefixes: &[Vec<u8>],
    etcd_client: &mut EtcdClient,
    plan: BootstrapPlan,
) -> Result<BootstrapSnapshot> {
    match plan {
        BootstrapPlan::LatestBackup => {
            if let Some(seed) = latest_backup_seed(config, cluster_id).await {
                return Ok(BootstrapSnapshot {
                    source: "latest_backup",
                    snapshot: seed,
                });
            }
            full_snapshot(etcd_client, prefixes).await
        }
        BootstrapPlan::FullSnapshot => full_snapshot(etcd_client, prefixes).await,
        BootstrapPlan::LocalSeed(snapshot) => Ok(BootstrapSnapshot {
            source: "local_seed",
            snapshot,
        }),
        BootstrapPlan::IncrementalScan(snapshot) => {
            incremental_snapshot(etcd_client, prefixes, snapshot).await
        }
    }
}

async fn latest_backup_seed(config: &BackupConfig, cluster_id: u64) -> Option<PdMetaSnapshot> {
    let dfs = new_dfs_from_config(config.dfs.clone());
    match get_latest_backup_meta(dfs.as_ref(), cluster_id).await {
        Ok(meta) => backup_meta_seed(meta).inspect(|seed| {
            info!("pd meta watch cache: seed from latest backup";
                "revision" => seed.revision,
                "kvs" => seed.kvs.len());
        }),
        Err(err) => {
            warn!("pd meta watch cache: latest backup seed unavailable"; "err" => ?err);
            None
        }
    }
}

fn backup_meta_seed(mut meta: ClusterBackupMeta) -> Option<PdMetaSnapshot> {
    if meta.meta_revision == 0 || meta.keyspace_meta.is_empty() {
        return None;
    }
    Some(PdMetaSnapshot {
        revision: meta.meta_revision,
        kvs: meta.take_keyspace_meta(),
    })
}

async fn full_snapshot(
    etcd_client: &mut EtcdClient,
    prefixes: &[Vec<u8>],
) -> Result<BootstrapSnapshot> {
    let revision = latest_etcd_revision(etcd_client, prefixes).await?;
    let mut kvs = HashMap::new();
    for prefix in prefixes {
        read_prefix_at_revision(etcd_client, prefix, revision, &mut kvs).await?;
    }
    Ok(BootstrapSnapshot {
        source: "full_snapshot",
        snapshot: PdMetaSnapshot { revision, kvs },
    })
}

async fn incremental_snapshot(
    etcd_client: &mut EtcdClient,
    prefixes: &[Vec<u8>],
    seed: PdMetaSnapshot,
) -> Result<BootstrapSnapshot> {
    let old_revision = seed.revision;
    let revision = latest_etcd_revision(etcd_client, prefixes).await?;
    let mut kvs = seed.kvs;
    for prefix in prefixes {
        read_prefix_changes_at_revision(etcd_client, prefix, old_revision, revision, &mut kvs)
            .await?;
    }
    Ok(BootstrapSnapshot {
        source: "incremental_scan",
        snapshot: PdMetaSnapshot { revision, kvs },
    })
}

async fn latest_etcd_revision(etcd_client: &mut EtcdClient, prefixes: &[Vec<u8>]) -> Result<i64> {
    let prefix = prefixes
        .first()
        .ok_or_else(|| Error::BackupError("pd meta prefixes are empty".to_owned()))?;
    let resp = etcd_client
        .get(
            prefix.clone(),
            Some(
                GetOptions::new()
                    .with_range(get_prefix_range_end(prefix))
                    .with_limit(1),
            ),
        )
        .await?;
    response_revision(resp.header())
}

async fn read_prefix_at_revision(
    etcd_client: &mut EtcdClient,
    prefix: &[u8],
    revision: i64,
    kvs: &mut HashMap<Vec<u8>, Vec<u8>>,
) -> Result<()> {
    let get_option = GetOptions::new()
        .with_range(get_prefix_range_end(prefix))
        .with_revision(revision)
        .with_limit(MAX_BATCH_GET_CNT);
    let mut seek_key = prefix.to_vec();
    loop {
        let resp = etcd_client
            .get(seek_key.clone(), Some(get_option.clone()))
            .await?;
        let mut more = resp.more();
        for kv in resp.kvs() {
            if !kv.key().starts_with(prefix) {
                more = false;
                break;
            }
            kvs.insert(kv.key().to_vec(), kv.value().to_vec());
        }
        if !more {
            break;
        }
        let Some(last_key) = resp.kvs().last().map(|kv| kv.key()) else {
            break;
        };
        seek_key = last_key.to_vec();
        seek_key.push(0);
    }
    Ok(())
}

async fn read_prefix_changes_at_revision(
    etcd_client: &mut EtcdClient,
    prefix: &[u8],
    min_mod_revision: i64,
    revision: i64,
    kvs: &mut HashMap<Vec<u8>, Vec<u8>>,
) -> Result<()> {
    let get_option = GetOptions::new()
        .with_range(get_prefix_range_end(prefix))
        .with_revision(revision)
        .with_min_mod_revision(min_mod_revision)
        .with_limit(MAX_BATCH_GET_CNT);
    let mut seek_key = prefix.to_vec();
    loop {
        let resp = etcd_client
            .get(seek_key.clone(), Some(get_option.clone()))
            .await?;
        let mut more = resp.more();
        for kv in resp.kvs() {
            if !kv.key().starts_with(prefix) {
                more = false;
                break;
            }
            kvs.insert(kv.key().to_vec(), kv.value().to_vec());
        }
        if !more {
            break;
        }
        let Some(last_key) = resp.kvs().last().map(|kv| kv.key()) else {
            break;
        };
        seek_key = last_key.to_vec();
        seek_key.push(0);
    }
    Ok(())
}

async fn start_watch(
    etcd_client: &mut EtcdClient,
    start_revision: i64,
) -> std::result::Result<(Watcher, WatchStream), WatchStartError> {
    let first_watch =
        AssertUnwindSafe(etcd_client.watch(Vec::new(), Some(watch_options(start_revision))))
            .catch_unwind()
            .await;
    match first_watch {
        Ok(Ok(watch)) => Ok(watch),
        Ok(Err(err)) => Err(WatchStartError::Retry("watch_create_error", err.into())),
        Err(_) => Err(WatchStartError::FullResync("watch_create_panic")),
    }
}

fn watch_options(start_revision: i64) -> WatchOptions {
    WatchOptions::new()
        .with_all_keys()
        .with_start_revision(start_revision)
}

async fn watch_loop(
    inner: Arc<CacheInner>,
    prefixes: &[Vec<u8>],
    mut watcher: Watcher,
    mut stream: WatchStream,
    command_rx: &mut UnboundedReceiver<WatchCommand>,
) -> WatchLoopExit {
    let mut pending_progress = VecDeque::new();
    loop {
        tokio::select! {
            command = command_rx.recv() => {
                match command {
                    Some(WatchCommand::RequestProgress(request_id)) => {
                        if let Err(err) = watcher.request_progress().await {
                            warn!("pd meta watch cache: request progress failed"; "err" => ?err);
                            return WatchLoopExit::Reconnect("progress_request_error");
                        }
                        pending_progress.push_back(request_id);
                    }
                    Some(WatchCommand::Stop) | None => return WatchLoopExit::Stopped,
                }
            }
            resp = stream.message() => {
                let resp = match resp {
                    Ok(Some(resp)) => resp,
                    Ok(None) => return WatchLoopExit::Reconnect("stream_eof"),
                    Err(err) => {
                        warn!("pd meta watch cache: watch stream error"; "err" => ?err);
                        return WatchLoopExit::Reconnect("stream_error");
                    }
                };
                if resp.canceled() {
                    warn!("pd meta watch cache: watch canceled";
                        "watch_id" => resp.watch_id(),
                        "compact_revision" => resp.compact_revision(),
                        "reason" => resp.cancel_reason());
                    return if resp.compact_revision() > 0 {
                        WatchLoopExit::IncrementalScan("compacted")
                    } else {
                        WatchLoopExit::Reconnect("canceled")
                    };
                }
                if resp.created() {
                    continue;
                }

                let is_progress = resp.events().is_empty();
                inner.apply_watch_response(prefixes, &resp);
                if is_progress {
                    if let Some(request_id) = pending_progress.pop_front() {
                        inner.ack_progress(request_id);
                    }
                }
            }
        }
    }
}

enum RawPdMetaEvent<'a> {
    Put { key: &'a [u8], value: &'a [u8] },
    Delete { key: &'a [u8] },
}

fn apply_raw_event(
    prefixes: &[Vec<u8>],
    kvs: &mut HashMap<Vec<u8>, Vec<u8>>,
    event: RawPdMetaEvent<'_>,
) -> bool {
    let key = match event {
        RawPdMetaEvent::Put { key, .. } | RawPdMetaEvent::Delete { key } => key,
    };
    if !prefixes.iter().any(|prefix| key.starts_with(prefix)) {
        NATIVE_BR_PD_META_WATCH_KEYS_TOTAL
            .with_label_values(&["ignored"])
            .inc();
        NATIVE_BR_PD_META_WATCH_EVENTS_TOTAL
            .with_label_values(&["ignored"])
            .inc();
        return false;
    }

    NATIVE_BR_PD_META_WATCH_KEYS_TOTAL
        .with_label_values(&["watched"])
        .inc();
    match event {
        RawPdMetaEvent::Put { key, value } => {
            kvs.insert(key.to_vec(), value.to_vec());
            NATIVE_BR_PD_META_WATCH_EVENTS_TOTAL
                .with_label_values(&["put"])
                .inc();
        }
        RawPdMetaEvent::Delete { key } => {
            kvs.remove(key);
            NATIVE_BR_PD_META_WATCH_EVENTS_TOTAL
                .with_label_values(&["delete"])
                .inc();
        }
    }
    true
}

pub(crate) async fn connect_etcd(config: &BackupConfig) -> Result<EtcdClient> {
    let option = generate_etcd_connect_opt(&config.security)?;
    EtcdClient::connect(&config.pd.endpoints, Some(option))
        .await
        .map_err(|err| {
            warn!("connect etcd failed"; "err" => ?err);
            Error::EtcdError(err)
        })
}

async fn connect_etcd_with_retry(config: &BackupConfig) -> Result<EtcdClient> {
    let mut backoff = Duration::ZERO;
    try_wait_result_async(
        || Box::pin(connect_etcd(config)),
        RETRY_CONNECT_TIMEOUT,
        move || {
            backoff += Duration::from_millis(200);
            backoff
        },
    )
    .await
}

// Get keyspace meta from etcd and populate them to ClusterBackupMeta.
pub async fn backup_pd_keyspace_meta(
    config: &BackupConfig,
    cluster_backup_meta: &mut ClusterBackupMeta,
    pd_meta_provider: Option<&dyn PdMetaProvider>,
) -> Result<()> {
    #[cfg(feature = "testexport")]
    if config.skip_keyspace_meta {
        warn!("skip backup keyspace meta");
        return Ok(());
    }

    if let Some(provider) = pd_meta_provider {
        match provider.snapshot(config.backup_pd_meta_watch_cache_progress_timeout.0) {
            Ok(snapshot) => {
                let kv_count = snapshot.kvs.len();
                let revision = snapshot.revision;
                cluster_backup_meta.set_keyspace_meta(snapshot.kvs);
                cluster_backup_meta.set_meta_revision(revision);
                NATIVE_BR_PD_META_BACKUP_SNAPSHOT_TOTAL
                    .with_label_values(&["cache"])
                    .inc();
                info!("Backed up pd meta from watch cache";
                    "kvs" => kv_count,
                    "revision" => revision);
                return Ok(());
            }
            Err(err) => {
                let reason = match &err {
                    Error::PdMetaWatchCache(err) => err.fallback_label(),
                    _ => "provider_error",
                };
                NATIVE_BR_PD_META_FALLBACK_TOTAL
                    .with_label_values(&[reason])
                    .inc();
                warn!("pd meta watch cache unavailable, fallback to scan"; "err" => ?err);
            }
        }
    } else {
        NATIVE_BR_PD_META_FALLBACK_TOTAL
            .with_label_values(&["provider_disabled"])
            .inc();
    }

    NATIVE_BR_PD_META_BACKUP_SNAPSHOT_TOTAL
        .with_label_values(&["scan"])
        .inc();
    scan_pd_keyspace_meta(config, cluster_backup_meta).await
}

async fn scan_pd_keyspace_meta(
    config: &BackupConfig,
    cluster_backup_meta: &mut ClusterBackupMeta,
) -> Result<()> {
    info!("start backup PD keyspace meta"; "backup_meta" => %cluster_backup_meta);
    let cluster_id = cluster_backup_meta.cluster_id;

    let mut etcd_client = connect_etcd_with_retry(config).await?;

    let old_revision = cluster_backup_meta.meta_revision;
    let mut min_revision = i64::MAX;
    let mut new_meta_cnt = 0;
    for prefix in pd_keyspace_meta_prefixes(cluster_id) {
        let get_option = GetOptions::new()
            .with_range(get_prefix_range_end(&prefix))
            .with_min_mod_revision(old_revision)
            .with_limit(MAX_BATCH_GET_CNT);
        let mut seek_key = prefix.clone();
        loop {
            let resp = etcd_client.get(seek_key, Some(get_option.clone())).await?;
            let mut more = resp.more();
            for kv in resp.kvs() {
                if !kv.key().starts_with(&prefix) {
                    more = false;
                    break;
                }
                new_meta_cnt += 1;
                cluster_backup_meta
                    .mut_keyspace_meta()
                    .insert(kv.key().to_vec(), kv.value().to_vec());
            }
            min_revision = std::cmp::min(min_revision, response_revision(resp.header())?);
            if !more {
                break;
            }
            seek_key = resp.kvs().last().unwrap().key().to_owned();
            seek_key.push(0);
        }
    }
    if min_revision != i64::MAX {
        cluster_backup_meta.set_meta_revision(min_revision);
    }
    info!(
        "Backed up {} pd meta kvs, including {} new meta, revision: {} -> {}",
        cluster_backup_meta.keyspace_meta.len(),
        new_meta_cnt,
        old_revision,
        cluster_backup_meta.meta_revision
    );
    Ok(())
}

fn response_revision(header: Option<&etcd_client::ResponseHeader>) -> Result<i64> {
    header
        .map(|header| header.revision())
        .ok_or_else(|| Error::BackupError("missing etcd response header".to_owned()))
}

fn cache_error(status: CacheStatus) -> Error {
    PdMetaWatchCacheError::State(status.label()).into()
}

fn set_state_metric(status: CacheStatus) {
    for state in CACHE_STATES {
        NATIVE_BR_PD_META_CACHE_STATE
            .with_label_values(&[state])
            .set(if state == status.label() { 1 } else { 0 });
    }
}

pub(crate) fn pd_keyspace_meta_prefixes(cluster_id: u64) -> Vec<Vec<u8>> {
    PD_KEY_SPACE_META_PATH
        .into_iter()
        .map(|(path, prefixed)| {
            if prefixed {
                format!("/pd/{cluster_id}/{path}").into_bytes()
            } else {
                path.as_bytes().to_vec()
            }
        })
        .collect()
}

/// Returns the PD metadata prefixes used by the watch cache for integration
/// tests that verify etcd contents directly.
#[cfg(feature = "testexport")]
pub fn pd_keyspace_meta_prefixes_for_test(cluster_id: u64) -> Vec<Vec<u8>> {
    pd_keyspace_meta_prefixes(cluster_id)
}

// Returns the exclusive upper bound for an etcd prefix range scan, matching
// Go's clientv3.GetPrefixRangeEnd:
// https://github.com/etcd-io/etcd/blob/e7b3bb6ccac840770f108ef9a0f013fa51b83256/client/v3/op.go#L366-L384
//
// Copyright 2016 The etcd Authors
// Modifications Copyright 2026 PingCAP, Inc.
//
// Increments the last non-0xFF byte and truncates the suffix. If every byte is
// 0xFF (no valid successor exists), returns `[0x00]`, which etcd interprets as
// "no upper bound" — consistent with the Go implementation.
pub(crate) fn get_prefix_range_end(prefix: &[u8]) -> Vec<u8> {
    for (i, &b) in prefix.iter().enumerate().rev() {
        if b < 0xFF {
            let mut end = prefix[..=i].to_vec();
            end[i] += 1;
            return end;
        }
    }
    vec![0x00]
}

/// Returns the exclusive etcd range end for a prefix scan in integration tests.
#[cfg(feature = "testexport")]
pub fn get_prefix_range_end_for_test(prefix: &[u8]) -> Vec<u8> {
    get_prefix_range_end(prefix)
}

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

    use etcd_client::proto::{PbEvent, PbKeyValue, PbResponseHeader, PbWatchResponse};

    use super::*;

    #[test]
    fn test_get_prefix_range_end() {
        assert_eq!(get_prefix_range_end(b"foo"), b"fop");
        assert_eq!(
            get_prefix_range_end(b"/pd/123/keyspaces/"),
            b"/pd/123/keyspaces0"
        );
        assert_eq!(get_prefix_range_end(b"foo\xFF"), b"fop");
        assert_eq!(get_prefix_range_end(b"\xFF"), b"\x00");
        assert_eq!(get_prefix_range_end(b"\xFF\xFF\xFF"), b"\x00");
    }

    #[test]
    fn test_pd_keyspace_meta_prefixes() {
        let prefixes = pd_keyspace_meta_prefixes(42);
        assert!(prefixes.contains(&b"/pd/42/keyspaces/".to_vec()));
        assert!(prefixes.contains(&b"resource_group/".to_vec()));
        assert!(prefixes.contains(&b"/tidb/remote/worker/".to_vec()));
        assert!(prefixes.contains(&b"/pd/42/config".to_vec()));
    }

    #[test]
    fn test_apply_put_and_delete_events_update_map() {
        let prefixes = vec![b"/pd/42/keyspaces/".to_vec()];
        let mut kvs = HashMap::new();
        assert!(apply_raw_event(
            &prefixes,
            &mut kvs,
            RawPdMetaEvent::Put {
                key: b"/pd/42/keyspaces/100",
                value: b"v1",
            },
        ));
        assert_eq!(
            kvs.get(b"/pd/42/keyspaces/100".as_slice()),
            Some(&b"v1".to_vec())
        );

        assert!(apply_raw_event(
            &prefixes,
            &mut kvs,
            RawPdMetaEvent::Delete {
                key: b"/pd/42/keyspaces/100",
            },
        ));
        assert!(!kvs.contains_key(b"/pd/42/keyspaces/100".as_slice()));
    }

    #[test]
    fn test_apply_event_ignores_keys_outside_scope() {
        let prefixes = vec![b"/pd/42/keyspaces/".to_vec()];
        let mut kvs = HashMap::new();
        assert!(!apply_raw_event(
            &prefixes,
            &mut kvs,
            RawPdMetaEvent::Put {
                key: b"/pd/43/keyspaces/100",
                value: b"v1",
            },
        ));
        assert!(kvs.is_empty());
    }

    #[test]
    fn test_apply_event_records_watched_and_ignored_keys() {
        let prefixes = vec![b"/pd/42/keyspaces/".to_vec()];
        let mut kvs = HashMap::new();
        let watched_before = NATIVE_BR_PD_META_WATCH_KEYS_TOTAL
            .with_label_values(&["watched"])
            .get();
        let ignored_before = NATIVE_BR_PD_META_WATCH_KEYS_TOTAL
            .with_label_values(&["ignored"])
            .get();

        assert!(apply_raw_event(
            &prefixes,
            &mut kvs,
            RawPdMetaEvent::Put {
                key: b"/pd/42/keyspaces/100",
                value: b"v1",
            },
        ));
        assert!(!apply_raw_event(
            &prefixes,
            &mut kvs,
            RawPdMetaEvent::Put {
                key: b"/pd/43/keyspaces/100",
                value: b"v1",
            },
        ));

        assert!(
            NATIVE_BR_PD_META_WATCH_KEYS_TOTAL
                .with_label_values(&["watched"])
                .get()
                > watched_before
        );
        assert!(
            NATIVE_BR_PD_META_WATCH_KEYS_TOTAL
                .with_label_values(&["ignored"])
                .get()
                > ignored_before
        );
    }

    #[test]
    fn test_replace_snapshot_is_atomic_for_callers() {
        let inner = CacheInner::new();
        let mut first = HashMap::new();
        first.insert(b"k1".to_vec(), b"v1".to_vec());
        inner.replace_snapshot(PdMetaSnapshot {
            revision: 10,
            kvs: first.clone(),
        });
        assert_eq!(inner.local_seed().unwrap().kvs, first);

        let mut second = HashMap::new();
        second.insert(b"k2".to_vec(), b"v2".to_vec());
        inner.replace_snapshot(PdMetaSnapshot {
            revision: 20,
            kvs: second.clone(),
        });
        let seed = inner.local_seed().unwrap();
        assert_eq!(seed.revision, 20);
        assert_eq!(seed.kvs, second);
    }

    #[test]
    fn test_apply_watch_response_resumes_from_last_event_revision() {
        let prefixes = vec![b"/pd/42/keyspaces/".to_vec()];
        let inner = CacheInner::new();
        inner.replace_snapshot(PdMetaSnapshot {
            revision: 10,
            kvs: HashMap::new(),
        });

        let resp = watch_response(100, vec![put_event(b"/pd/42/keyspaces/100", b"v1", 40)]);

        inner.apply_watch_response(&prefixes, &resp);
        let seed = inner.local_seed().unwrap();
        assert_eq!(seed.revision, 40);
        assert_eq!(seed.revision + 1, 41);
        assert_eq!(
            seed.kvs.get(b"/pd/42/keyspaces/100".as_slice()),
            Some(&b"v1".to_vec())
        );
    }

    #[test]
    fn test_apply_watch_response_progress_uses_header_revision() {
        let prefixes = vec![b"/pd/42/keyspaces/".to_vec()];
        let inner = CacheInner::new();
        inner.replace_snapshot(PdMetaSnapshot {
            revision: 10,
            kvs: HashMap::new(),
        });

        let resp = watch_response(100, vec![]);

        inner.apply_watch_response(&prefixes, &resp);
        let seed = inner.local_seed().unwrap();
        assert_eq!(seed.revision, 100);
        assert_eq!(seed.revision + 1, 101);
    }

    #[test]
    fn test_backup_meta_seed_requires_revision_and_meta() {
        let mut meta = ClusterBackupMeta::new();
        meta.set_meta_revision(10);
        assert!(backup_meta_seed(meta.clone()).is_none());

        meta.mut_keyspace_meta()
            .insert(b"k".to_vec(), b"v".to_vec());
        let seed = backup_meta_seed(meta).unwrap();
        assert_eq!(seed.revision, 10);
        assert_eq!(seed.kvs.get(b"k".as_slice()), Some(&b"v".to_vec()));

        let mut zero_revision = ClusterBackupMeta::new();
        zero_revision
            .mut_keyspace_meta()
            .insert(b"k".to_vec(), b"v".to_vec());
        assert!(backup_meta_seed(zero_revision).is_none());
    }

    struct StaticProvider {
        calls: AtomicUsize,
    }

    impl PdMetaProvider for StaticProvider {
        fn snapshot(&self, _: Duration) -> Result<PdMetaSnapshot> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Ok(PdMetaSnapshot {
                revision: 12,
                kvs: HashMap::from([(b"k".to_vec(), b"v".to_vec())]),
            })
        }
    }

    #[test]
    fn test_backup_pd_keyspace_meta_uses_provider_snapshot() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let provider = StaticProvider {
            calls: AtomicUsize::new(0),
        };
        let mut meta = ClusterBackupMeta::new();
        meta.set_cluster_id(42);

        runtime
            .block_on(backup_pd_keyspace_meta(
                &BackupConfig::default(),
                &mut meta,
                Some(&provider),
            ))
            .unwrap();

        assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
        assert_eq!(meta.meta_revision, 12);
        assert_eq!(
            meta.keyspace_meta.get(b"k".as_slice()),
            Some(&b"v".to_vec())
        );
    }

    fn watch_response(header_revision: i64, events: Vec<PbEvent>) -> WatchResponse {
        WatchResponse(PbWatchResponse {
            header: Some(PbResponseHeader {
                revision: header_revision,
                ..Default::default()
            }),
            events,
            ..Default::default()
        })
    }

    fn put_event(key: &[u8], value: &[u8], mod_revision: i64) -> PbEvent {
        PbEvent {
            r#type: EventType::Put as i32,
            kv: Some(PbKeyValue {
                key: key.to_vec(),
                value: value.to_vec(),
                mod_revision,
                ..Default::default()
            }),
            ..Default::default()
        }
    }
}