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

mod metrics;

use std::{
    collections::{BTreeMap, HashMap},
    error::Error as StdError,
    fmt,
    sync::{
        Arc, Mutex,
        atomic::{AtomicU64, Ordering},
    },
    time::{Duration, Instant},
};

use api_version::ApiV2;
use cloud_encryption::KeyspaceEncryptionConfig;
use hyper::{Body, Request, StatusCode, Uri};
use kvengine::dfs;
use kvenginepb::EncryptionMeta;
use kvproto::metapb::StoreState;
use log::{error, info};
use pd_client::{Error as PdError, PdClient, RegionInfo};
use protobuf::Message;
use security::{HttpClient, SecurityManager};
use serde::{Deserialize, Serialize};
use tokio::sync::watch;
use txn_types::Key;

use crate::metrics::{
    ENCRYPTION_CONF_CHANGE_FAILED_JOBS_COUNTER, ENCRYPTION_CONF_CHANGE_RUNNING_JOBS_GAUGE,
    ENCRYPTION_CONF_CHANGE_STUCK_JOBS_GAUGE,
};

pub const DEFAULT_SYNC_LATEST_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_REMOTE_START_ROLLOUT_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_REMOTE_GEN_META_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_SYNC_REGION_RETRY_LIMIT: usize = 16;
const DEFAULT_SYNC_REGION_RETRY_DELAY: Duration = Duration::from_millis(500);
const DEFAULT_STUCK_JOB_THRESHOLD: Duration = Duration::from_secs(24 * 60 * 60);

const KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID: &str = "encryption_meta_current_file_id";
const KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID: &str = "encryption_meta_next_file_id";
const KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID: &str =
    "encryption_meta_job_owner_store_id";
const KEYSPACE_CONFIG_ENCRYPTION: &str = "encryption";

/// Controls whether this store should enforce the "rollout job owner" recorded
/// in PD keyspace config.
///
/// `Disabled` is typically used by tools/clients: they may start the job on the
/// owner store via HTTP, and don't need to prove they are the owner.
///
/// `Enforced` is used on TiKV status servers: it rejects starting/canceling
/// jobs on a non-owner store to avoid duplicate work and conflicting updates.
#[derive(Clone, Debug)]
pub enum OwnerCheckMode {
    Disabled,
    Enforced { owner_store_id: u64 },
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StartRolloutJobRequest {
    pub keyspace_id: u32,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CancelRolloutJobRequest {
    pub keyspace_id: u32,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StartRolloutJobResponse {
    pub keyspace_id: u32,
    pub started: bool,
    pub running: bool,
    pub current_file_id: u64,
    pub next_file_id: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub note: Option<String>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CancelRolloutJobResponse {
    pub keyspace_id: u32,
    pub canceled: bool,
}

#[derive(Clone, Debug)]
pub struct GenMetaRequest {
    pub keyspace_id: u32,
    pub mode: EncryptionMetaGenMode,
}

#[derive(Clone, Debug)]
pub struct GenMetaResult {
    pub keyspace_id: u32,
    pub pd_before: PdEncryptionMetaState,
    pub pd_after: PdEncryptionMetaState,
    pub owner_store_id: Option<u64>,
    pub gen_meta: EncryptionMetaGenResponse,
    pub pd_updated: bool,
}

#[derive(Debug)]
struct KeyspaceFirstRegion {
    region_id: u64,
    store_status_addr: String,
    // Leader store ID for the first region.
    store_id: u64,
}

#[derive(Clone, Copy, Debug, Default, Serialize)]
pub struct RolloutJobSummary {
    pub regions_success: u64,
    pub regions_failed: u64,
}

#[derive(Debug, Clone)]
pub enum EncryptionConfChangeError {
    Canceled,
    KeyspaceNotFound(u32),
    WrongOwner {
        expected_owner_store_id: u64,
        pd_owner_store_id: u64,
    },
    Failed(String),
}

impl fmt::Display for EncryptionConfChangeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            EncryptionConfChangeError::Canceled => write!(f, "canceled"),
            EncryptionConfChangeError::KeyspaceNotFound(keyspace_id) => {
                write!(f, "keyspace not found: {}", keyspace_id)
            }
            EncryptionConfChangeError::WrongOwner {
                expected_owner_store_id,
                pd_owner_store_id,
            } => write!(
                f,
                "wrong owner: expected store_id {}, PD has store_id {}",
                expected_owner_store_id, pd_owner_store_id
            ),
            EncryptionConfChangeError::Failed(msg) => write!(f, "{}", msg),
        }
    }
}

impl StdError for EncryptionConfChangeError {}

impl EncryptionConfChangeError {
    fn failed(msg: impl Into<String>) -> Self {
        EncryptionConfChangeError::Failed(msg.into())
    }

    pub fn is_keyspace_not_found(&self) -> bool {
        matches!(self, Self::KeyspaceNotFound(_))
    }
}

#[derive(Debug)]
pub enum EncryptionSyncError {
    Canceled,
    Retryable(String),
    Failed(String),
}

impl fmt::Display for EncryptionSyncError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            EncryptionSyncError::Canceled => write!(f, "canceled"),
            EncryptionSyncError::Retryable(msg) => write!(f, "{}", msg),
            EncryptionSyncError::Failed(msg) => write!(f, "{}", msg),
        }
    }
}

impl StdError for EncryptionSyncError {}

impl EncryptionSyncError {
    fn retryable(msg: impl Into<String>) -> Self {
        EncryptionSyncError::Retryable(msg.into())
    }

