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

use bytes::Bytes;
use cloud_encryption::KeyspaceEncryptionConfig;
use encryption_conf_change::{
    EncryptionConfChangeManager, EncryptionMetaGenMode as ConfChangeGenMetaMode, GenMetaRequest,
    OwnerCheckMode, StartRolloutJobRequest,
};
use http::{Request, StatusCode};
use hyper::{Body, Client, client::HttpConnector};
use kvengine::dfs::{Dfs, FileType, Options};
use kvenginepb::EncryptionMeta;
use pd_client::PdClient;
use protobuf::Message;
use test_cloud_server::{ServerClusterBuilder, alloc_node_id_vec, oss::prepare_dfs, try_wait};
use tikv::config::TikvConfig;
use tikv_util::info;
use tokio::runtime::Runtime;

use crate::{create_keyspace_meta, new_security_config};

fn get_region_for_keyspace(pd_client: &dyn PdClient, keyspace_id: u32) -> kvproto::metapb::Region {
    let ks_start_key = api_version::ApiV2::get_keyspace_prefix_by_id(keyspace_id);
    let encoded_key = txn_types::Key::from_raw(&ks_start_key).into_encoded();
    pd_client.get_region(&encoded_key).unwrap()
}

/// Helper function to get region_id for a keyspace
fn get_region_id_for_keyspace(pd_client: &dyn PdClient, keyspace_id: u32) -> u64 {
    get_region_for_keyspace(pd_client, keyspace_id).get_id()
}

