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

use std::{
    collections::BTreeMap,
    path::{Path, PathBuf},
};

use api_version::ApiV2;
use chrono::{DateTime, NaiveDateTime, Utc};
use clap::{Args, Subcommand, ValueEnum};
use cloud_encryption::KeyspaceEncryptionConfig;
use encryption_conf_change::{
    EncryptionConfChangeManager, EncryptionMetaGenMode, GenMetaRequest, OwnerCheckMode,
    PdKeyspaceEncryptionStatus, StartRolloutJobRequest,
};
use hyper::Body;
use kvengine::{ENCRYPTION_META_KEY, get_shard_property};
use kvenginepb::EncryptionMeta;
use kvproto::{
    keyspacepb::KeyspaceMeta,
    metapb::{Peer, Region, Store},
};
use pd_client::PdClient;
use protobuf::Message;
use txn_types::Key;

use crate::{
    common::{CommonConfig, CtlContext},
    http::build_uri,
};

const DEFAULT_PD_ENDPOINT: &str = "https://db-pd:2379";

fn resolve_pd(config: &Path, pd: &str) -> String {
    if !pd.is_empty() {
        return pd.to_owned();
    }
    // If user explicitly provided a config path, do not override its PD config.
    if !config.as_os_str().is_empty() {
        return String::new();
    }
    DEFAULT_PD_ENDPOINT.to_owned()
}

fn build_ctx(config: &Path, pd: &str) -> CtlContext {
    let pd = resolve_pd(config, pd);
    let config = CommonConfig::from_args(config, &pd);
    config.create_context()
}

#[derive(Args)]
pub struct EncryptionCommand {
    #[clap(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Show encryption status for a region or keyspace.
    Show(ShowEncryptionArgs),
    /// Generate the latest encryption meta and update next_file_id in PD.
    GenMeta(GenMetaArgs),
    /// Sync all keyspace regions to PD next_file_id and finalize
    /// current_file_id.
    SyncMeta(SyncMetaArgs),
    /// Show keyspace encryption meta current/next file IDs stored in PD.
    Status(StatusArgs),
}

pub fn execute_encryption_command(cmd: EncryptionCommand) {
    match cmd.command {
        Commands::Show(args) => execute_show_encryption(args),
        Commands::GenMeta(args) => execute_gen_meta(args),
        Commands::SyncMeta(args) => execute_sync_meta(args),
        Commands::Status(args) => execute_status(args),
    }
}

#[derive(Args)]
pub struct ShowEncryptionArgs {
    /// The path of the config file.
    #[clap(long, default_value = "")]
    pub config: PathBuf,
    /// PD endpoints, use `,` to separate multiple PDs.
    #[clap(long, default_value_t = String::new())]
    pub pd: String,
    /// Region ID to query encryption status.
    #[clap(long)]
    pub region_id: Option<u64>,
    #[clap(flatten)]
    keyspace: KeyspaceArgs,
    /// Show all regions in the keyspace (only used with keyspace selector) (may
    /// cost too much time).
    #[clap(long)]
    pub all: bool,
}

#[derive(Args, Clone, Debug, Default)]
struct KeyspaceArgs {
    /// Keyspace ID.
    #[clap(long)]
    pub keyspace_id: Option<u32>,
    /// Keyspace name.
    #[clap(long)]
    pub keyspace_name: Option<String>,
}

#[derive(Args)]
pub struct GenMetaArgs {
    #[clap(long, default_value = "")]
    pub config: PathBuf,
    /// PD endpoints, use `,` to separate multiple PDs
    #[clap(long, default_value_t = String::new())]
    pub pd: String,
    #[clap(flatten)]
    keyspace: KeyspaceArgs,
    /// Generation mode.
    #[clap(long, value_enum, default_value_t = GenMetaMode::RotateDataKey)]
    mode: GenMetaMode,
    /// Maximum age of data key in seconds before generating a new data key.
    /// Only valid for `rotate-data-key` mode.
    #[clap(long)]
    max_data_key_age: Option<u64>,
    /// CMEK ID. Required for `enable-cmek` mode.
    #[clap(long)]
    cmek_id: Option<String>,
    /// KMS vendor. Required for `enable-cmek` mode.
    #[clap(long)]
    vendor: Option<String>,
    /// KMS region. Optional for `enable-cmek` mode.
    #[clap(long)]
    region: Option<String>,
    /// KMS endpoint. Optional for `enable-cmek` mode.
    #[clap(long)]
    endpoint: Option<String>,
    /// KMS encryption context as a JSON object of string->string, for example:
    /// `{"ExternalId":"customer-42","tenant":"acme"}`.
    ///
    /// The JSON string length limit is 8192 characters.
    ///
    /// Note: the CLI flag value itself is a string. In bash/zsh, the easiest
    /// way to pass a JSON object is to wrap it with single quotes, e.g.
    /// `--encryption-context '{"ExternalId":"customer-42","tenant":"acme"}'`.
    ///
    /// If you are embedding the encryption config JSON inside another JSON
    /// payload (e.g. PD create-keyspace API where `config["encryption"]` is a
    /// JSON string), the inner quotes will be escaped by the outer JSON
    /// encoder. Optional for `enable-cmek` mode.
    #[clap(long)]
    encryption_context: Option<String>,
}