    fn failed(msg: impl Into<String>) -> Self {
        EncryptionSyncError::Failed(msg.into())
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PdEncryptionMetaState {
    pub current_file_id: u64,
    pub next_file_id: u64,
}

#[derive(Clone, Debug)]
pub struct PdKeyspaceEncryptionStatus {
    pub keyspace_id: u32,
    pub keyspace_name: String,
    pub encryption_config: Option<KeyspaceEncryptionConfig>,
    pub meta_state: PdEncryptionMetaState,
    pub encryption_meta_job_owner_store_id: Option<u64>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KeyspaceEncryptionMetaId {
    /// Keyspace is not encrypted in the selected (stable/latest) view.
    Unencrypted,
    /// PD indicates the encryption meta can be fetched via DFS with this file
    /// id.
    MetaFileId(u64),
    /// Encryption is enabled, but PD does not have a usable file id yet.
    ///
    /// Callers may choose to fall back to in-memory cached encryption meta.
    UnavailableInPd,
}

impl PdKeyspaceEncryptionStatus {
    /// Selects the encryption meta file id according to PD keyspace state.
    ///
    /// - `latest_unstable=false` returns the stable view (usually
    ///   `current_file_id`)
    /// - `latest_unstable=true` returns the latest view (uses `next_file_id` if
    ///   set)
    pub fn select_meta_id(&self, latest_unstable: bool) -> KeyspaceEncryptionMetaId {
        use KeyspaceEncryptionMetaId::*;

        let encryption_enabled = self
            .encryption_config
            .as_ref()
            .is_some_and(|config| config.enabled);
        if !encryption_enabled {
            return Unencrypted;
        }

        match (
            self.meta_state.current_file_id,
            self.meta_state.next_file_id,
            latest_unstable,
        ) {
            // This only happens when a keyspace is created with CMEK enabled
            // and has not gone through a key rotation yet.
            // TODO: Backfill `current_file_id` later to eliminate this case.
            (0, 0, _) => UnavailableInPd,
            // Stable view has not switched to an encryption meta yet, even if
            // PD has already staged the first `next_file_id`.
            (0, _, false) => Unencrypted,
            // Latest view prefers the staged `next_file_id` when PD has one.
            (_, next, true) if next != 0 => MetaFileId(next),
            // Remaining cases always have `current_file_id != 0`, so both the
            // stable view and latest-without-next fall back to `current`.
            (current, ..) => MetaFileId(current),
        }
    }
}

#[derive(Debug, thiserror::Error)]
pub enum EncryptionMetaLoadError {
    #[error("dfs error: {0}")]
    Dfs(#[from] dfs::Error),
    #[error("failed to decode encryption meta file_id {file_id}: {source}")]
    Decode {
        file_id: u64,
        source: protobuf::ProtobufError,
    },
}

pub async fn load_encryption_meta_from_dfs(
    fs: Arc<dyn dfs::Dfs>,
    meta_file_id: u64,
) -> Result<Arc<EncryptionMeta>, EncryptionMetaLoadError> {
    let opts = dfs::Options {
        file_type: dfs::FileType::EncryptionDict,
        ..Default::default()
    };

    let meta_bytes = fs.read_file(meta_file_id, opts).await?;
    let mut meta = EncryptionMeta::default();
    meta.merge_from_bytes(&meta_bytes)
        .map_err(|e| EncryptionMetaLoadError::Decode {
            file_id: meta_file_id,
            source: e,
        })?;

    // Keep this permissive: after restore, a cluster may intentionally reuse an
    // existing encryption meta file for a different keyspace.
    Ok(Arc::new(meta))
}

#[derive(Clone)]
pub struct EncryptionConfChangeManager {
    inner: Arc<EncryptionConfChangeManagerInner>,
}

type PdHttpContext = (Arc<dyn PdClient>, Arc<SecurityManager>, HttpClient);

struct EncryptionConfChangeManagerInner {
    owner_check: OwnerCheckMode,
    pd_client: Arc<dyn PdClient>,
    jobs: Mutex<HashMap<u32, RunningJob>>,
    next_job_id: AtomicU64,
}

#[derive(Clone, Debug)]
enum JobStatus {
    Running,
    Finished(Result<RolloutJobSummary, EncryptionConfChangeError>),
}

struct RunningJob {
    job_id: u64,
    is_stuck: bool,
    cancel_tx: watch::Sender<bool>,
    status_rx: watch::Receiver<JobStatus>,
}

pub struct StartRolloutJobResult {
    pub response: StartRolloutJobResponse,
    pub handle: Option<RolloutJobHandle>,
}

pub struct RolloutJobHandle {
    status_rx: watch::Receiver<JobStatus>,
}

impl RolloutJobHandle {
    pub async fn wait(self) -> Result<RolloutJobSummary, EncryptionConfChangeError> {
        wait_job_finished(self.status_rx).await
    }
}

impl EncryptionConfChangeManager {
    pub fn new(pd_client: Arc<dyn PdClient>, owner_check: OwnerCheckMode) -> Self {
        Self {
            inner: Arc::new(EncryptionConfChangeManagerInner {
                owner_check,
                pd_client,
                jobs: Mutex::new(HashMap::new()),
                next_job_id: AtomicU64::new(1),
            }),
        }
    }

    pub fn load_keyspace_encryption_status_from_pd(
        &self,
        keyspace_id: u32,
    ) -> Result<PdKeyspaceEncryptionStatus, EncryptionConfChangeError> {
        let meta = self.load_keyspace_meta_from_pd(keyspace_id)?;
        let config = meta.get_config();

        let current_file_id =
            parse_keyspace_config_u64(config, KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID)?
                .unwrap_or(0);
        let next_file_id =
            parse_keyspace_config_u64(config, KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID)?
                .unwrap_or(0);
        validate_ids(PdEncryptionMetaState {
            current_file_id,
            next_file_id,
        })?;

        let encryption_config = match config.get(KEYSPACE_CONFIG_ENCRYPTION) {
            None => None,
            Some(value) => Some(
                KeyspaceEncryptionConfig::from_pd_config_value(value).map_err(|err| {
                    EncryptionConfChangeError::failed(format!(
                        "invalid PD keyspace encryption config JSON for keyspace {}: {}",
                        keyspace_id, err
                    ))
                })?,
            ),
        };

        let encryption_meta_job_owner_store_id =
            parse_keyspace_config_u64(config, KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID)?
                .filter(|id| *id != 0);

        Ok(PdKeyspaceEncryptionStatus {
            keyspace_id: meta.get_id(),
            keyspace_name: meta.get_name().to_owned(),
            encryption_config,
            meta_state: PdEncryptionMetaState {
                current_file_id,
                next_file_id,
            },
            encryption_meta_job_owner_store_id,
        })
    }

    async fn resolve_effective_current_file_id_for_gen_meta(
        &self,
        status: &PdKeyspaceEncryptionStatus,
        mode: &EncryptionMetaGenMode,
        security_mgr: &SecurityManager,
        http_client: &HttpClient,
        first_region: &KeyspaceFirstRegion,
    ) -> Result<u64, EncryptionConfChangeError> {
        let mut effective_current_file_id = status.meta_state.current_file_id;

        // Keyspaces created with CMEK already have shard-level encryption meta,
        // but PD may still show `current_file_id=0` until the first explicit
        // rotation backfills it. For rotate-data-key, resolve the real current
        // file id from the anchor region and feed that into the existing
        // current/next rollout state machine.
        if matches!(mode, EncryptionMetaGenMode::RotateDataKey { .. })
            && effective_current_file_id == 0
        {
            let cmek_enabled = status.encryption_config.as_ref().is_some_and(|c| c.enabled);
            if !cmek_enabled {
                return Err(EncryptionConfChangeError::failed(format!(
                    "keyspace {} has no current encryption meta (PD current_file_id=0); use enable-cmek first",
                    status.keyspace_id
                )));
            }

            effective_current_file_id = self
                .request_get_region_meta(
                    security_mgr,
                    http_client,
                    &first_region.store_status_addr,
                    status.keyspace_id,
                    first_region.region_id,
                )
                .await?;
        }

        Ok(effective_current_file_id)
    }

    fn build_pd_http_client(&self) -> Result<PdHttpContext, EncryptionConfChangeError> {
        let pd_client = self.inner.pd_client.clone();
        let security_mgr = pd_client.get_security_mgr();
        let http_client = security_mgr
            .http_client(hyper::Client::builder())
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("failed to build http client: {}", e))
            })?;
        Ok((pd_client, security_mgr, http_client))
    }

    /// Resolves the first region in the keyspace together with its leader
    /// store metadata.
    async fn resolve_first_region_of_keyspace(
        &self,
        pd_client: &dyn PdClient,
        keyspace_id: u32,
    ) -> Result<KeyspaceFirstRegion, EncryptionConfChangeError> {
        let ks_start_key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        let encoded_start_key = Key::from_raw(&ks_start_key).into_encoded();
        let first_region_info = pd_client
            .get_region_info_async(&encoded_start_key)
            .await
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("failed to get first region: {}", e))
            })?;

        let region_id = first_region_info.region.get_id();
        let leader = first_region_info.leader.as_ref().ok_or_else(|| {
            EncryptionConfChangeError::failed(format!(
                "no leader found for first region {}",
                region_id
            ))
        })?;
        let leader_store_id = leader.get_store_id();
        let leader_store = pd_client
            .get_store_async(leader_store_id)
            .await
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("get_store_async failed: {}", e))
            })?;

        let store_status_addr = leader_store.get_status_address().to_owned();
        if store_status_addr.is_empty() {
            return Err(EncryptionConfChangeError::failed(format!(
                "store {} has empty status address",
                leader_store_id
            )));
        }

        Ok(KeyspaceFirstRegion {
            region_id,
            store_status_addr,
            store_id: leader_store_id,
        })
    }

    async fn request_gen_meta(
        &self,
        security_mgr: &SecurityManager,
        http_client: &HttpClient,
        store_status_addr: &str,
        keyspace_id: u32,
        region_id: u64,
        mode: &EncryptionMetaGenMode,
    ) -> Result<EncryptionMetaGenResponse, EncryptionConfChangeError> {
        generate_encryption_meta(
            security_mgr,
            http_client,
            store_status_addr,
            keyspace_id,
            region_id,
            mode.clone(),
        )
        .await
    }

    async fn request_get_region_meta(
        &self,
        security_mgr: &SecurityManager,
        http_client: &HttpClient,
        store_status_addr: &str,
        keyspace_id: u32,
        region_id: u64,
    ) -> Result<u64, EncryptionConfChangeError> {
        let uri = build_uri(
            security_mgr,
            store_status_addr,
            &format!("/encryption/get-meta?region_id={}", region_id),
        )?;
        let req = Request::get(uri).body(Body::empty()).map_err(|e| {
            EncryptionConfChangeError::failed(format!("failed to build request: {}", e))
        })?;

        let resp = tokio::time::timeout(DEFAULT_REMOTE_GEN_META_TIMEOUT, http_client.request(req))
            .await
            .map_err(|_| EncryptionConfChangeError::failed("request timeout".to_owned()))?
            .map_err(|e| EncryptionConfChangeError::failed(format!("request failed: {}", e)))?;
        let status = resp.status();
        let body = hyper::body::to_bytes(resp.into_body()).await.map_err(|e| {
            EncryptionConfChangeError::failed(format!("failed to read response body: {}", e))
        })?;

        if status == StatusCode::NO_CONTENT {
            return Err(EncryptionConfChangeError::failed(format!(
                "keyspace {} has encryption enabled in PD but region {} has no current encryption meta",
                keyspace_id, region_id
            )));
        }
        if !status.is_success() {
            return Err(EncryptionConfChangeError::failed(format!(
                "encryption/get-meta failed ({}): {}",
                status,
                String::from_utf8_lossy(&body)
            )));
        }

        let mut meta = EncryptionMeta::default();
        meta.merge_from_bytes(&body).map_err(|e| {
            EncryptionConfChangeError::failed(format!(
                "failed to parse encryption meta for region {}: {}",
                region_id, e
            ))
        })?;
        if meta.get_current().get_file_id() == 0 {
            return Err(EncryptionConfChangeError::failed(format!(
                "region {} returned encryption meta with current.file_id=0",
                region_id
            )));
        }

        Ok(meta.get_current().get_file_id())
    }

    async fn apply_gen_meta_to_pd(
        &self,
        keyspace_id: u32,
        pd_before: PdEncryptionMetaState,
        effective_current_file_id: u64,
        owner_store_id: u64,
        mode: &EncryptionMetaGenMode,
        gen_resp: EncryptionMetaGenResponse,
    ) -> Result<GenMetaResult, EncryptionConfChangeError> {
        // Persist the generated file ID into PD keyspace config as `next_file_id`, and
        // record `owner_store_id` so subsequent rollout requests can be routed
        // to a single owner.
        if gen_resp.file_id == 0 {
            return Err(EncryptionConfChangeError::failed(
                "invalid generated file_id=0".to_owned(),
            ));
        }

        if gen_resp.new_data_key_generated == Some(false) {
            if gen_resp.file_id != effective_current_file_id {
                return Err(EncryptionConfChangeError::failed(format!(
                    "unexpected generated file_id={} when rotation is skipped; expected current_file_id={}",
                    gen_resp.file_id, effective_current_file_id
                )));
            }

            return Ok(GenMetaResult {
                keyspace_id,
                pd_before,
                pd_after: pd_before,
                owner_store_id: None,
                gen_meta: gen_resp,
                pd_updated: false,
            });
        }

        if gen_resp.file_id <= effective_current_file_id {
            return Err(EncryptionConfChangeError::failed(format!(
                "invalid generated file_id={}: must be greater than current_file_id={}",
                gen_resp.file_id, effective_current_file_id
            )));
        }

        let target_state = PdEncryptionMetaState {
            current_file_id: effective_current_file_id,
            next_file_id: gen_resp.file_id,
        };

        let mut extra_config = HashMap::new();
        extra_config.insert(
            KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID.to_owned(),
            Some(owner_store_id.to_string()),
        );

        if let EncryptionMetaGenMode::EnableCmek {
            cmek_id,
            vendor,
            region,
            endpoint,
            encryption_context,
        } = mode
        {
            let encryption_cfg =
                build_enable_cmek_config(cmek_id, vendor, region, endpoint, encryption_context);
            let encryption_json = encryption_cfg.to_pd_config_value().map_err(|e| {
                EncryptionConfChangeError::failed(format!(
                    "failed to serialize encryption config: {}",
                    e
                ))
            })?;
            extra_config.insert(KEYSPACE_CONFIG_ENCRYPTION.to_owned(), Some(encryption_json));
        }

        let pd_after = self
            .update_keyspace_config(keyspace_id, pd_before, target_state, extra_config)
            .await?;

        Ok(GenMetaResult {
            keyspace_id,
            pd_before,
            pd_after,
            owner_store_id: Some(owner_store_id),
            gen_meta: gen_resp,
            pd_updated: true,
        })
    }

    pub async fn handle_gen_meta(
        &self,
        req: GenMetaRequest,
    ) -> Result<GenMetaResult, EncryptionConfChangeError> {
        // Orchestrates a gen-meta operation:
        // 1) precheck PD state (no pending `next_file_id`, etc)
        // 2) resolve a target TiKV store to generate meta on (anchor region leader)
        // 3) request gen-meta from that store's status endpoint
        // 4) update PD keyspace config (file IDs, owner store id, and optional
        //    encryption config)
        let GenMetaRequest { keyspace_id, mode } = req;

        let (pd_client, security_mgr, http_client) = self.build_pd_http_client()?;
        let status = self.load_keyspace_encryption_status_from_pd(keyspace_id)?;
        let pd_before = status.meta_state;

        if let EncryptionMetaGenMode::EnableCmek {
            cmek_id,
            vendor,
            region,
            endpoint,
            encryption_context,
        } = &mode
        {
            if let Some(current_cfg) = status.encryption_config.as_ref() {
                let requested_cfg =
                    build_enable_cmek_config(cmek_id, vendor, region, endpoint, encryption_context);
                if cmek_configs_match(current_cfg, &requested_cfg) {
                    let file_id = if pd_before.next_file_id != 0 {
                        pd_before.next_file_id
                    } else {
                        pd_before.current_file_id
                    };
                    return Ok(GenMetaResult {
                        keyspace_id,
                        pd_before,
                        pd_after: pd_before,
                        owner_store_id: status.encryption_meta_job_owner_store_id,
                        gen_meta: EncryptionMetaGenResponse {
                            file_id,
                            data_key_id: 0,
                            new_data_key_generated: Some(pd_before.next_file_id != 0),
                        },
                        pd_updated: false,
                    });
                }
            }
        }

        if pd_before.next_file_id != 0 {
            return Err(EncryptionConfChangeError::failed(format!(
                "keyspace {} already has a pending conf change: current_file_id={}, next_file_id={}",
                keyspace_id, pd_before.current_file_id, pd_before.next_file_id
            )));
        }

        let first_region = self
            .resolve_first_region_of_keyspace(pd_client.as_ref(), keyspace_id)
            .await?;
        let effective_current_file_id = self
            .resolve_effective_current_file_id_for_gen_meta(
                &status,
                &mode,
                security_mgr.as_ref(),
                &http_client,
                &first_region,
            )
            .await?;
        let gen_resp = self
            .request_gen_meta(
                security_mgr.as_ref(),
                &http_client,
                &first_region.store_status_addr,
                keyspace_id,
                first_region.region_id,
                &mode,
            )
            .await?;

        self.apply_gen_meta_to_pd(
            keyspace_id,
            pd_before,
            effective_current_file_id,
            first_region.store_id,
            &mode,
            gen_resp,
        )
        .await
    }

    pub async fn update_keyspace_config(
        &self,
        keyspace_id: u32,
        expected: PdEncryptionMetaState,
        target: PdEncryptionMetaState,
        extra_config: HashMap<String, Option<String>>,
    ) -> Result<PdEncryptionMetaState, EncryptionConfChangeError> {
        // Patch PD keyspace config with server-side compare-and-swap semantics:
        // 1) verify current PD state matches `expected`.
        // 2) build `preconditions` pinned to PD's current raw values so the patch is
        //    applied atomically only if PD state has not changed in between.
        // 3) reload and verify PD state matches `target` as a defensive check (covers
        //    older PD versions that ignore `preconditions`).
        let expected = validate_ids(expected)?;
        let target = validate_ids(target)?;

        let meta = self.load_keyspace_meta_from_pd(keyspace_id)?;
        let current = parse_pd_meta_state(meta.get_config())?;
        if current != expected {
            return Err(EncryptionConfChangeError::failed(format!(
                "PD encryption meta file ID state changed concurrently: expected current_file_id={}, next_file_id={}, actual current_file_id={}, next_file_id={}",
                expected.current_file_id,
                expected.next_file_id,
                current.current_file_id,
                current.next_file_id
            )));
        }
        let preconditions = build_ids_preconditions(meta.get_config());

        let mut config = build_ids_config(target);
        config.extend(extra_config);

        self.inner
            .pd_client
            .patch_keyspace_config(keyspace_id, config, Some(preconditions))
            .await
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("patch_keyspace_config failed: {}", e))
            })?;

        let updated = self
            .load_keyspace_encryption_status_from_pd(keyspace_id)?
            .meta_state;
        if updated != target {
            return Err(EncryptionConfChangeError::failed(format!(
                "PD encryption meta file ID state verification failed after update: expected current_file_id={}, next_file_id={}, actual current_file_id={}, next_file_id={}",
                target.current_file_id,
                target.next_file_id,
                updated.current_file_id,
                updated.next_file_id
            )));
        }

        Ok(updated)
    }

    fn load_keyspace_meta_from_pd(
        &self,
        keyspace_id: u32,
    ) -> Result<kvproto::keyspacepb::KeyspaceMeta, EncryptionConfChangeError> {
        let meta = self
            .inner
            .pd_client
            .load_keyspace_by_id(keyspace_id)
            .map_err(|e| {
                if e.is_keyspace_not_found() {
                    EncryptionConfChangeError::KeyspaceNotFound(keyspace_id)
                } else {
                    EncryptionConfChangeError::failed(format!("load_keyspace_by_id failed: {}", e))
                }
            })?;

        if meta.get_id() != keyspace_id {
            return Err(EncryptionConfChangeError::KeyspaceNotFound(keyspace_id));
        }
        Ok(meta)
    }

    /// Starts a rollout job on the owner store recorded in PD.
    ///
    /// This is a client-side helper: it resolves the owner store ID from PD and
    /// sends an HTTP request to that store's status server endpoint
    /// (`/encryption/start_rollout_job`). The actual job is executed on the
    /// owner store.
    pub async fn request_start_rollout_job_on_owner(
        &self,
        keyspace_id: u32,
    ) -> Result<StartRolloutJobResponse, EncryptionConfChangeError> {
        let status = self.load_keyspace_encryption_status_from_pd(keyspace_id)?;
        let expected_state = status.meta_state;
        if expected_state.next_file_id == 0 {
            return Ok(StartRolloutJobResponse {
                keyspace_id,
                started: false,
                running: false,
                current_file_id: expected_state.current_file_id,
                next_file_id: expected_state.next_file_id,
                note: Some("no_pending_conf_change".to_owned()),
            });
        }

        let owner_store_id = status.encryption_meta_job_owner_store_id.unwrap_or(0);
        if owner_store_id == 0 {
            return Err(EncryptionConfChangeError::failed(format!(
                "PD keyspace config missing {}",
                KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID
            )));
        }

        let pd_client = self.inner.pd_client.clone();
        let security_mgr = pd_client.get_security_mgr();
        let http_client = security_mgr
            .http_client(hyper::Client::builder())
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("failed to build http client: {}", e))
            })?;

        // The owner store is persisted to PD during gen-meta so that the cloud
        // control plane can trigger a rollout asynchronously without scanning
        // all regions itself. If the owner becomes unavailable before sync,
        // remote execution can be blocked until owner recovers.
        //
        // To mitigate this, we opportunistically re-select the owner from the
        // keyspace anchor region leader when PD marks the previous owner as
        // non-up (or missing status address), and update PD config accordingly.
        let store = pd_client
            .get_store_async(owner_store_id)
            .await
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!(
                    "get_store_async failed: store_id={}, err={}",
                    owner_store_id, e
                ))
            })?;
        let mut target_store_id = owner_store_id;
        let mut store_status_addr = store.get_status_address().to_owned();

        if store.get_state() != StoreState::Up || store_status_addr.is_empty() {
            let first_region = self
                .resolve_first_region_of_keyspace(pd_client.as_ref(), keyspace_id)
                .await?;
            if first_region.store_id != owner_store_id {
                let mut extra_config = HashMap::new();
                extra_config.insert(
                    KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID.to_owned(),
                    Some(first_region.store_id.to_string()),
                );
                self.update_keyspace_config(
                    keyspace_id,
                    expected_state,
                    expected_state,
                    extra_config,
                )
                .await?;
            }
            target_store_id = first_region.store_id;
            store_status_addr = first_region.store_status_addr;
        }
        if store_status_addr.is_empty() {
            return Err(EncryptionConfChangeError::failed(format!(
                "store {} has empty status address",
                target_store_id
            )));
        }

        let uri = build_uri(
            security_mgr.as_ref(),
            store_status_addr.as_str(),
            "/encryption/start_rollout_job",
        )?;
        let body_bytes =
            serde_json::to_vec(&StartRolloutJobRequest { keyspace_id }).map_err(|e| {
                EncryptionConfChangeError::failed(format!("failed to serialize request: {}", e))
            })?;
        let req = Request::post(uri)
            .header(hyper::header::CONTENT_TYPE, "application/json")
            .body(Body::from(body_bytes))
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("failed to build request: {}", e))
            })?;

        let resp = tokio::time::timeout(
            DEFAULT_REMOTE_START_ROLLOUT_TIMEOUT,
            http_client.request(req),
        )
        .await
        .map_err(|_| EncryptionConfChangeError::failed("request timeout".to_owned()))?
        .map_err(|e| EncryptionConfChangeError::failed(format!("request failed: {}", e)))?;
        let status = resp.status();
        let body = hyper::body::to_bytes(resp.into_body()).await.map_err(|e| {
            EncryptionConfChangeError::failed(format!("failed to read response body: {}", e))
        })?;

        if !status.is_success() {
            return Err(EncryptionConfChangeError::failed(format!(
                "encryption/start_rollout_job failed ({}): {}",
                status,
                String::from_utf8_lossy(&body)
            )));
        }

        serde_json::from_slice(&body).map_err(|e| {
            EncryptionConfChangeError::failed(format!(
                "failed to parse start-rollout response: {}, body={}",
                e,
                String::from_utf8_lossy(&body)
            ))
        })
    }

    /// Starts a local rollout job if needed.
    ///
    /// This is the TiKV-side entrypoint used by the status server endpoint
    /// (`/encryption/start_rollout_job`). The job itself runs asynchronously
    /// and performs HTTP requests to other TiKV status endpoints via
    /// `/encryption/sync-latest` to sync region encryption meta, then finalizes
    /// PD state.
    pub async fn handle_start_rollout_job(
        &self,
        req: StartRolloutJobRequest,
    ) -> Result<StartRolloutJobResult, EncryptionConfChangeError> {
        let keyspace_id = req.keyspace_id;

        let status = self.load_keyspace_encryption_status_from_pd(keyspace_id)?;
        let meta_state = status.meta_state;

        if meta_state.next_file_id == 0 {
            return Ok(StartRolloutJobResult {
                response: StartRolloutJobResponse {
                    keyspace_id,
                    started: false,
                    running: false,
                    current_file_id: meta_state.current_file_id,
                    next_file_id: meta_state.next_file_id,
                    note: Some("no_pending_conf_change".to_owned()),
                },
                handle: None,
            });
        }

        self.check_owner(&status)?;

        {
            let job_guard = self.inner.jobs.lock().unwrap();
            if job_guard.contains_key(&keyspace_id) {
                return Ok(StartRolloutJobResult {
                    response: StartRolloutJobResponse {
                        keyspace_id,
                        started: false,
                        running: true,
                        current_file_id: meta_state.current_file_id,
                        next_file_id: meta_state.next_file_id,
                        note: Some("already_running".to_owned()),
                    },
                    handle: None,
                });
            }
        }

        let (cancel_tx, cancel_rx) = watch::channel(false);
        let (status_tx, status_rx) = watch::channel::<JobStatus>(JobStatus::Running);
        let handle = RolloutJobHandle {
            status_rx: status_rx.clone(),
        };
        let job_id = self.inner.next_job_id.fetch_add(1, Ordering::Relaxed);
        let job_start = Instant::now();

        {
            let mut job_guard = self.inner.jobs.lock().unwrap();
            if job_guard.contains_key(&keyspace_id) {
                return Ok(StartRolloutJobResult {
                    response: StartRolloutJobResponse {
                        keyspace_id,
                        started: false,
                        running: true,
                        current_file_id: meta_state.current_file_id,
                        next_file_id: meta_state.next_file_id,
                        note: Some("already_running".to_owned()),
                    },
                    handle: None,
                });
            }
            job_guard.insert(
                keyspace_id,
                RunningJob {
                    job_id,
                    is_stuck: false,
                    cancel_tx,
                    status_rx,
                },
            );
            update_job_metrics(&job_guard);
        }

        let mgr = self.clone();
        tokio::spawn(async move {
            let outcome = mgr
                .do_rollout(keyspace_id, job_id, meta_state, cancel_rx, job_start)
                .await;
            let elapsed = job_start.elapsed();

            match &outcome {
                Ok(summary) => {
                    info!(
                        "[CMEK] encryption conf-change job finished: keyspace_id={}, regions_success={}, elapsed_ms={}",
                        keyspace_id,
                        summary.regions_success,
                        elapsed.as_millis() as u64
                    );
                }
                Err(EncryptionConfChangeError::Canceled) => {
                    info!(
                        "[CMEK] encryption conf-change job canceled: keyspace_id={}, elapsed_ms={}",
                        keyspace_id,
                        elapsed.as_millis() as u64
                    );
                }
                Err(e) => {
                    ENCRYPTION_CONF_CHANGE_FAILED_JOBS_COUNTER.inc();
                    error!(
                        "[CMEK] encryption conf-change job failed: keyspace_id={}, elapsed_ms={}, err={}",
                        keyspace_id,
                        elapsed.as_millis() as u64,
                        e
                    );
                }
            }

            let _ = status_tx.send(JobStatus::Finished(outcome));
            mgr.remove_job_if_match(keyspace_id, job_id);
        });

        Ok(StartRolloutJobResult {
            response: StartRolloutJobResponse {
                keyspace_id,
                started: true,
                running: true,
                current_file_id: meta_state.current_file_id,
                next_file_id: meta_state.next_file_id,
                note: None,
            },
            handle: Some(handle),
        })
    }

    fn check_owner(
        &self,
        status: &PdKeyspaceEncryptionStatus,
    ) -> Result<(), EncryptionConfChangeError> {
        let OwnerCheckMode::Enforced { owner_store_id } = &self.inner.owner_check else {
            return Ok(());
        };

        let Some(pd_owner_store_id) = status.encryption_meta_job_owner_store_id else {
            return Ok(());
        };
        if pd_owner_store_id == 0 || pd_owner_store_id == *owner_store_id {
            return Ok(());
        }

        Err(EncryptionConfChangeError::WrongOwner {
            expected_owner_store_id: *owner_store_id,
            pd_owner_store_id,
        })
    }

    /// Performs the actual rollout procedure for a keyspace.
    ///
    /// This scans all regions in the keyspace via PD, and for each region sends
    /// an HTTP request to the region leader's status endpoint
    /// (`/encryption/sync-latest`) to ensure the latest encryption meta is
    /// applied. After all regions succeed, it finalizes PD keyspace config by
    /// advancing `current_file_id` and clearing `next_file_id`.
    async fn do_rollout(
        &self,
        keyspace_id: u32,
        job_id: u64,
        meta_state: PdEncryptionMetaState,
        cancel_rx: watch::Receiver<bool>,
        job_start: Instant,
    ) -> Result<RolloutJobSummary, EncryptionConfChangeError> {
        if *cancel_rx.borrow() {
            return Err(EncryptionConfChangeError::Canceled);
        }

        let status = self.load_keyspace_encryption_status_from_pd(keyspace_id)?;
        if status.meta_state != meta_state {
            return Err(EncryptionConfChangeError::failed(format!(
                "PD keyspace config changed concurrently: expected current_file_id={}, next_file_id={}, actual current_file_id={}, next_file_id={}",
                meta_state.current_file_id,
                meta_state.next_file_id,
                status.meta_state.current_file_id,
                status.meta_state.next_file_id
            )));
        }
        if meta_state.next_file_id == 0 {
            return Ok(RolloutJobSummary::default());
        }

        let pd_client = self.inner.pd_client.clone();
        let security_mgr = pd_client.get_security_mgr();
        let http_client = security_mgr
            .http_client(hyper::Client::builder())
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("failed to build http client: {}", e))
            })?;

        let current_file_id = meta_state.current_file_id;
        let next_file_id = meta_state.next_file_id;

        let mut summary = RolloutJobSummary::default();
        let mut scanner = KeyspaceRegionScanner::for_keyspace(keyspace_id);
        let sync_start = Instant::now();
        let mut sync_region_retry_count = 0usize;

        loop {
            self.mark_job_stuck_if_needed(keyspace_id, job_id, job_start);
            let checkpoint = scanner.checkpoint();
            let Some(region_info) = (match scanner.next_region_async(pd_client.as_ref()).await {
                Ok(region_info) => region_info,
                Err(EncryptionSyncError::Retryable(msg)) => {
                    sync_region_retry_count += 1;
                    scanner.rewind(checkpoint);
                    if sync_region_retry_count > DEFAULT_SYNC_REGION_RETRY_LIMIT {
                        return Err(EncryptionConfChangeError::failed(format!(
                            "failed to sync region after {} retries: err={}",
                            DEFAULT_SYNC_REGION_RETRY_LIMIT, msg
                        )));
                    }
                    wait_sync_region_retry_delay(&cancel_rx).await?;
                    continue;
                }
                Err(EncryptionSyncError::Failed(msg)) => {
                    return Err(EncryptionConfChangeError::failed(msg));
                }
                Err(EncryptionSyncError::Canceled) => {
                    return Err(EncryptionConfChangeError::Canceled);
                }
            }) else {
                break;
            };
            if *cancel_rx.borrow() {
                return Err(EncryptionConfChangeError::Canceled);
            }

            let region_id = region_info.region.get_id();
            let leader = match region_info.leader.as_ref() {
                Some(leader) => leader,
                None => {
                    sync_region_retry_count += 1;
                    scanner.rewind(checkpoint);
                    if sync_region_retry_count > DEFAULT_SYNC_REGION_RETRY_LIMIT {
                        return Err(EncryptionConfChangeError::failed(format!(
                            "failed to sync region after {} retries: region_id={}, err=no leader in PD region info",
                            DEFAULT_SYNC_REGION_RETRY_LIMIT, region_id
                        )));
                    }
                    wait_sync_region_retry_delay(&cancel_rx).await?;
                    continue;
                }
            };

            match sync_region_encryption_meta(SyncRegionEncryptionMetaRequest {
                pd_client: pd_client.as_ref(),
                security_mgr: security_mgr.as_ref(),
                http_client: &http_client,
                keyspace_id,
                region_id,
                leader_store_id: leader.get_store_id(),
                current_file_id,
                next_file_id,
                expected_region_version: region_info.region.get_region_epoch().get_version(),
                expected_region_conf_ver: region_info.region.get_region_epoch().get_conf_ver(),
                timeout: DEFAULT_SYNC_LATEST_TIMEOUT,
                cancel_rx: Some(cancel_rx.clone()),
            })
            .await
            {
                Ok(_) => {
                    sync_region_retry_count = 0;
                    summary.regions_success += 1;
                }
                Err(EncryptionSyncError::Canceled) => {
                    return Err(EncryptionConfChangeError::Canceled);
                }
                Err(EncryptionSyncError::Retryable(msg)) => {
                    sync_region_retry_count += 1;
                    scanner.rewind(checkpoint);
                    if sync_region_retry_count > DEFAULT_SYNC_REGION_RETRY_LIMIT {
                        return Err(EncryptionConfChangeError::failed(format!(
                            "failed to sync region after {} retries: region_id={}, err={}",
                            DEFAULT_SYNC_REGION_RETRY_LIMIT, region_id, msg
                        )));
                    }
                    wait_sync_region_retry_delay(&cancel_rx).await?;
                }
                Err(EncryptionSyncError::Failed(msg)) => {
                    return Err(EncryptionConfChangeError::failed(format!(
                        "failed to sync region {}: {}",
                        region_id, msg
                    )));
                }
            }
        }

        info!(
            "[CMEK] encryption conf-change job sync finished: keyspace_id={}, current_file_id={}, next_file_id={}, regions_success={}, elapsed_ms={}",
            keyspace_id,
            current_file_id,
            next_file_id,
            summary.regions_success,
            sync_start.elapsed().as_millis() as u64
        );

        if *cancel_rx.borrow() {
            return Err(EncryptionConfChangeError::Canceled);
        }

        let meta = self.load_keyspace_meta_from_pd(keyspace_id)?;
        let current = parse_pd_meta_state(meta.get_config())?;
        let expected = PdEncryptionMetaState {
            current_file_id,
            next_file_id,
        };
        if current != expected {
            return Err(EncryptionConfChangeError::failed(format!(
                "PD keyspace config changed concurrently before finalize: expected current_file_id={}, next_file_id={}, actual current_file_id={}, next_file_id={}",
                expected.current_file_id,
                expected.next_file_id,
                current.current_file_id,
                current.next_file_id
            )));
        }
        let preconditions = build_ids_preconditions(meta.get_config());

        let mut config = HashMap::new();
        config.insert(
            KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID.to_owned(),
            Some(next_file_id.to_string()),
        );
        config.insert(
            KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID.to_owned(),
            None,
        );
        config.insert(
            KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID.to_owned(),
            None,
        );

        pd_client
            .patch_keyspace_config(keyspace_id, config, Some(preconditions))
            .await
            .map_err(|e| {
                EncryptionConfChangeError::failed(format!("patch_keyspace_config failed: {}", e))
            })?;

        let final_state = self
            .load_keyspace_encryption_status_from_pd(keyspace_id)?
            .meta_state;
        if final_state.current_file_id != next_file_id || final_state.next_file_id != 0 {
            return Err(EncryptionConfChangeError::failed(format!(
                "PD keyspace config verification failed after finalize: expected current_file_id={}, next_file_id=0, actual current_file_id={}, next_file_id={}",
                next_file_id, final_state.current_file_id, final_state.next_file_id
            )));
        }

        Ok(summary)
    }

    /// Cancels a local rollout job (if any) and waits for it to stop.
    pub async fn handle_cancel_rollout_job(
        &self,
        req: CancelRolloutJobRequest,
    ) -> Result<CancelRolloutJobResponse, EncryptionConfChangeError> {
        let keyspace_id = req.keyspace_id;
        let (job_id, status_rx) = {
            let job_guard = self.inner.jobs.lock().unwrap();
            let Some(job) = job_guard.get(&keyspace_id) else {
                return Ok(CancelRolloutJobResponse {
                    keyspace_id,
                    canceled: false,
                });
            };
            let _ = job.cancel_tx.send(true);
            (job.job_id, job.status_rx.clone())
        };

        let _ = wait_job_finished(status_rx).await;
        self.remove_job_if_match(keyspace_id, job_id);

        Ok(CancelRolloutJobResponse {
            keyspace_id,
            canceled: true,
        })
    }

    fn remove_job_if_match(&self, keyspace_id: u32, job_id: u64) {
        let mut job_guard = self.inner.jobs.lock().unwrap();
        let should_remove = job_guard
            .get(&keyspace_id)
            .is_some_and(|job| job.job_id == job_id);
        if should_remove {
            job_guard.remove(&keyspace_id);
            update_job_metrics(&job_guard);
        }
    }

    fn mark_job_stuck_if_needed(&self, keyspace_id: u32, job_id: u64, job_start: Instant) {
        if job_start.elapsed() < DEFAULT_STUCK_JOB_THRESHOLD {
            return;
        }

        let mut job_guard = self.inner.jobs.lock().unwrap();
        let Some(job) = job_guard.get_mut(&keyspace_id) else {
            return;
        };
        if job.job_id != job_id || job.is_stuck {
            return;
        }
        job.is_stuck = true;
        update_job_metrics(&job_guard);
    }
}