/// Test the `/encryption/gen-meta` API end-to-end with JSON body.
///
/// This test validates:
/// 1. API correctly generates new encryption meta with rotated data key
/// 2. The generated meta is correctly written to S3
/// 3. The meta can be read back from S3 and parsed correctly
#[test]
fn test_encryption_meta_gen_api() {
    test_util::init_log_for_test();

    let keyspace_id = 1u32;

    // Setup: start cluster with S3 mock
    let (_base_dir, mut oss, dfs_config) = prepare_dfs("test_encryption_meta_gen");
    let nodes = alloc_node_id_vec(1);
    let dfs_config_clone = dfs_config.clone();

    let cluster = ServerClusterBuilder::new(nodes.clone(), move |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config_clone.clone();
        conf.security = new_security_config();
    })
    .pd_server_cnt(1)
    .build();

    let node_id = nodes[0];
    let status_addr = cluster.status_addr(node_id);
    let pd_client = cluster.get_pd_client();

    // Setup encryption config for the keyspace
    let encryption_cfg = KeyspaceEncryptionConfig {
        enabled: true,
        cmek_id: Some("random".into()),
        vendor: Some("test".into()),
        ..Default::default()
    };
    pd_client
        .set_keyspace_encryption(keyspace_id, encryption_cfg)
        .unwrap();

    // Split keyspace and wait for region to be ready
    let mut client = cluster.new_client();
    client.split_keyspaces(keyspace_id..keyspace_id + 1);
    cluster.wait_pd_region_count(3); // default region + keyspace region + next region

    // Wait for shard to have encryption meta and get region_id
    let region_id = {
        let ok = try_wait(
            || {
                let kvengine = cluster.get_kvengine(node_id);
                let ks_start_key = api_version::ApiV2::get_keyspace_prefix_by_id(keyspace_id);
                let encoded_key = txn_types::Key::from_raw(&ks_start_key).into_encoded();
                if let Ok(region) = pd_client.get_region(&encoded_key) {
                    if let Some(shard) = kvengine.get_shard(region.get_id()) {
                        return shard.get_property(kvengine::ENCRYPTION_META_KEY).is_some();
                    }
                }
                false
            },
            30,
        );
        assert!(ok, "Shard should have encryption meta");
        get_region_id_for_keyspace(pd_client.as_ref(), keyspace_id)
    };

    // Create HTTP client and runtime
    let http_client: Client<HttpConnector> = Client::new();
    let runtime = Runtime::new().unwrap();

    // Call the encryption/gen-meta API with JSON body
    let api_url = format!("http://{}/encryption/gen-meta", status_addr);
    let request_body = serde_json::json!({
        "keyspace_id": keyspace_id,
        "region_id": region_id,
        "mode": "rotate_data_key"
    });
    info!("Calling API: {} with body: {}", api_url, request_body);

    let resp = runtime.block_on(async {
        let req = Request::post(&api_url)
            .header("Content-Type", "application/json")
            .body(Body::from(request_body.to_string()))
            .unwrap();
        http_client.request(req).await.unwrap()
    });

    assert_eq!(resp.status(), StatusCode::OK, "API should return 200 OK");

    // Parse response
    let body_bytes = runtime
        .block_on(hyper::body::to_bytes(resp.into_body()))
        .unwrap();
    let response: serde_json::Value =
        serde_json::from_slice(&body_bytes).expect("Failed to parse API response");

    let file_id = response["file_id"].as_u64().expect("file_id should exist");
    let data_key_id = response["data_key_id"]
        .as_u64()
        .expect("data_key_id should exist");
    let new_data_key_generated = response["new_data_key_generated"]
        .as_bool()
        .expect("new_data_key_generated should exist");
    assert!(
        new_data_key_generated,
        "new_data_key_generated should be true"
    );

    info!(
        "API response: file_id={}, data_key_id={}, new_data_key_generated={}",
        file_id, data_key_id, new_data_key_generated
    );

    // Verify: read the meta from S3
    let s3_fs = kvengine::dfs::S3Fs::new(
        dfs_config.prefix.clone(),
        dfs_config.s3_endpoint.clone(),
        dfs_config.s3_key_id.clone(),
        dfs_config.s3_secret_key.clone(),
        dfs_config.s3_region.clone(),
        dfs_config.s3_bucket.clone(),
        dfs_config.conn_options.clone(),
        dfs_config.read_only,
    );

    let mut opts = Options::default();
    opts.file_type = FileType::EncryptionDict;

    let meta_bytes: Bytes = s3_fs
        .get_runtime()
        .block_on(s3_fs.read_file(file_id, opts))
        .expect("Failed to read encryption meta from S3");

    // Parse and verify the meta
    let mut loaded_meta = EncryptionMeta::default();
    loaded_meta
        .merge_from_bytes(&meta_bytes)
        .expect("Failed to parse encryption meta");

    assert_eq!(
        loaded_meta.keyspace_id, keyspace_id,
        "keyspace_id should match"
    );
    assert_eq!(
        loaded_meta.get_current().file_id,
        file_id,
        "current epoch file_id should match"
    );
    assert_eq!(
        loaded_meta.get_current().data_key_id as u64,
        data_key_id,
        "current epoch data_key_id should match"
    );
    assert!(
        loaded_meta
            .get_data_keys()
            .contains_key(&(data_key_id as u32)),
        "data_keys should contain the new key"
    );
    assert!(
        !loaded_meta.get_history().is_empty(),
        "history should not be empty after rotation"
    );

    info!(
        "Verified encryption meta: keyspace_id={}, file_id={}, data_key_id={}, history_len={}",
        loaded_meta.keyspace_id,
        loaded_meta.get_current().file_id,
        loaded_meta.get_current().data_key_id,
        loaded_meta.get_history().len()
    );

    // Cleanup
    drop(s3_fs);
    drop(cluster);
    oss.graceful_shutdown();

    info!("test_encryption_meta_gen_api completed successfully!");
}