#[derive(Args)]
pub struct SyncMetaArgs {
    #[clap(long, default_value = "")]
    pub config: PathBuf,
    /// PD endpoints, use `,` to separate multiple PDs
    #[clap(long, default_value_t = String::new())]
    pub pd: String,
    #[clap(flatten)]
    keyspace: KeyspaceArgs,
    /// Sync mode.
    #[clap(long, value_enum, default_value = "local")]
    mode: SyncMetaMode,
    /// Verbose output.
    #[clap(long)]
    verbose: bool,
}

#[derive(Args)]
pub struct StatusArgs {
    #[clap(long, default_value = "")]
    pub config: PathBuf,
    /// PD endpoints, use `,` to separate multiple PDs
    #[clap(long, default_value_t = String::new())]
    pub pd: String,
    #[clap(flatten)]
    keyspace: KeyspaceArgs,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)]
enum SyncMetaMode {
    Local,
    Remote,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)]
enum GenMetaMode {
    RotateDataKey,
    EnableCmek,
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum KeyspaceSelector {
    Id(u32),
    Name(String),
    IdAndName { id: u32, name: String },
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ResolvedKeyspace {
    id: u32,
    name: String,
}

impl KeyspaceArgs {
    fn parse_selector(&self) -> Result<Option<KeyspaceSelector>, String> {
        let keyspace_name = match &self.keyspace_name {
            Some(name) => {
                let trimmed = name.trim();
                if trimmed.is_empty() {
                    return Err("--keyspace-name must not be empty".to_owned());
                }
                Some(trimmed.to_owned())
            }
            None => None,
        };

        Ok(match (self.keyspace_id, keyspace_name) {
            (Some(id), Some(name)) => Some(KeyspaceSelector::IdAndName { id, name }),
            (Some(id), None) => Some(KeyspaceSelector::Id(id)),
            (None, Some(name)) => Some(KeyspaceSelector::Name(name)),
            (None, None) => None,
        })
    }

    fn parse_required_selector(&self) -> Result<KeyspaceSelector, String> {
        self.parse_selector()?
            .ok_or_else(|| "Either --keyspace-id or --keyspace-name must be specified.".to_owned())
    }
}

fn ensure_non_system_keyspace(keyspace_id: u32) {
    if keyspace_id == 0 {
        exit_with_error("SYSTEM keyspace (keyspace_id=0) does not support encryption.");
    }
}

fn load_keyspace_by_name(ctx: &CtlContext, keyspace_name: &str) -> KeyspaceMeta {
    let meta = match ctx.pd_client.load_keyspace(keyspace_name.to_owned()) {
        Ok(meta) => meta,
        Err(e) => {
            exit_with_error(&format!(
                "Failed to load keyspace '{}' from PD: {}",
                keyspace_name, e
            ));
        }
    };
    ensure_non_system_keyspace(meta.get_id());
    meta
}

fn load_keyspace_by_id(ctx: &CtlContext, keyspace_id: u32) -> KeyspaceMeta {
    ensure_non_system_keyspace(keyspace_id);

    let meta = match ctx.pd_client.load_keyspace_by_id(keyspace_id) {
        Ok(meta) => meta,
        Err(e) => {
            exit_with_error(&format!(
                "Failed to load keyspace id {} from PD: {}",
                keyspace_id, e
            ));
        }
    };
    meta
}

fn resolved_keyspace_from_meta(meta: KeyspaceMeta) -> ResolvedKeyspace {
    let keyspace_id = meta.get_id();
    ensure_non_system_keyspace(keyspace_id);
    ResolvedKeyspace {
        id: keyspace_id,
        name: meta.get_name().to_owned(),
    }
}

fn resolve_keyspace_selector(ctx: &CtlContext, selector: KeyspaceSelector) -> ResolvedKeyspace {
    match selector {
        KeyspaceSelector::Id(keyspace_id) => {
            resolved_keyspace_from_meta(load_keyspace_by_id(ctx, keyspace_id))
        }
        KeyspaceSelector::Name(keyspace_name) => {
            resolved_keyspace_from_meta(load_keyspace_by_name(ctx, &keyspace_name))
        }
        KeyspaceSelector::IdAndName { id, name } => {
            let resolved = resolved_keyspace_from_meta(load_keyspace_by_name(ctx, &name));
            if resolved.id != id {
                exit_with_error(&format!(
                    "--keyspace-id {} does not match --keyspace-name '{}' (resolved id={}).",
                    id, name, resolved.id
                ));
            }
            resolved
        }
    }
}

fn resolve_optional_keyspace(
    ctx: &CtlContext,
    keyspace: &KeyspaceArgs,
) -> Option<ResolvedKeyspace> {
    let selector = match keyspace.parse_selector() {
        Ok(selector) => selector,
        Err(e) => exit_with_error(&e),
    };
    selector.map(|selector| resolve_keyspace_selector(ctx, selector))
}

fn resolve_required_keyspace(ctx: &CtlContext, keyspace: &KeyspaceArgs) -> ResolvedKeyspace {
    let selector = match keyspace.parse_required_selector() {
        Ok(selector) => selector,
        Err(e) => exit_with_error(&e),
    };
    resolve_keyspace_selector(ctx, selector)
}

fn execute_show_encryption(args: ShowEncryptionArgs) {
    let ctx = build_ctx(&args.config, &args.pd);
    let runtime = ctx.dfs.get_runtime();
    let keyspace = resolve_optional_keyspace(&ctx, &args.keyspace);

    match (args.region_id, keyspace) {
        (Some(region_id), keyspace) => {
            show_region_encryption(
                &ctx,
                runtime,
                region_id,
                keyspace.map(|keyspace| keyspace.id),
            );
        }
        (None, Some(keyspace)) => {
            let keyspace_id = keyspace.id;
            if args.all {
                show_keyspace_all_regions_encryption(&ctx, runtime, keyspace_id);
            } else {
                show_keyspace_first_region_encryption(&ctx, runtime, keyspace_id);
            }
        }
        (None, None) => {
            print_error("Either --region-id, --keyspace-id, or --keyspace-name must be specified.");
            std::process::exit(1);
        }
    }
}

fn show_region_encryption(
    ctx: &CtlContext,
    runtime: &tokio::runtime::Runtime,
    region_id: u64,
    expected_keyspace_id: Option<u32>,
) {
    let (region, leader) = match get_region_leader(ctx, runtime, region_id) {
        Some(result) => result,
        None => {
            print_error(&format!("Region {} not found or has no leader.", region_id));
            return;
        }
    };

    let store = match ctx.pd_client.get_store(leader.get_store_id()) {
        Ok(store) => store,
        Err(e) => {
            print_error(&format!(
                "Failed to get store {}: {}",
                leader.get_store_id(),
                e
            ));
            return;
        }
    };

    let encryption_meta = match runtime.block_on(fetch_encryption_meta(ctx, &store, region_id)) {
        Ok(meta) => meta,
        Err(e) => {
            print_error(&e);
            return;
        }
    };
    let keyspace_id = extract_keyspace_id(&region);

    // Validate keyspace_id if expected
    if let Some(expected_ks) = expected_keyspace_id {
        if keyspace_id != Some(expected_ks) {
            print_error(&format!("Keyspace {} not found.", expected_ks));
            return;
        }
    }

    print_json_format(region_id, keyspace_id, &leader, &store, encryption_meta);
}

fn show_keyspace_first_region_encryption(
    ctx: &CtlContext,
    runtime: &tokio::runtime::Runtime,
    keyspace_id: u32,
) {
    let ks_start_key = Key::from_raw(&ApiV2::get_keyspace_prefix_by_id(keyspace_id));
    let region = match ctx.pd_client.get_region(ks_start_key.as_encoded()) {
        Ok(region) => region,
        Err(e) => {
            print_error(&format!(
                "Failed to get region for keyspace {}: {}",
                keyspace_id, e
            ));
            return;
        }
    };

    // Validate that the region belongs to the requested keyspace
    let actual_keyspace_id = extract_keyspace_id(&region);
    if actual_keyspace_id != Some(keyspace_id) {
        print_error(&format!("Keyspace {} not found.", keyspace_id));
        return;
    }

    show_region_encryption(ctx, runtime, region.get_id(), Some(keyspace_id));
}

fn show_keyspace_all_regions_encryption(
    ctx: &CtlContext,
    runtime: &tokio::runtime::Runtime,
    keyspace_id: u32,
) {
    let ks_start_key = Key::from_raw(&ApiV2::get_keyspace_prefix_by_id(keyspace_id));
    let next_keyspace_id = match keyspace_id.checked_add(1) {
        Some(id) => id,
        None => {
            print_error(&format!(
                "Invalid keyspace_id {}, cannot compute end range",
                keyspace_id
            ));
            return;
        }
    };
    let ks_end_key = Key::from_raw(&ApiV2::get_keyspace_prefix_by_id(next_keyspace_id));
    let ks_end_encoded = ks_end_key.as_encoded().as_slice();

    let mut current_key = ks_start_key.as_encoded().clone();
    let mut region_count = 0;
    let mut results: Vec<serde_json::Value> = Vec::new();

    while current_key.as_slice() < ks_end_encoded {
        let region = match ctx.pd_client.get_region(&current_key) {
            Ok(region) => region,
            Err(e) => {
                print_error(&format!("Failed to get region: {}", e));
                break;
            }
        };

        // Validate that the region belongs to the requested keyspace
        let actual_keyspace_id = extract_keyspace_id(&region);
        if actual_keyspace_id != Some(keyspace_id) {
            break;
        }

        let region_id = region.get_id();
        let end_key = region.get_end_key().to_vec();

        if let Some(json) = get_region_encryption_json(ctx, runtime, region_id, Some(keyspace_id)) {
            results.push(json);
            region_count += 1;
        }

        if end_key.is_empty() || end_key.as_slice() >= ks_end_encoded {
            break;
        }
        current_key = end_key;
    }

    if region_count == 0 {
        print_error(&format!("Keyspace {} not found.", keyspace_id));
        return;
    }

    let output = serde_json::json!({
        "keyspace_id": keyspace_id,
        "total_regions": region_count,
        "regions": results,
    });
    println!("{}", serde_json::to_string_pretty(&output).unwrap());
}

fn get_region_encryption_json(
    ctx: &CtlContext,
    runtime: &tokio::runtime::Runtime,
    region_id: u64,
    expected_keyspace_id: Option<u32>,
) -> Option<serde_json::Value> {
    let (region, leader) = get_region_leader(ctx, runtime, region_id)?;

    let store = ctx.pd_client.get_store(leader.get_store_id()).ok()?;

    let encryption_meta = match runtime.block_on(fetch_encryption_meta(ctx, &store, region_id)) {
        Ok(meta) => meta,
        Err(e) => {
            eprintln!("Warning: {}", e);
            return None;
        }
    };
    let keyspace_id = extract_keyspace_id(&region);

    // Validate keyspace_id if expected
    if let Some(expected_ks) = expected_keyspace_id {
        if keyspace_id != Some(expected_ks) {
            eprintln!(
                "Warning: region {} does not belong to keyspace {}.",
                region_id, expected_ks
            );
            return None;
        }
    }

    Some(build_json_output(
        region_id,
        keyspace_id,
        &leader,
        &store,
        encryption_meta,
    ))
}

fn get_region_leader(
    ctx: &CtlContext,
    runtime: &tokio::runtime::Runtime,
    region_id: u64,
) -> Option<(Region, Peer)> {
    runtime
        .block_on(ctx.pd_client.get_region_leader_by_id(region_id))
        .ok()
        .flatten()
}

async fn fetch_encryption_meta(
    ctx: &CtlContext,
    store: &Store,
    region_id: u64,
) -> Result<Option<EncryptionMeta>, String> {
    let uri = build_uri(ctx, store, &format!("kvengine/meta/{}", region_id));
    let req = hyper::Request::get(uri)
        .body(Body::empty())
        .map_err(|e| format!("Failed to build request: {}", e))?;

    let resp = ctx.http_client.request(req).await.map_err(|e| {
        format!(
            "Failed to fetch encryption meta for region {}: {}",
            region_id, e
        )
    })?;

    let status = resp.status();
    let body = hyper::body::to_bytes(resp.into_body()).await.map_err(|e| {
        format!(
            "Failed to read response body for region {}: {}",
            region_id, e
        )
    })?;

    if !status.is_success() {
        return Err(format!(
            "Failed to fetch encryption meta for region {}: HTTP {}",
            region_id, status
        ));
    }

    if body.is_empty() {
        return Ok(None);
    }

    let mut changeset = kvenginepb::ChangeSet::default();
    if changeset.merge_from_bytes(&body).is_err() {
        return Err(format!(
            "Failed to parse ChangeSet for region {}",
            region_id
        ));
    }

    let snapshot = changeset.get_snapshot();
    let meta =
        get_shard_property(ENCRYPTION_META_KEY, snapshot.get_properties()).and_then(|bytes| {
            let mut meta = EncryptionMeta::default();
            meta.merge_from_bytes(&bytes).ok()?;
            Some(meta)
        });
    Ok(meta)
}

fn extract_keyspace_id(region: &Region) -> Option<u32> {
    let start_key = region.get_start_key();
    if start_key.len() >= 4 {
        let raw_key = Key::from_encoded_slice(start_key).into_raw().ok()?;
        if raw_key.len() >= 4 {
            return ApiV2::get_u32_keyspace_id_by_key(&raw_key);
        }
    }
    None
}

fn print_error(message: &str) {
    let error_json = serde_json::json!({
        "error": message
    });
    eprintln!("{}", serde_json::to_string_pretty(&error_json).unwrap());
}

fn exit_with_error(message: &str) -> ! {
    print_error(message);
    std::process::exit(1);
}

fn print_json_format(
    region_id: u64,
    keyspace_id: Option<u32>,
    leader: &Peer,
    store: &Store,
    encryption_meta: Option<EncryptionMeta>,
) {
    let output = build_json_output(region_id, keyspace_id, leader, store, encryption_meta);
    println!("{}", serde_json::to_string_pretty(&output).unwrap());
}

fn print_json(output: &serde_json::Value) {
    println!("{}", serde_json::to_string_pretty(output).unwrap());
}

fn build_json_output(
    region_id: u64,
    keyspace_id: Option<u32>,
    leader: &Peer,
    store: &Store,
    encryption_meta: Option<EncryptionMeta>,
) -> serde_json::Value {
    let mut output = serde_json::json!({
        "region_id": region_id,
        "leader_store_id": leader.get_store_id(),
        "leader_address": store.get_address(),
    });

    if let Some(ks_id) = keyspace_id {
        output["keyspace_id"] = serde_json::json!(ks_id);
    }

    match encryption_meta {
        Some(meta) => {
            output["encrypted"] = serde_json::json!(true);

            if meta.has_master_key() {
                let master_key = meta.get_master_key();
                output["master_key"] = serde_json::json!({
                    "vendor": master_key.get_vendor(),
                    "cmek_id": master_key.get_cmek_id(),
                    "region": master_key.get_region(),
                    "endpoint": master_key.get_endpoint(),
                });
            }

            if meta.has_current() {
                let current = meta.get_current();
                let created_at = current.get_created_at();
                output["current_epoch"] = serde_json::json!({
                    "file_id": current.get_file_id(),
                    "data_key_id": current.get_data_key_id(),
                    "created_at": created_at,
                    "created_at_rfc3339": unix_secs_to_rfc3339(created_at),
                });
            }

            let data_keys = meta.get_data_keys();
            output["data_keys"] = serde_json::json!({
                "count": data_keys.len(),
                "key_ids": data_keys.keys().copied().collect::<Vec<u32>>(),
            });

            let history = meta.get_history();
            if !history.is_empty() {
                let history_json: Vec<serde_json::Value> = history
                    .iter()
                    .map(|epoch| {
                        let created_at = epoch.get_created_at();
                        serde_json::json!({
                            "file_id": epoch.get_file_id(),
                            "data_key_id": epoch.get_data_key_id(),
                            "created_at": created_at,
                            "created_at_rfc3339": unix_secs_to_rfc3339(created_at),
                        })
                    })
                    .collect();
                output["history"] = serde_json::json!(history_json);
            }
        }
        None => {
            output["encrypted"] = serde_json::json!(false);
        }
    }

    output
}

fn unix_secs_to_rfc3339(unix_secs: u64) -> Option<String> {
    if unix_secs == 0 {
        return None;
    }
    let unix_secs: i64 = unix_secs.try_into().ok()?;
    let dt = NaiveDateTime::from_timestamp_opt(unix_secs, 0)?;
    Some(DateTime::<Utc>::from_naive_utc_and_offset(dt, Utc).to_rfc3339())
}

fn build_status_output(status: PdKeyspaceEncryptionStatus) -> serde_json::Value {
    let encryption_enabled = status.encryption_config.as_ref().is_some_and(|c| c.enabled);

    serde_json::json!({
        "keyspace_id": status.keyspace_id,
        "keyspace_name": status.keyspace_name,
        "encryption_enabled": encryption_enabled,
        "meta_state": {
            "current_file_id": status.meta_state.current_file_id,
            "next_file_id": status.meta_state.next_file_id,
        },
        "encryption_meta_job_owner_store_id": status.encryption_meta_job_owner_store_id,
        "encryption_config": status.encryption_config,
    })
}

fn build_gen_meta_output(
    mode: &EncryptionMetaGenMode,
    result: encryption_conf_change::GenMetaResult,
) -> serde_json::Value {
    let mut output = serde_json::json!({
        "keyspace_id": result.keyspace_id,
        "mode": gen_meta_mode_name(mode),
        "generated": {
            "file_id": result.gen_meta.file_id,
            "data_key_id": result.gen_meta.data_key_id,
            "new_data_key_generated": result.gen_meta.new_data_key_generated,
        },
        "pd_updated": result.pd_updated,
        "pd_after": {
            "current_file_id": result.pd_after.current_file_id,
            "next_file_id": result.pd_after.next_file_id,
        },
        "owner_store_id": result.owner_store_id,
    });

    if let EncryptionMetaGenMode::EnableCmek {
        cmek_id,
        vendor,
        region,
        endpoint,
        encryption_context,
    } = mode
    {
        output["encryption_config"] = serde_json::json!({
            "enabled": true,
            "cmek_id": cmek_id,
            "vendor": vendor,
            "region": region,
            "endpoint": endpoint,
            "encryption_context": (!encryption_context.is_empty()).then_some(encryption_context),
        });
    }

    output
}

fn option_non_empty(name: &str, v: &Option<String>) -> Result<Option<String>, String> {
    match v {
        Some(s) => {
            let trimmed = s.trim();
            if trimmed.is_empty() {
                Err(format!("--{} must not be empty", name))
            } else {
                Ok(Some(trimmed.to_owned()))
            }
        }
        None => Ok(None),
    }
}

fn parse_encryption_context(value: &str) -> Result<BTreeMap<String, String>, String> {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return Err("--encryption-context must not be empty".to_owned());
    }

    if trimmed.len() > KeyspaceEncryptionConfig::MAX_ENCRYPTION_CONTEXT_JSON_LEN {
        return Err(format!(
            "--encryption-context JSON length must be <= {} chars, got {}",
            KeyspaceEncryptionConfig::MAX_ENCRYPTION_CONTEXT_JSON_LEN,
            trimmed.len()
        ));
    }

    let parsed = serde_json::from_str::<BTreeMap<String, String>>(trimmed).map_err(|err| {
        format!(
            "--encryption-context must be a JSON object with string keys and string values (for example, 'k1' maps to 'v1' and 'k2' maps to 'v2'); total length must be <= {} chars: {}",
            KeyspaceEncryptionConfig::MAX_ENCRYPTION_CONTEXT_JSON_LEN,
            err
        )
    })?;

    KeyspaceEncryptionConfig::validate_encryption_context(&parsed)
        .map_err(|err| format!("--encryption-context invalid: {}", err))?;

    Ok(parsed)
}

fn parse_encryption_context_arg(
    value: &Option<String>,
) -> Result<BTreeMap<String, String>, String> {
    let Some(value) = option_non_empty("encryption-context", value)? else {
        return Ok(BTreeMap::new());
    };
    parse_encryption_context(&value)
}

fn build_gen_meta_mode(args: &GenMetaArgs) -> Result<EncryptionMetaGenMode, String> {
    match args.mode {
        GenMetaMode::RotateDataKey => {
            let mut invalid_flags = Vec::new();
            if args.cmek_id.is_some() {
                invalid_flags.push("--cmek-id");
            }
            if args.vendor.is_some() {
                invalid_flags.push("--vendor");
            }
            if args.region.is_some() {
                invalid_flags.push("--region");
            }
            if args.endpoint.is_some() {
                invalid_flags.push("--endpoint");
            }
            if args.encryption_context.is_some() {
                invalid_flags.push("--encryption-context");
            }
            if !invalid_flags.is_empty() {
                return Err(format!(
                    "{} can only be used with --mode enable-cmek",
                    invalid_flags.join(", ")
                ));
            }
            Ok(EncryptionMetaGenMode::RotateDataKey {
                max_data_key_age: args.max_data_key_age,
            })
        }
        GenMetaMode::EnableCmek => {
            if args.max_data_key_age.is_some() {
                return Err(
                    "--max-data-key-age can only be used with --mode rotate-data-key".to_owned(),
                );
            }

            let cmek_id = option_non_empty("cmek-id", &args.cmek_id)?
                .ok_or_else(|| "--cmek-id is required when --mode enable-cmek".to_owned())?;
            let vendor = option_non_empty("vendor", &args.vendor)?
                .ok_or_else(|| "--vendor is required when --mode enable-cmek".to_owned())?;
            let region = option_non_empty("region", &args.region)?;
            let endpoint = option_non_empty("endpoint", &args.endpoint)?;
            let encryption_context = parse_encryption_context_arg(&args.encryption_context)?;

            Ok(EncryptionMetaGenMode::EnableCmek {
                cmek_id,
                vendor,
                region,
                endpoint,
                encryption_context,
            })
        }
    }
}

fn gen_meta_mode_name(mode: &EncryptionMetaGenMode) -> &'static str {
    match mode {
        EncryptionMetaGenMode::RotateDataKey { .. } => "rotate_data_key",
        EncryptionMetaGenMode::EnableCmek { .. } => "enable_cmek",
    }
}