fn update_job_metrics(jobs: &HashMap<u32, RunningJob>) {
    ENCRYPTION_CONF_CHANGE_RUNNING_JOBS_GAUGE.set(jobs.len() as i64);
    ENCRYPTION_CONF_CHANGE_STUCK_JOBS_GAUGE
        .set(jobs.values().filter(|job| job.is_stuck).count() as i64);
}

fn validate_ids(
    state: PdEncryptionMetaState,
) -> Result<PdEncryptionMetaState, EncryptionConfChangeError> {
    if state.next_file_id != 0 && state.next_file_id <= state.current_file_id {
        return Err(EncryptionConfChangeError::failed(format!(
            "invalid encryption meta file ID state: current_file_id={}, next_file_id={}",
            state.current_file_id, state.next_file_id
        )));
    }
    Ok(state)
}

fn build_enable_cmek_config(
    cmek_id: &str,
    vendor: &str,
    region: &Option<String>,
    endpoint: &Option<String>,
    encryption_context: &BTreeMap<String, String>,
) -> KeyspaceEncryptionConfig {
    KeyspaceEncryptionConfig::enabled(
        Some(cmek_id.to_owned()),
        Some(vendor.to_owned()),
        region.clone(),
        endpoint.clone(),
        encryption_context.clone(),
    )
}