/// Test the `/encryption/sync-latest` API end-to-end with JSON body.
///
/// This test validates:
/// 1. A new encryption meta is generated and uploaded to S3 via
///    `/encryption/gen-meta`
/// 2. `/encryption/sync-latest` updates shard ENCRYPTION_META_KEY through
///    raftstore
/// 3. The shard property and in-memory encryption key are updated to the new
///    meta
#[test]
fn test_sync_latest_encryption_meta_api() {
    test_util::init_log_for_test();

    let keyspace_id = 1u32;

    // Setup: start cluster with S3 mock
    let (_base_dir, mut oss, dfs_config) = prepare_dfs("test_sync_latest_encryption_meta");
    let nodes = alloc_node_id_vec(1);
    let dfs_config_clone = dfs_config.clone();

    let cluster = ServerClusterBuilder::new(nodes.clone(), move |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config_clone.clone();
        conf.security = new_security_config();
    })
    .pd_server_cnt(1)
    .build();

    let node_id = nodes[0];
    let status_addr = cluster.status_addr(node_id);
    let pd_client = cluster.get_pd_client();

    // Setup encryption config for the keyspace
    let encryption_cfg = KeyspaceEncryptionConfig {
        enabled: true,
        cmek_id: Some("random".into()),
        vendor: Some("test".into()),
        ..Default::default()
    };
    pd_client
        .set_keyspace_encryption(keyspace_id, encryption_cfg)
        .unwrap();

    // Split keyspace and wait for region to be ready
    let mut client = cluster.new_client();
    client.split_keyspaces(keyspace_id..keyspace_id + 1);
    cluster.wait_pd_region_count(3); // default region + keyspace region + next region

    // Wait for shard to have encryption meta and get region_id
    let region_id = {
        let ok = try_wait(
            || {
                let kvengine = cluster.get_kvengine(node_id);
                let ks_start_key = api_version::ApiV2::get_keyspace_prefix_by_id(keyspace_id);
                let encoded_key = txn_types::Key::from_raw(&ks_start_key).into_encoded();
                if let Ok(region) = pd_client.get_region(&encoded_key) {
                    if let Some(shard) = kvengine.get_shard(region.get_id()) {
                        return shard.get_property(kvengine::ENCRYPTION_META_KEY).is_some();
                    }
                }
                false
            },
            30,
        );
        assert!(ok, "Shard should have encryption meta");
        get_region_id_for_keyspace(pd_client.as_ref(), keyspace_id)
    };

    // Read current encryption meta from shard
    let (current_file_id, current_data_key_id) = {
        let kvengine = cluster.get_kvengine(node_id);
        let shard = kvengine.get_shard(region_id).unwrap();
        let meta_bytes = shard
            .get_property(kvengine::ENCRYPTION_META_KEY)
            .expect("Shard should have encryption meta");
        let mut meta = EncryptionMeta::default();
        meta.merge_from_bytes(&meta_bytes)
            .expect("Failed to parse shard encryption meta");
        (
            meta.get_current().get_file_id(),
            meta.get_current().get_data_key_id(),
        )
    };

    // Create HTTP client and runtime
    let http_client: Client<HttpConnector> = Client::new();
    let runtime = Runtime::new().unwrap();

    // Step 1: Call `/encryption/gen-meta` to generate next meta in S3
    let gen_url = format!("http://{}/encryption/gen-meta", status_addr);
    let gen_body = serde_json::json!({
        "keyspace_id": keyspace_id,
        "region_id": region_id,
        "mode": "rotate_data_key"
    });
    info!("Calling API: {} with body: {}", gen_url, gen_body);

    let gen_resp = runtime.block_on(async {
        let req = Request::post(&gen_url)
            .header("Content-Type", "application/json")
            .body(Body::from(gen_body.to_string()))
            .unwrap();
        http_client.request(req).await.unwrap()
    });
    assert_eq!(
        gen_resp.status(),
        StatusCode::OK,
        "encryption/gen-meta should return 200 OK"
    );

    let gen_body_bytes = runtime
        .block_on(hyper::body::to_bytes(gen_resp.into_body()))
        .unwrap();
    let gen_response: serde_json::Value =
        serde_json::from_slice(&gen_body_bytes).expect("Failed to parse encryption/gen-meta resp");

    let next_file_id = gen_response["file_id"]
        .as_u64()
        .expect("file_id should exist");
    let next_data_key_id = gen_response["data_key_id"]
        .as_u64()
        .expect("data_key_id should exist");
    let new_data_key_generated = gen_response["new_data_key_generated"]
        .as_bool()
        .expect("new_data_key_generated should exist");
    assert!(
        new_data_key_generated,
        "new_data_key_generated should be true"
    );
    assert!(
        next_file_id > current_file_id,
        "next_file_id should be newer than current_file_id"
    );
    assert_ne!(
        next_data_key_id as u32, current_data_key_id,
        "data key should be rotated"
    );

    // Read new encryption meta from S3 for later verification.
    let s3_fs = kvengine::dfs::S3Fs::new(
        dfs_config.prefix.clone(),
        dfs_config.s3_endpoint.clone(),
        dfs_config.s3_key_id.clone(),
        dfs_config.s3_secret_key.clone(),
        dfs_config.s3_region.clone(),
        dfs_config.s3_bucket.clone(),
        dfs_config.conn_options.clone(),
        dfs_config.read_only,
    );
    let mut opts = Options::default();
    opts.file_type = FileType::EncryptionDict;
    let new_meta_bytes: Bytes = s3_fs
        .get_runtime()
        .block_on(s3_fs.read_file(next_file_id, opts))
        .expect("Failed to read new encryption meta from S3");

    // Step 2: Call `/encryption/sync-latest` to apply the new meta through
    // raftstore.
    let sync_url = format!("http://{}/encryption/sync-latest", status_addr);
    let region = get_region_for_keyspace(pd_client.as_ref(), keyspace_id);
    let sync_body = serde_json::json!({
        "keyspace_id": keyspace_id,
        "region_id": region_id,
        "current_file_id": current_file_id,
        "next_file_id": next_file_id,
        "expected_region_version": region.get_region_epoch().get_version(),
        "expected_region_conf_ver": region.get_region_epoch().get_conf_ver()
    });
    info!("Calling API: {} with body: {}", sync_url, sync_body);

    let sync_resp = runtime.block_on(async {
        let req = Request::post(&sync_url)
            .header("Content-Type", "application/json")
            .body(Body::from(sync_body.to_string()))
            .unwrap();
        http_client.request(req).await.unwrap()
    });
    assert_eq!(
        sync_resp.status(),
        StatusCode::OK,
        "encryption/sync-latest should return 200 OK"
    );

    let sync_body_bytes = runtime
        .block_on(hyper::body::to_bytes(sync_resp.into_body()))
        .unwrap();
    let sync_response: serde_json::Value = serde_json::from_slice(&sync_body_bytes)
        .expect("Failed to parse encryption/sync-latest resp");

    assert_eq!(
        sync_response["file_id"].as_u64().unwrap(),
        next_file_id,
        "file_id should be updated to next_file_id"
    );
    assert_eq!(
        sync_response["data_key_id"].as_u64().unwrap(),
        next_data_key_id,
        "data_key_id should match new meta current key"
    );
    assert!(
        sync_response["synced"].as_bool().unwrap(),
        "synced should be true on first sync"
    );

    // Verify: shard property is updated and matches S3 content.
    let ok = try_wait(
        || {
            let kvengine = cluster.get_kvengine(node_id);
            let Some(shard) = kvengine.get_shard(region_id) else {
                return false;
            };
            let Some(meta_bytes) = shard.get_property(kvengine::ENCRYPTION_META_KEY) else {
                return false;
            };
            let mut meta = EncryptionMeta::default();
            if meta.merge_from_bytes(&meta_bytes).is_err() {
                return false;
            }
            meta.get_current().get_file_id() == next_file_id
        },
        30,
    );
    assert!(ok, "Shard encryption meta should be updated after sync");

    let kvengine = cluster.get_kvengine(node_id);
    let shard = kvengine.get_shard(region_id).unwrap();
    let synced_meta_bytes = shard
        .get_property(kvengine::ENCRYPTION_META_KEY)
        .expect("Shard should have encryption meta after sync");
    assert_eq!(
        synced_meta_bytes.as_ref(),
        new_meta_bytes.as_ref(),
        "Shard ENCRYPTION_META_KEY should match the meta in S3"
    );
    let encryption_key = shard
        .try_get_encryption_key()
        .unwrap()
        .expect("Shard should have encryption key after sync");
    assert_eq!(
        encryption_key.get_key_id().unwrap() as u64,
        next_data_key_id,
        "In-memory encryption key should match new meta current key"
    );

    // Verify idempotency: calling again with the same next_file_id returns
    // synced=false.
    let sync_resp2 = runtime.block_on(async {
        let req = Request::post(&sync_url)
            .header("Content-Type", "application/json")
            .body(Body::from(sync_body.to_string()))
            .unwrap();
        http_client.request(req).await.unwrap()
    });
    assert_eq!(
        sync_resp2.status(),
        StatusCode::OK,
        "second encryption/sync-latest should return 200 OK"
    );
    let sync_body_bytes2 = runtime
        .block_on(hyper::body::to_bytes(sync_resp2.into_body()))
        .unwrap();
    let sync_response2: serde_json::Value = serde_json::from_slice(&sync_body_bytes2)
        .expect("Failed to parse second encryption/sync-latest resp");
    assert_eq!(sync_response2["file_id"].as_u64().unwrap(), next_file_id);
    assert_eq!(
        sync_response2["data_key_id"].as_u64().unwrap(),
        next_data_key_id
    );
    assert!(
        !sync_response2["synced"].as_bool().unwrap(),
        "synced should be false when already at next_file_id"
    );

    // Cleanup
    drop(s3_fs);
    drop(cluster);
    oss.graceful_shutdown();
}