fn get_keyspace_range(keyspace_id: u32) -> Result<(Key, Key), String> {
    let start_key = Key::from_raw(&ApiV2::get_keyspace_prefix_by_id(keyspace_id));
    let next_keyspace_id = keyspace_id.checked_add(1).ok_or_else(|| {
        format!(
            "Invalid keyspace_id {}, cannot compute end range",
            keyspace_id
        )
    })?;
    let end_key = Key::from_raw(&ApiV2::get_keyspace_prefix_by_id(next_keyspace_id));
    Ok((start_key, end_key))
}

fn execute_status(args: StatusArgs) {
    let ctx = build_ctx(&args.config, &args.pd);
    let keyspace = resolve_required_keyspace(&ctx, &args.keyspace);
    let keyspace_id = keyspace.id;

    let mgr = EncryptionConfChangeManager::new(ctx.pd_client.clone(), OwnerCheckMode::Disabled);
    let status = match mgr.load_keyspace_encryption_status_from_pd(keyspace_id) {
        Ok(status) => status,
        Err(e) => {
            exit_with_error(&format!(
                "Failed to load keyspace '{}' (id={}) from PD: {}",
                keyspace.name, keyspace_id, e
            ));
        }
    };

    print_json(&build_status_output(status));
}

fn execute_gen_meta(args: GenMetaArgs) {
    let ctx = build_ctx(&args.config, &args.pd);
    let runtime = ctx.dfs.get_runtime();

    let keyspace_id = resolve_required_keyspace(&ctx, &args.keyspace).id;
    let mode = match build_gen_meta_mode(&args) {
        Ok(mode) => mode,
        Err(e) => {
            exit_with_error(&format!("Invalid gen-meta arguments: {}", e));
        }
    };
    let mode_for_print = mode.clone();
    let mode_name = gen_meta_mode_name(&mode_for_print);

    let mgr = EncryptionConfChangeManager::new(ctx.pd_client.clone(), OwnerCheckMode::Disabled);
    let result = match runtime.block_on(async {
        mgr.handle_gen_meta(GenMetaRequest { keyspace_id, mode })
            .await
    }) {
        Ok(result) => result,
        Err(e) => {
            exit_with_error(&format!(
                "Failed to generate encryption meta with mode={}: {}",
                mode_name, e
            ));
        }
    };

    print_json(&build_gen_meta_output(&mode_for_print, result));
}