fn cmek_configs_match(lhs: &KeyspaceEncryptionConfig, rhs: &KeyspaceEncryptionConfig) -> bool {
    lhs.enabled == rhs.enabled
        && lhs.cmek_id == rhs.cmek_id
        && lhs.vendor == rhs.vendor
        && lhs.region == rhs.region
        && lhs.endpoint == rhs.endpoint
        && lhs.encryption_context == rhs.encryption_context
}

fn build_ids_config(state: PdEncryptionMetaState) -> HashMap<String, Option<String>> {
    let mut config = HashMap::new();
    config.insert(
        KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID.to_string(),
        Some(state.current_file_id.to_string()),
    );
    config.insert(
        KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID.to_string(),
        if state.next_file_id == 0 {
            None
        } else {
            Some(state.next_file_id.to_string())
        },
    );
    config
}

/// Builds keyspace-config preconditions for a server-side CAS update by
/// pinning the current raw values of the encryption meta ID keys observed in
/// PD. Callers are responsible for verifying that those raw values match the
/// expected state before submitting the patch.
fn build_ids_preconditions(config: &HashMap<String, String>) -> HashMap<String, Option<String>> {
    let mut preconditions = HashMap::new();
    preconditions.insert(
        KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID.to_owned(),
        config
            .get(KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID)
            .cloned(),
    );
    preconditions.insert(
        KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID.to_owned(),
        config
            .get(KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID)
            .cloned(),
    );
    preconditions
}