/// Test enabling CMEK for a previously unencrypted keyspace.
#[test]
fn test_sync_latest_encryption_meta_initial_enable_cmek() {
    test_util::init_log_for_test();

    let keyspace_id = 1u32;

    // Setup: start cluster with S3 mock
    let (_base_dir, mut oss, dfs_config) = prepare_dfs("test_sync_latest_encryption_meta_enable");
    let nodes = alloc_node_id_vec(1);
    let dfs_config_clone = dfs_config.clone();

    let cluster = ServerClusterBuilder::new(nodes.clone(), move |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config_clone.clone();
        conf.security = new_security_config();
    })
    .pd_server_cnt(1)
    .build();

    let node_id = nodes[0];
    let status_addr = cluster.status_addr(node_id);
    let pd_client = cluster.get_pd_client();

    // Split keyspace and wait for region to be ready
    let mut client = cluster.new_client();
    client.split_keyspaces(keyspace_id..keyspace_id + 1);
    cluster.wait_pd_region_count(3); // default region + keyspace region + next region

    let region_id = get_region_id_for_keyspace(pd_client.as_ref(), keyspace_id);

    // Wait for shard to be created and ensure it has no encryption meta initially.
    let ok = try_wait(
        || {
            let kvengine = cluster.get_kvengine(node_id);
            kvengine
                .get_shard(region_id)
                .map(|s| s.get_property(kvengine::ENCRYPTION_META_KEY).is_none())
                .unwrap_or(false)
        },
        30,
    );
    assert!(ok, "Shard should exist and have no encryption meta");

    // Create HTTP client and runtime
    let http_client: Client<HttpConnector> = Client::new();
    let runtime = Runtime::new().unwrap();

    // Step 1: Call `/encryption/gen-meta` to generate initial meta in S3
    let gen_url = format!("http://{}/encryption/gen-meta", status_addr);
    let gen_body = serde_json::json!({
        "keyspace_id": keyspace_id,
        "region_id": region_id,
        "mode": "enable_cmek",
        "cmek_id": "random",
        "vendor": "test"
    });
    info!("Calling API: {} with body: {}", gen_url, gen_body);

    let gen_resp = runtime.block_on(async {
        let req = Request::post(&gen_url)
            .header("Content-Type", "application/json")
            .body(Body::from(gen_body.to_string()))
            .unwrap();
        http_client.request(req).await.unwrap()
    });
    assert_eq!(
        gen_resp.status(),
        StatusCode::OK,
        "encryption/gen-meta should return 200 OK"
    );

    let gen_body_bytes = runtime
        .block_on(hyper::body::to_bytes(gen_resp.into_body()))
        .unwrap();
    let gen_response: serde_json::Value =
        serde_json::from_slice(&gen_body_bytes).expect("Failed to parse encryption/gen-meta resp");

    let next_file_id = gen_response["file_id"]
        .as_u64()
        .expect("file_id should exist");
    let next_data_key_id = gen_response["data_key_id"]
        .as_u64()
        .expect("data_key_id should exist");

    // Read new encryption meta from S3 for verification.
    let s3_fs = kvengine::dfs::S3Fs::new(
        dfs_config.prefix.clone(),
        dfs_config.s3_endpoint.clone(),
        dfs_config.s3_key_id.clone(),
        dfs_config.s3_secret_key.clone(),
        dfs_config.s3_region.clone(),
        dfs_config.s3_bucket.clone(),
        dfs_config.conn_options.clone(),
        dfs_config.read_only,
    );
    let mut opts = Options::default();
    opts.file_type = FileType::EncryptionDict;
    let new_meta_bytes: Bytes = s3_fs
        .get_runtime()
        .block_on(s3_fs.read_file(next_file_id, opts))
        .expect("Failed to read new encryption meta from S3");

    // Step 2: Call `/encryption/sync-latest` to apply the new meta through
    // raftstore.
    let sync_url = format!("http://{}/encryption/sync-latest", status_addr);
    let region = get_region_for_keyspace(pd_client.as_ref(), keyspace_id);
    let sync_body = serde_json::json!({
        "keyspace_id": keyspace_id,
        "region_id": region_id,
        "current_file_id": 0u64,
        "next_file_id": next_file_id,
        "expected_region_version": region.get_region_epoch().get_version(),
        "expected_region_conf_ver": region.get_region_epoch().get_conf_ver()
    });
    info!("Calling API: {} with body: {}", sync_url, sync_body);

    let sync_resp = runtime.block_on(async {
        let req = Request::post(&sync_url)
            .header("Content-Type", "application/json")
            .body(Body::from(sync_body.to_string()))
            .unwrap();
        http_client.request(req).await.unwrap()
    });
    assert_eq!(
        sync_resp.status(),
        StatusCode::OK,
        "encryption/sync-latest should return 200 OK"
    );

    let sync_body_bytes = runtime
        .block_on(hyper::body::to_bytes(sync_resp.into_body()))
        .unwrap();
    let sync_response: serde_json::Value = serde_json::from_slice(&sync_body_bytes)
        .expect("Failed to parse encryption/sync-latest resp");

    assert_eq!(
        sync_response["file_id"].as_u64().unwrap(),
        next_file_id,
        "file_id should be updated to next_file_id"
    );
    assert_eq!(
        sync_response["data_key_id"].as_u64().unwrap(),
        next_data_key_id,
        "data_key_id should match new meta current key"
    );
    assert!(
        sync_response["synced"].as_bool().unwrap(),
        "synced should be true on first sync"
    );

    // Verify shard property and in-memory key are updated.
    let ok = try_wait(
        || {
            let kvengine = cluster.get_kvengine(node_id);
            let Some(shard) = kvengine.get_shard(region_id) else {
                return false;
            };
            let Some(meta_bytes) = shard.get_property(kvengine::ENCRYPTION_META_KEY) else {
                return false;
            };
            let mut meta = EncryptionMeta::default();
            if meta.merge_from_bytes(&meta_bytes).is_err() {
                return false;
            }
            meta.get_current().get_file_id() == next_file_id
        },
        30,
    );
    assert!(ok, "Shard encryption meta should be updated after sync");

    let kvengine = cluster.get_kvengine(node_id);
    let shard = kvengine.get_shard(region_id).unwrap();
    let synced_meta_bytes = shard
        .get_property(kvengine::ENCRYPTION_META_KEY)
        .expect("Shard should have encryption meta after sync");
    assert_eq!(
        synced_meta_bytes.as_ref(),
        new_meta_bytes.as_ref(),
        "Shard ENCRYPTION_META_KEY should match the meta in S3"
    );
    let encryption_key = shard
        .try_get_encryption_key()
        .unwrap()
        .expect("Shard should have encryption key after sync");
    assert_eq!(
        encryption_key.get_key_id().unwrap() as u64,
        next_data_key_id,
        "In-memory encryption key should match new meta current key"
    );

    // Cleanup
    drop(s3_fs);
    drop(cluster);
    oss.graceful_shutdown();
}