fn execute_sync_meta(args: SyncMetaArgs) {
    let ctx = build_ctx(&args.config, &args.pd);
    let runtime = ctx.dfs.get_runtime();

    let keyspace_id = resolve_required_keyspace(&ctx, &args.keyspace).id;
    let verbose = args.verbose;
    let mgr = EncryptionConfChangeManager::new(ctx.pd_client.clone(), OwnerCheckMode::Disabled);
    match args.mode {
        SyncMetaMode::Local => {
            if verbose {
                match get_keyspace_range(keyspace_id) {
                    Ok((ks_start_key, ks_end_key)) => {
                        println!(
                            "Keyspace range: {:?} - {:?}",
                            ks_start_key.as_encoded(),
                            ks_end_key.as_encoded()
                        );
                    }
                    Err(e) => {
                        eprintln!("{}", e);
                        std::process::exit(1);
                    }
                }
            }

            let start_result = match runtime.block_on(async {
                mgr.handle_start_rollout_job(StartRolloutJobRequest { keyspace_id })
                    .await
            }) {
                Ok(result) => result,
                Err(e) => {
                    eprintln!("Failed to start encryption conf-change job: {}", e);
                    std::process::exit(1);
                }
            };

            let resp = start_result.response;
            if !resp.started {
                if resp.running {
                    println!(
                        "Encryption conf-change job is already running for keyspace {}.",
                        keyspace_id
                    );
                } else {
                    println!(
                        "No pending encryption meta rotation for keyspace {} (next_file_id is not set).",
                        keyspace_id
                    );
                }
                return;
            }

            let Some(handle) = start_result.handle else {
                eprintln!("Started encryption conf-change job but missing wait handle");
                std::process::exit(1);
            };

            println!(
                "Syncing all regions from current_file_id={} to next_file_id={} in sequential mode...",
                resp.current_file_id, resp.next_file_id
            );

            let summary = match runtime.block_on(async { handle.wait().await }) {
                Ok(summary) => summary,
                Err(e) => {
                    eprintln!("Sync failed: {}", e);
                    std::process::exit(1);
                }
            };

            println!(
                "Sync completed: {} succeeded, {} failed",
                summary.regions_success, summary.regions_failed
            );
            println!(
                "Sync completed successfully. keyspace_id={}, current_file_id={}, next_file_id=0",
                keyspace_id, resp.next_file_id
            );
        }
        SyncMetaMode::Remote => {
            let resp = match runtime
                .block_on(async { mgr.request_start_rollout_job_on_owner(keyspace_id).await })
            {
                Ok(resp) => resp,
                Err(e) => {
                    eprintln!("Remote sync failed: {}", e);
                    std::process::exit(1);
                }
            };

            match serde_json::to_string_pretty(&resp) {
                Ok(json) => println!("{}", json),
                Err(e) => {
                    eprintln!("Failed to serialize remote sync response to JSON: {}", e);
                    std::process::exit(1);
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use clap::Parser;
    use cloud_encryption::KeyspaceEncryptionConfig;

    use super::*;

    #[derive(Parser)]
    struct TestGenMetaCli {
        #[clap(flatten)]
        args: GenMetaArgs,
    }

    fn base_gen_meta_args() -> GenMetaArgs {
        GenMetaArgs {
            config: PathBuf::new(),
            pd: String::new(),
            keyspace: KeyspaceArgs {
                keyspace_id: None,
                keyspace_name: Some("test".to_owned()),
            },
            mode: GenMetaMode::RotateDataKey,
            max_data_key_age: None,
            cmek_id: None,
            vendor: None,
            region: None,
            endpoint: None,
            encryption_context: None,
        }
    }

    #[test]
    fn test_gen_meta_cli_mode_default_is_rotate_data_key() {
        let parsed = TestGenMetaCli::parse_from(["test", "--keyspace-name", "test_keyspace"]);
        assert_eq!(parsed.args.mode, GenMetaMode::RotateDataKey);

        match build_gen_meta_mode(&parsed.args).unwrap() {
            EncryptionMetaGenMode::RotateDataKey { max_data_key_age } => {
                assert!(max_data_key_age.is_none());
            }
            EncryptionMetaGenMode::EnableCmek { .. } => panic!("unexpected enable-cmek mode"),
        }
    }

    #[test]
    fn test_gen_meta_cli_accepts_keyspace_id() {
        let parsed = TestGenMetaCli::parse_from(["test", "--keyspace-id", "42"]);
        assert_eq!(
            parsed.args.keyspace.parse_required_selector().unwrap(),
            KeyspaceSelector::Id(42)
        );
    }

    #[test]
    fn test_keyspace_args_accept_both_id_and_name() {
        let args = KeyspaceArgs {
            keyspace_id: Some(42),
            keyspace_name: Some("test_keyspace".to_owned()),
        };
        assert_eq!(
            args.parse_required_selector().unwrap(),
            KeyspaceSelector::IdAndName {
                id: 42,
                name: "test_keyspace".to_owned(),
            }
        );
    }

    #[test]
    fn test_keyspace_args_require_non_empty_name() {
        let args = KeyspaceArgs {
            keyspace_id: None,
            keyspace_name: Some("   ".to_owned()),
        };
        assert_eq!(
            args.parse_required_selector().unwrap_err(),
            "--keyspace-name must not be empty"
        );
    }

    #[test]
    fn test_build_gen_meta_mode_validation() {
        let mut args = base_gen_meta_args();
        args.cmek_id = Some("key-a".to_string());
        args.vendor = Some("aws".to_string());
        let err = build_gen_meta_mode(&args).unwrap_err();
        assert!(err.contains("--cmek-id"));
        assert!(err.contains("--vendor"));

        args = base_gen_meta_args();
        args.mode = GenMetaMode::EnableCmek;
        let missing_cmek_id = build_gen_meta_mode(&args).unwrap_err();
        assert!(missing_cmek_id.contains("--cmek-id is required"));

        args.cmek_id = Some("k".to_owned());
        let missing_vendor = build_gen_meta_mode(&args).unwrap_err();
        assert!(missing_vendor.contains("--vendor is required"));

        args.vendor = Some("aws".to_owned());
        args.max_data_key_age = Some(10);
        let max_age_err = build_gen_meta_mode(&args).unwrap_err();
        assert!(max_age_err.contains("--max-data-key-age"));
    }

    #[test]
    fn test_build_gen_meta_mode_enable_success() {
        let mut args = base_gen_meta_args();
        args.mode = GenMetaMode::EnableCmek;
        args.cmek_id = Some("k".to_owned());
        args.vendor = Some("aws".to_owned());
        args.region = Some("us-west-2".to_owned());
        args.endpoint = Some("https://kms.example".to_owned());
        args.encryption_context =
            Some(r#"{"ExternalId":"customer-42","tenant":"acme"}"#.to_owned());

        match build_gen_meta_mode(&args).unwrap() {
            EncryptionMetaGenMode::EnableCmek {
                cmek_id,
                vendor,
                region,
                endpoint,
                encryption_context,
            } => {
                assert_eq!(cmek_id, "k");
                assert_eq!(vendor, "aws");
                assert_eq!(region.as_deref(), Some("us-west-2"));
                assert_eq!(endpoint.as_deref(), Some("https://kms.example"));
                assert_eq!(
                    encryption_context,
                    BTreeMap::from([
                        ("ExternalId".to_owned(), "customer-42".to_owned()),
                        ("tenant".to_owned(), "acme".to_owned()),
                    ])
                );
            }
            EncryptionMetaGenMode::RotateDataKey { .. } => panic!("unexpected rotate mode"),
        }
    }

    #[test]
    fn test_encryption_config_json_round_trip() {
        let mut encryption_context = BTreeMap::new();
        encryption_context.insert("ExternalId".to_owned(), "customer-42".to_owned());
        encryption_context.insert("tenant".to_owned(), "acme".to_owned());
        let encryption_cfg = KeyspaceEncryptionConfig {
            enabled: true,
            cmek_id: Some("arn:aws:kms:us-west-2:123456:key/abc-def".to_owned()),
            vendor: Some("aws".to_owned()),
            region: Some("us-west-2".to_owned()),
            endpoint: None,
            encryption_context: encryption_context.clone(),
            ..Default::default()
        };
        let encryption_json = serde_json::to_string(&encryption_cfg).unwrap();
        let parsed: KeyspaceEncryptionConfig = serde_json::from_str(&encryption_json).unwrap();
        assert!(parsed.enabled);
        assert_eq!(
            parsed.cmek_id.as_deref(),
            Some("arn:aws:kms:us-west-2:123456:key/abc-def")
        );
        assert_eq!(parsed.encryption_context, encryption_context);
        assert_eq!(parsed.vendor.as_deref(), Some("aws"));
        assert_eq!(parsed.region.as_deref(), Some("us-west-2"));
        assert_eq!(parsed.endpoint, None);
    }

    #[test]
    fn test_parse_encryption_context() {
        let parsed = parse_encryption_context(r#"{"k1":"v1","k2":"v2"}"#).unwrap();
        assert_eq!(
            parsed,
            BTreeMap::from([
                ("k1".to_owned(), "v1".to_owned()),
                ("k2".to_owned(), "v2".to_owned()),
            ])
        );
    }

    #[test]
    fn test_parse_encryption_context_rejects_empty_input() {
        assert_eq!(
            parse_encryption_context("").unwrap_err(),
            "--encryption-context must not be empty"
        );
        assert_eq!(
            parse_encryption_context("   ").unwrap_err(),
            "--encryption-context must not be empty"
        );
    }

    #[test]
    fn test_parse_encryption_context_rejects_invalid_input() {
        let err = parse_encryption_context(r#"{"k1":"v1""#).unwrap_err();
        assert!(err.contains("JSON object"));
        assert!(!err.contains("\\\""));
        let err = parse_encryption_context(r#"{"k1":1}"#).unwrap_err();
        assert!(err.contains("JSON object"));
        assert!(!err.contains("\\\""));
    }

    #[test]
    fn test_parse_encryption_context_rejects_too_long_value() {
        let too_long = "a".repeat(KeyspaceEncryptionConfig::MAX_ENCRYPTION_CONTEXT_JSON_LEN + 1);
        let err = parse_encryption_context(&too_long).unwrap_err();
        assert!(err.contains("length"));
        assert!(err.contains("8192"));
    }
}