/// Parses encryption meta ID state from the raw PD keyspace config map.
fn parse_pd_meta_state(
    config: &HashMap<String, String>,
) -> Result<PdEncryptionMetaState, EncryptionConfChangeError> {
    let current_file_id =
        parse_keyspace_config_u64(config, KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID)?
            .unwrap_or(0);
    let next_file_id =
        parse_keyspace_config_u64(config, KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID)?
            .unwrap_or(0);
    Ok(PdEncryptionMetaState {
        current_file_id,
        next_file_id,
    })
}

fn parse_keyspace_config_u64(
    config: &HashMap<String, String>,
    key: &str,
) -> Result<Option<u64>, EncryptionConfChangeError> {
    let Some(value) = config.get(key) else {
        return Ok(None);
    };
    if value.is_empty() {
        return Ok(None);
    }
    let parsed = value.parse::<u64>().map_err(|e| {
        EncryptionConfChangeError::failed(format!("failed to parse {}={}: {}", key, value, e))
    })?;
    Ok(Some(parsed))
}

fn build_uri(
    security_mgr: &SecurityManager,
    addr: &str,
    path: &str,
) -> Result<Uri, EncryptionConfChangeError> {
    let addr = normalize_addr(addr);
    if addr.is_empty() {
        return Err(EncryptionConfChangeError::failed(
            "empty http address".to_owned(),
        ));
    }
    let path = path.trim_start_matches('/');
    security_mgr
        .build_uri(format!("{}/{}", addr, path))
        .map_err(|e| EncryptionConfChangeError::failed(e.to_string()))
}