#[test]
fn test_conf_change_manager_first_rotate_backfills_pd_current_and_rolls_out() {
    test_util::init_log_for_test();

    let keyspace_id = 1u32;

    let (_base_dir, mut oss, dfs_config) = prepare_dfs("test_conf_change_manager_first_rotate");
    let nodes = alloc_node_id_vec(1);
    let dfs_config_clone = dfs_config.clone();

    let cluster = ServerClusterBuilder::new(nodes.clone(), move |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config_clone.clone();
        conf.security = new_security_config();
    })
    .pd_server_cnt(1)
    .build();

    let runtime = Runtime::new().unwrap();
    create_keyspace_meta(&cluster, keyspace_id, &runtime);

    let node_id = nodes[0];
    let pd_client = cluster.get_pd_client();

    let encryption_cfg = KeyspaceEncryptionConfig {
        enabled: true,
        cmek_id: Some("random".into()),
        vendor: Some("test".into()),
        ..Default::default()
    };
    pd_client
        .set_keyspace_encryption(keyspace_id, encryption_cfg)
        .unwrap();

    let mut client = cluster.new_client();
    client.split_keyspaces(keyspace_id..keyspace_id + 1);
    cluster.wait_pd_region_count(3);

    let region_id = {
        let ok = try_wait(
            || {
                let kvengine = cluster.get_kvengine(node_id);
                let ks_start_key = api_version::ApiV2::get_keyspace_prefix_by_id(keyspace_id);
                let encoded_key = txn_types::Key::from_raw(&ks_start_key).into_encoded();
                if let Ok(region) = pd_client.get_region(&encoded_key) {
                    if let Some(shard) = kvengine.get_shard(region.get_id()) {
                        return shard.get_property(kvengine::ENCRYPTION_META_KEY).is_some();
                    }
                }
                false
            },
            30,
        );
        assert!(ok, "Shard should have encryption meta");
        get_region_id_for_keyspace(pd_client.as_ref(), keyspace_id)
    };

    let initial_current_file_id = {
        let kvengine = cluster.get_kvengine(node_id);
        let shard = kvengine.get_shard(region_id).unwrap();
        let meta_bytes = shard
            .get_property(kvengine::ENCRYPTION_META_KEY)
            .expect("Shard should have encryption meta");
        let mut meta = EncryptionMeta::default();
        meta.merge_from_bytes(&meta_bytes)
            .expect("Failed to parse shard encryption meta");
        meta.get_current().get_file_id()
    };
    assert!(
        initial_current_file_id > 0,
        "Initial current_file_id should be non-zero"
    );

    let mgr = EncryptionConfChangeManager::new(pd_client.clone(), OwnerCheckMode::Disabled);

    let gen_result = runtime
        .block_on(async {
            mgr.handle_gen_meta(GenMetaRequest {
                keyspace_id,
                mode: ConfChangeGenMetaMode::RotateDataKey {
                    max_data_key_age: None,
                },
            })
            .await
        })
        .unwrap();

    assert_eq!(gen_result.pd_before.current_file_id, 0);
    assert_eq!(gen_result.pd_before.next_file_id, 0);
    assert_eq!(gen_result.pd_after.current_file_id, initial_current_file_id);
    assert_eq!(
        gen_result.pd_after.next_file_id,
        gen_result.gen_meta.file_id
    );
    assert!(gen_result.pd_updated);
    assert!(
        gen_result.gen_meta.file_id > initial_current_file_id,
        "new file_id should be newer than the preexisting shard current_file_id"
    );

    let start_result = runtime
        .block_on(async {
            mgr.handle_start_rollout_job(StartRolloutJobRequest { keyspace_id })
                .await
        })
        .unwrap();
    assert!(start_result.response.started);
    assert!(start_result.response.running);

    runtime
        .block_on(start_result.handle.expect("missing rollout handle").wait())
        .unwrap();

    let final_status = mgr
        .load_keyspace_encryption_status_from_pd(keyspace_id)
        .unwrap();
    assert_eq!(
        final_status.meta_state.current_file_id,
        gen_result.gen_meta.file_id
    );
    assert_eq!(final_status.meta_state.next_file_id, 0);

    let ok = try_wait(
        || {
            let kvengine = cluster.get_kvengine(node_id);
            let Some(shard) = kvengine.get_shard(region_id) else {
                return false;
            };
            let Some(meta_bytes) = shard.get_property(kvengine::ENCRYPTION_META_KEY) else {
                return false;
            };
            let mut meta = EncryptionMeta::default();
            if meta.merge_from_bytes(&meta_bytes).is_err() {
                return false;
            }
            meta.get_current().get_file_id() == gen_result.gen_meta.file_id
        },
        30,
    );
    assert!(
        ok,
        "Shard encryption meta should be rolled out to the new file_id"
    );

    drop(cluster);
    oss.graceful_shutdown();
}