fn normalize_addr(addr: &str) -> String {
    addr.trim_start_matches("http://")
        .trim_start_matches("https://")
        .trim_end_matches('/')
        .to_owned()
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EncryptionMetaGenRequest {
    pub keyspace_id: u32,
    pub region_id: u64,
    #[serde(flatten)]
    pub mode: EncryptionMetaGenMode,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "mode", rename_all = "snake_case")]
pub enum EncryptionMetaGenMode {
    RotateDataKey {
        #[serde(skip_serializing_if = "Option::is_none")]
        max_data_key_age: Option<u64>,
    },
    EnableCmek {
        cmek_id: String,
        vendor: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        region: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        endpoint: Option<String>,
        #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
        encryption_context: BTreeMap<String, String>,
    },
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EncryptionMetaGenResponse {
    pub file_id: u64,
    pub data_key_id: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub new_data_key_generated: Option<bool>,
}

pub async fn generate_encryption_meta(
    security_mgr: &SecurityManager,
    http_client: &HttpClient,
    store_status_addr: &str,
    keyspace_id: u32,
    region_id: u64,
    mode: EncryptionMetaGenMode,
) -> Result<EncryptionMetaGenResponse, EncryptionConfChangeError> {
    let gen_req = EncryptionMetaGenRequest {
        keyspace_id,
        region_id,
        mode,
    };
    let body_bytes = serde_json::to_vec(&gen_req).map_err(|e| {
        EncryptionConfChangeError::failed(format!("failed to serialize gen-meta request: {}", e))
    })?;
    let uri = build_uri(security_mgr, store_status_addr, "/encryption/gen-meta")?;
    let http_req = Request::post(uri)
        .header(hyper::header::CONTENT_TYPE, "application/json")
        .body(Body::from(body_bytes))
        .map_err(|e| {
            EncryptionConfChangeError::failed(format!("failed to build request: {}", e))
        })?;

    let resp = tokio::time::timeout(
        DEFAULT_REMOTE_GEN_META_TIMEOUT,
        http_client.request(http_req),
    )
    .await
    .map_err(|_| EncryptionConfChangeError::failed("request timeout".to_owned()))?
    .map_err(|e| EncryptionConfChangeError::failed(format!("request failed: {}", e)))?;
    let status = resp.status();
    let body = hyper::body::to_bytes(resp.into_body()).await.map_err(|e| {
        EncryptionConfChangeError::failed(format!("failed to read response body: {}", e))
    })?;
    if !status.is_success() {
        return Err(EncryptionConfChangeError::failed(format!(
            "encryption/gen-meta failed ({}): {}",
            status,
            String::from_utf8_lossy(&body)
        )));
    }

    serde_json::from_slice(&body).map_err(|e| {
        EncryptionConfChangeError::failed(format!(
            "failed to parse gen-meta response: {}, body={}",
            e,
            String::from_utf8_lossy(&body)
        ))
    })
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EncryptionMetaSyncLatestRequest {
    pub keyspace_id: u32,
    pub region_id: u64,
    pub current_file_id: u64,
    pub next_file_id: u64,
    pub expected_region_version: u64,
    pub expected_region_conf_ver: u64,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EncryptionMetaSyncLatestResponse {
    pub file_id: u64,
    pub data_key_id: u32,
    pub synced: bool,
}

pub struct SyncRegionEncryptionMetaRequest<'a> {
    pub pd_client: &'a dyn PdClient,
    pub security_mgr: &'a SecurityManager,
    pub http_client: &'a HttpClient,
    pub keyspace_id: u32,
    pub region_id: u64,
    pub leader_store_id: u64,
    pub current_file_id: u64,
    pub next_file_id: u64,
    pub timeout: Duration,
    pub expected_region_version: u64,
    pub expected_region_conf_ver: u64,
    pub cancel_rx: Option<watch::Receiver<bool>>,
}

/// Sends an HTTP request to a region leader store to sync encryption meta to
/// the latest file ID.
///
/// This calls the target store's status endpoint (`/encryption/sync-latest`).
pub async fn sync_region_encryption_meta(
    req: SyncRegionEncryptionMetaRequest<'_>,
) -> Result<EncryptionMetaSyncLatestResponse, EncryptionSyncError> {
    let req_body = EncryptionMetaSyncLatestRequest {
        keyspace_id: req.keyspace_id,
        region_id: req.region_id,
        current_file_id: req.current_file_id,
        next_file_id: req.next_file_id,
        expected_region_version: req.expected_region_version,
        expected_region_conf_ver: req.expected_region_conf_ver,
    };
    let body_bytes = serde_json::to_vec(&req_body)
        .map_err(|e| EncryptionSyncError::failed(format!("failed to serialize request: {}", e)))?;

    if let Some(rx) = req.cancel_rx.as_ref() {
        if *rx.borrow() {
            return Err(EncryptionSyncError::Canceled);
        }
    }

    let store = req
        .pd_client
        .get_store_async(req.leader_store_id)
        .await
        .map_err(|e| EncryptionSyncError::retryable(format!("failed to get store: {}", e)))?;

    let uri = build_uri(
        req.security_mgr,
        store.get_status_address(),
        "/encryption/sync-latest",
    )
    .map_err(|e| EncryptionSyncError::failed(e.to_string()))?;

    let http_req = Request::post(uri)
        .header(hyper::header::CONTENT_TYPE, "application/json")
        .body(Body::from(body_bytes.clone()))
        .map_err(|e| EncryptionSyncError::failed(format!("failed to build request: {}", e)))?;

    let resp = tokio::time::timeout(req.timeout, req.http_client.request(http_req))
        .await
        .map_err(|_| EncryptionSyncError::retryable("request timeout".to_owned()))?
        .map_err(|e| EncryptionSyncError::retryable(format!("request failed: {}", e)))?;

    let status = resp.status();
    let body = hyper::body::to_bytes(resp.into_body())
        .await
        .map_err(|e| EncryptionSyncError::retryable(format!("failed to read response: {}", e)))?;

    if status.is_success() {
        return serde_json::from_slice(&body)
            .map_err(|e| EncryptionSyncError::failed(format!("failed to parse response: {}", e)));
    }

    if status == StatusCode::CONFLICT || status == StatusCode::NOT_FOUND {
        return Err(EncryptionSyncError::retryable(format!(
            "encryption/sync-latest failed ({}): {}",
            status,
            String::from_utf8_lossy(&body)
        )));
    }

    Err(EncryptionSyncError::failed(format!(
        "encryption/sync-latest failed ({}): {}",
        status,
        String::from_utf8_lossy(&body)
    )))
}

pub struct KeyspaceRegionScanner {
    current_key: Vec<u8>,
    end_key: Vec<u8>,
    finished: bool,
}

impl KeyspaceRegionScanner {
    pub fn for_keyspace(keyspace_id: u32) -> Self {
        let start_key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        let end_key = ApiV2::get_keyspace_end_by_id(keyspace_id);
        Self {
            current_key: Key::from_raw(&start_key).into_encoded(),
            end_key: Key::from_raw(&end_key).into_encoded(),
            finished: false,
        }
    }

    pub async fn next_region_async(
        &mut self,
        pd_client: &dyn PdClient,
    ) -> Result<Option<RegionInfo>, EncryptionSyncError> {
        if self.finished
            || self.current_key.is_empty()
            || self.current_key.as_slice() >= self.end_key.as_slice()
        {
            self.finished = true;
            return Ok(None);
        }

        let current_key = self.current_key.clone();
        let region_info = pd_client
            .get_region_info_async(&current_key)
            .await
            .map_err(|e| match e {
                // PD may have a transient range gap while split heartbeats arrive
                // one-by-one. Retry from the same scan key instead of failing the job.
                PdError::RegionNotFound(_) => EncryptionSyncError::retryable(format!(
                    "failed to get region info at scan_key {:?}: {}",
                    current_key, e
                )),
                _ => EncryptionSyncError::failed(format!("failed to get region info: {}", e)),
            })?;

        let region_id = region_info.region.get_id();
        let region_end = region_info.region.get_end_key().to_vec();
        let at_keyspace_end =
            !region_end.is_empty() && region_end.as_slice() >= self.end_key.as_slice();

        if at_keyspace_end || region_end.is_empty() {
            self.finished = true;
            return Ok(Some(region_info));
        }

        if region_end == self.current_key {
            return Err(EncryptionSyncError::failed(format!(
                "region {} end key does not advance scan progress",
                region_id
            )));
        }

        self.current_key = region_end;
        Ok(Some(region_info))
    }

    pub fn checkpoint(&self) -> Vec<u8> {
        self.current_key.clone()
    }

    pub fn rewind(&mut self, checkpoint: Vec<u8>) {
        self.current_key = checkpoint;
        self.finished = false;
    }
}

async fn wait_sync_region_retry_delay(
    cancel_rx: &watch::Receiver<bool>,
) -> Result<(), EncryptionConfChangeError> {
    if *cancel_rx.borrow() {
        return Err(EncryptionConfChangeError::Canceled);
    }

    let mut cancel_rx = cancel_rx.clone();
    tokio::select! {
        _ = tokio::time::sleep(DEFAULT_SYNC_REGION_RETRY_DELAY) => Ok(()),
        changed = cancel_rx.changed() => match changed {
            Ok(()) if *cancel_rx.borrow() => Err(EncryptionConfChangeError::Canceled),
            Ok(()) | Err(_) => Ok(()),
        },
    }
}

async fn wait_job_finished(
    mut status_rx: watch::Receiver<JobStatus>,
) -> Result<RolloutJobSummary, EncryptionConfChangeError> {
    loop {
        match status_rx.borrow().clone() {
            JobStatus::Running => {}
            JobStatus::Finished(res) => return res,
        }

        if status_rx.changed().await.is_err() {
            return Err(EncryptionConfChangeError::failed(
                "job status channel closed unexpectedly".to_owned(),
            ));
        }
    }
}

#[cfg(test)]
mod tests {
    use pd_client::PdClient;
    use security::GetSecurityManager;

    use super::*;

    #[derive(Default)]
    struct DummyPdClient;

    impl GetSecurityManager for DummyPdClient {}

    impl PdClient for DummyPdClient {}

    struct KeyspaceMetaPdClient {
        keyspace_id: u32,
        current_file_id: u64,
        next_file_id: u64,
        encryption_config: Option<KeyspaceEncryptionConfig>,
        owner_store_id: Option<u64>,
    }

    impl GetSecurityManager for KeyspaceMetaPdClient {}

    impl PdClient for KeyspaceMetaPdClient {
        fn load_keyspace_by_id(
            &self,
            keyspace_id: u32,
        ) -> pd_client::Result<kvproto::keyspacepb::KeyspaceMeta> {
            if keyspace_id != self.keyspace_id {
                return Err(pd_client::Error::KeyspaceNotFound(keyspace_id));
            }
            let mut meta = kvproto::keyspacepb::KeyspaceMeta::default();
            meta.set_id(self.keyspace_id);
            meta.set_name(format!("test_{}", self.keyspace_id));
            meta.mut_config().insert(
                KEYSPACE_CONFIG_ENCRYPTION_META_CURRENT_FILE_ID.to_owned(),
                self.current_file_id.to_string(),
            );
            meta.mut_config().insert(
                KEYSPACE_CONFIG_ENCRYPTION_META_NEXT_FILE_ID.to_owned(),
                self.next_file_id.to_string(),
            );
            if let Some(encryption_config) = self.encryption_config.as_ref() {
                meta.mut_config().insert(
                    KEYSPACE_CONFIG_ENCRYPTION.to_owned(),
                    encryption_config.to_pd_config_value().unwrap(),
                );
            }
            if let Some(owner_store_id) = self.owner_store_id {
                meta.mut_config().insert(
                    KEYSPACE_CONFIG_ENCRYPTION_META_JOB_OWNER_STORE_ID.to_owned(),
                    owner_store_id.to_string(),
                );
            }
            Ok(meta)
        }
    }

    fn runtime() -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_current_thread()
            .enable_time()
            .build()
            .unwrap()
    }

    #[test]
    fn test_load_keyspace_encryption_status_classifies_missing_keyspace() {
        let pd_client: Arc<dyn PdClient> = Arc::new(KeyspaceMetaPdClient {
            keyspace_id: 9,
            current_file_id: 0,
            next_file_id: 0,
            encryption_config: None,
            owner_store_id: None,
        });
        let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

        let err = mgr.load_keyspace_encryption_status_from_pd(10).unwrap_err();
        assert!(err.is_keyspace_not_found(), "unexpected error: {}", err);
    }

    #[test]
    fn test_select_meta_id() {
        use KeyspaceEncryptionMetaId::*;

        let cases = [
            (false, 0, 0, false, Unencrypted),
            (false, 10, 20, true, Unencrypted),
            (true, 0, 0, false, UnavailableInPd),
            (true, 0, 0, true, UnavailableInPd),
            (true, 0, 20, false, Unencrypted),
            (true, 0, 20, true, MetaFileId(20)),
            (true, 10, 0, false, MetaFileId(10)),
            (true, 10, 0, true, MetaFileId(10)),
            (true, 10, 20, false, MetaFileId(10)),
            (true, 10, 20, true, MetaFileId(20)),
        ];

        for (enabled, current_file_id, next_file_id, latest_unstable, expected) in cases {
            let status = PdKeyspaceEncryptionStatus {
                keyspace_id: 1,
                keyspace_name: "test".to_owned(),
                encryption_config: enabled.then(|| KeyspaceEncryptionConfig {
                    enabled: true,
                    ..Default::default()
                }),
                meta_state: PdEncryptionMetaState {
                    current_file_id,
                    next_file_id,
                },
                encryption_meta_job_owner_store_id: None,
            };
            assert_eq!(status.select_meta_id(latest_unstable), expected);
        }
    }

    #[test]
    fn test_owner_check_enforced_rejects_wrong_owner() {
        let pd_client: Arc<dyn PdClient> = Arc::new(DummyPdClient);
        let mgr = EncryptionConfChangeManager::new(
            pd_client,
            OwnerCheckMode::Enforced { owner_store_id: 7 },
        );
        let status = PdKeyspaceEncryptionStatus {
            keyspace_id: 7,
            keyspace_name: "test".to_owned(),
            encryption_config: None,
            meta_state: PdEncryptionMetaState {
                current_file_id: 10,
                next_file_id: 11,
            },
            encryption_meta_job_owner_store_id: Some(42),
        };

        match mgr.check_owner(&status) {
            Ok(_) => panic!("expected WrongOwner error"),
            Err(EncryptionConfChangeError::WrongOwner {
                expected_owner_store_id,
                pd_owner_store_id,
            }) => {
                assert_eq!(expected_owner_store_id, 7);
                assert_eq!(pd_owner_store_id, 42);
            }
            Err(e) => panic!("unexpected error: {}", e),
        }
    }

    #[test]
    fn test_handle_gen_meta_rejects_pending_rotate() {
        let rt = runtime();
        rt.block_on(async {
            let pd_client: Arc<dyn PdClient> = Arc::new(KeyspaceMetaPdClient {
                keyspace_id: 9,
                current_file_id: 100,
                next_file_id: 200,
                encryption_config: Some(KeyspaceEncryptionConfig::enabled(
                    Some("cmek-a".to_owned()),
                    Some("test".to_owned()),
                    Some("us-west-2".to_owned()),
                    None,
                    BTreeMap::new(),
                )),
                owner_store_id: Some(7),
            });
            let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

            let err = mgr
                .handle_gen_meta(GenMetaRequest {
                    keyspace_id: 9,
                    mode: EncryptionMetaGenMode::RotateDataKey {
                        max_data_key_age: None,
                    },
                })
                .await
                .unwrap_err();

            assert!(
                err.to_string()
                    .contains("already has a pending conf change")
                    && err.to_string().contains("next_file_id=200"),
                "unexpected error: {}",
                err
            );
        });
    }

    #[test]
    fn test_handle_gen_meta_checks_pending_enable() {
        let rt = runtime();
        rt.block_on(async {
            let pd_client: Arc<dyn PdClient> = Arc::new(KeyspaceMetaPdClient {
                keyspace_id: 11,
                current_file_id: 0,
                next_file_id: 300,
                encryption_config: Some(KeyspaceEncryptionConfig::enabled(
                    Some("cmek-a".to_owned()),
                    Some("test".to_owned()),
                    Some("us-west-2".to_owned()),
                    Some("kms-endpoint".to_owned()),
                    BTreeMap::new(),
                )),
                owner_store_id: Some(5),
            });
            let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

            let result = mgr
                .handle_gen_meta(GenMetaRequest {
                    keyspace_id: 11,
                    mode: EncryptionMetaGenMode::EnableCmek {
                        cmek_id: "cmek-a".to_owned(),
                        vendor: "test".to_owned(),
                        region: Some("us-west-2".to_owned()),
                        endpoint: Some("kms-endpoint".to_owned()),
                        encryption_context: BTreeMap::new(),
                    },
                })
                .await
                .unwrap();

            assert_eq!(
                result.pd_before,
                PdEncryptionMetaState {
                    current_file_id: 0,
                    next_file_id: 300,
                }
            );
            assert_eq!(result.pd_after, result.pd_before);
            assert_eq!(result.gen_meta.file_id, 300);
            assert_eq!(result.gen_meta.new_data_key_generated, Some(true));
            assert_eq!(result.owner_store_id, Some(5));
            assert!(!result.pd_updated);

            let pd_client: Arc<dyn PdClient> = Arc::new(KeyspaceMetaPdClient {
                keyspace_id: 13,
                current_file_id: 0,
                next_file_id: 400,
                encryption_config: Some(KeyspaceEncryptionConfig::enabled(
                    Some("cmek-a".to_owned()),
                    Some("test".to_owned()),
                    Some("us-west-2".to_owned()),
                    None,
                    BTreeMap::new(),
                )),
                owner_store_id: Some(8),
            });
            let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

            let err = mgr
                .handle_gen_meta(GenMetaRequest {
                    keyspace_id: 13,
                    mode: EncryptionMetaGenMode::EnableCmek {
                        cmek_id: "cmek-b".to_owned(),
                        vendor: "test".to_owned(),
                        region: Some("us-west-2".to_owned()),
                        endpoint: None,
                        encryption_context: BTreeMap::new(),
                    },
                })
                .await
                .unwrap_err();

            assert!(
                err.to_string()
                    .contains("already has a pending conf change")
                    && err.to_string().contains("next_file_id=400"),
                "unexpected error: {}",
                err
            );
        });
    }

    #[test]
    fn test_remove_job_if_match_only_removes_same_job() {
        let pd_client: Arc<dyn PdClient> = Arc::new(DummyPdClient);
        let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

        let (cancel_tx, _cancel_rx) = watch::channel(false);
        let (_status_tx, status_rx) = watch::channel::<JobStatus>(JobStatus::Running);
        {
            let mut jobs = mgr.inner.jobs.lock().unwrap();
            jobs.insert(
                1,
                RunningJob {
                    job_id: 2,
                    is_stuck: true,
                    cancel_tx,
                    status_rx,
                },
            );
        }

        // Wrong job id should not remove.
        mgr.remove_job_if_match(1, 1);
        {
            let jobs = mgr.inner.jobs.lock().unwrap();
            assert!(jobs.contains_key(&1));
            assert_eq!(jobs.get(&1).unwrap().job_id, 2);
            assert!(jobs.get(&1).unwrap().is_stuck);
        }

        // Correct job id should remove.
        mgr.remove_job_if_match(1, 2);
        {
            let jobs = mgr.inner.jobs.lock().unwrap();
            assert!(!jobs.contains_key(&1));
        }
    }

    #[test]
    fn test_mark_job_stuck_if_needed_ignores_replaced_job() {
        let pd_client: Arc<dyn PdClient> = Arc::new(DummyPdClient);
        let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

        let (cancel_tx, _cancel_rx) = watch::channel(false);
        let (_status_tx, status_rx) = watch::channel::<JobStatus>(JobStatus::Running);
        {
            let mut jobs = mgr.inner.jobs.lock().unwrap();
            jobs.insert(
                1,
                RunningJob {
                    job_id: 2,
                    is_stuck: false,
                    cancel_tx,
                    status_rx,
                },
            );
        }

        mgr.mark_job_stuck_if_needed(1, 1, Instant::now() - DEFAULT_STUCK_JOB_THRESHOLD);

        let jobs = mgr.inner.jobs.lock().unwrap();
        let job = jobs.get(&1).unwrap();
        assert_eq!(job.job_id, 2);
        assert!(!job.is_stuck);
    }

    #[test]
    fn test_handle_cancel_rollout_job_does_not_remove_new_job() {
        let rt = runtime();
        rt.block_on(async {
            let pd_client: Arc<dyn PdClient> = Arc::new(DummyPdClient);
            let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

            let keyspace_id = 42;
            let (cancel_tx, mut cancel_rx) = watch::channel(false);
            let (status_tx, status_rx) = watch::channel::<JobStatus>(JobStatus::Running);
            {
                let mut jobs = mgr.inner.jobs.lock().unwrap();
                jobs.insert(
                    keyspace_id,
                    RunningJob {
                        job_id: 1,
                        is_stuck: false,
                        cancel_tx,
                        status_rx,
                    },
                );
            }

            let mgr2 = mgr.clone();
            let cancel_task = tokio::spawn(async move {
                mgr2.handle_cancel_rollout_job(CancelRolloutJobRequest { keyspace_id })
                    .await
                    .unwrap()
            });

            // Wait until cancel signal is observed, ensuring the cancel path is in
            // progress.
            loop {
                if *cancel_rx.borrow() {
                    break;
                }
                cancel_rx.changed().await.unwrap();
            }

            // Simulate a new job starting before the cancel handler removes from the map.
            let (cancel_tx2, _cancel_rx2) = watch::channel(false);
            let (_status_tx2, status_rx2) = watch::channel::<JobStatus>(JobStatus::Running);
            {
                let mut jobs = mgr.inner.jobs.lock().unwrap();
                jobs.insert(
                    keyspace_id,
                    RunningJob {
                        job_id: 2,
                        is_stuck: false,
                        cancel_tx: cancel_tx2,
                        status_rx: status_rx2,
                    },
                );
            }

            let _ = status_tx.send(JobStatus::Finished(Err(
                EncryptionConfChangeError::Canceled,
            )));
            let _ = cancel_task.await.unwrap();

            let jobs = mgr.inner.jobs.lock().unwrap();
            assert!(jobs.contains_key(&keyspace_id));
            assert_eq!(jobs.get(&keyspace_id).unwrap().job_id, 2);
        });
    }

    #[test]
    fn test_keyspace_region_scanner_rewind_restores_scan_key_and_finished_state() {
        let mut scanner = KeyspaceRegionScanner::for_keyspace(9);
        let checkpoint = scanner.checkpoint();

        scanner.current_key = b"advanced".to_vec();
        scanner.finished = true;
        scanner.rewind(checkpoint.clone());

        assert_eq!(scanner.current_key, checkpoint);
        assert!(!scanner.finished);
    }

    #[test]
    fn test_handle_start_rollout_job_returns_already_running() {
        let rt = runtime();
        rt.block_on(async {
            let pd_client: Arc<dyn PdClient> = Arc::new(KeyspaceMetaPdClient {
                keyspace_id: 7,
                current_file_id: 100,
                next_file_id: 200,
                encryption_config: None,
                owner_store_id: None,
            });
            let mgr = EncryptionConfChangeManager::new(pd_client, OwnerCheckMode::Disabled);

            let (cancel_tx, _cancel_rx) = watch::channel(false);
            let (_status_tx, status_rx) = watch::channel::<JobStatus>(JobStatus::Running);
            {
                let mut jobs = mgr.inner.jobs.lock().unwrap();
                jobs.insert(
                    7,
                    RunningJob {
                        job_id: 999,
                        is_stuck: false,
                        cancel_tx,
                        status_rx,
                    },
                );
            }

            let result = mgr
                .handle_start_rollout_job(StartRolloutJobRequest { keyspace_id: 7 })
                .await
                .unwrap();
            assert!(!result.response.started);
            assert!(result.response.running);
            assert_eq!(result.response.note.as_deref(), Some("already_running"));
            assert!(result.handle.is_none());
        });
    }

    #[test]
    fn test_wait_job_finished_returns_result() {
        let rt = runtime();
        rt.block_on(async {
            let (status_tx, status_rx) = watch::channel::<JobStatus>(JobStatus::Running);
            let wait_task = tokio::spawn(async move { wait_job_finished(status_rx).await });

            let summary = RolloutJobSummary {
                regions_success: 1,
                regions_failed: 2,
            };
            let _ = status_tx.send(JobStatus::Finished(Ok(summary)));

            let res = wait_task.await.unwrap().unwrap();
            assert_eq!(res.regions_success, 1);
            assert_eq!(res.regions_failed, 2);
        });
    }
}