// Copyright 2025 PingCAP, Inc.

use std::{
    collections::{HashMap, HashSet},
    sync::Mutex,
    time::Duration,
};

use bytes::Bytes;
use hyper::{Body, Request, StatusCode};
use kvengine::{
    self, STORAGE_CLASS_KEY, StorageClassStatusResp,
    table::{BoundedDataSet, DataBound, sstable::BlockCacheType},
    table_id::encode_table_prefix_key,
};
use kvproto::kvrpcpb;
use pd_client::PdClient;
use rstest::rstest;
use schema::{
    generate_storage_class_schema_data_for_test,
    schema::{StorageClass, StorageClassSpec, StorageClassTransitRule},
};
use test_cloud_server::{
    ServerCluster, ServerClusterBuilder,
    client::ClusterClient,
    keyspace::{CreateKeyspaceOptions, make_index_key},
    must_wait,
    oss::prepare_dfs,
    util::{Mutation, get_keyspace_split_keys},
};
use test_pd_client::{PdClientExt, PdWrapper};
use tidb_query_datatype::{codec::table::decode_table_id, expr::EvalContext};
use tikv::config::TikvConfig;
use tikv_util::{
    codec::{bytes::encode_bytes, number::NumberEncoder},
    config::{AbsoluteOrPercentSize, ReadableDuration, ReadableSize},
};

use crate::{
    alloc_node_id,
    columnar::{gen_row_key, gen_row_val},
    new_security_config,
};

const SEGMENT_SIZE: i64 = 64;
const FREQ_UPDATE_INTERVAL: Duration = Duration::from_millis(200);

#[rstest]
#[case::ordinary(false)]
fn test_storage_class_with_schema_manager(#[case] partitioned: bool) {
    let node_id = alloc_node_id();
    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_storage_class_with_schema_manager");
    let security_config = new_security_config();
    let pd_wrapper = PdWrapper::new_test(1, &security_config, None);
    let mut cluster = ServerClusterBuilder::new(vec![node_id], |_, conf: &mut TikvConfig| {
        conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(1);
        conf.rocksdb.writecf.block_size = ReadableSize(512);
        conf.rocksdb.writecf.target_file_size_base = ReadableSize::kb(2);
        conf.coprocessor.region_split_size = ReadableSize::kb(512);
        conf.coprocessor.region_bucket_size = ReadableSize::kb(64);
        conf.storage.block_cache.capacity = Some(ReadableSize::kb(16));
        conf.dfs = dfs_config.clone();
        conf.kvengine.ia.segment_size = SEGMENT_SIZE;
        conf.kvengine.ia.freq_update_interval = ReadableDuration(FREQ_UPDATE_INTERVAL);
        conf.kvengine.ia.mem_cap = AbsoluteOrPercentSize::Abs(ReadableSize::kb(64));
        conf.kvengine.ia.disk_cap = AbsoluteOrPercentSize::Abs(ReadableSize::kb(1024));
        conf.kvengine.ia.auto_ia_check_interval = ReadableDuration::millis(100);
        conf.kvengine.block_cache_type = BlockCacheType::Quick;
        conf.security = security_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_schema_manager(alloc_node_id());
    let dfs = cluster.get_dfs().unwrap();
    let keyspace_id = 9;
    let db_id = 5;
    let table_count = 3;
    let table_ids = dfs
        .get_runtime()
        .block_on(create_keyspace_region_without_splitting_table(
            &mut cluster,
            keyspace_id,
            table_count,
        ));
    let kvengine = cluster.get_kvengine(node_id);
    let status_addr = cluster.status_addr(node_id);
    assert!(kvengine.is_ia_enabled());
    let mut client = cluster.new_client();
    let ctx = Mutex::new(EvalContext::default());

    let table_id = table_ids[1];
    let partition_id = partitioned.then_some(table_ids[2]);
    let backend_table_id = partition_id.unwrap_or(table_id);

    let mut tables = None;
    let mut partitions = partition_id.map(|id| vec![id]);

    let table_start_key = encode_table_prefix_key(backend_table_id);
    let table_end_key = encode_table_prefix_key(backend_table_id + 1);
    let table_bound = DataBound::new(table_start_key.as_ref(), table_end_key.as_ref(), false);
    let table_region_id = client.get_region_id(&gen_row_key(keyspace_id, backend_table_id, 0));

    let mut schema_version = 10;
    let mut update_schema = |client: &mut ClusterClient, sc_spec: StorageClassSpec| {
        let kv_pairs = generate_storage_class_schema_data_for_test(
            keyspace_id,
            db_id,
            table_id,
            partition_id,
            partitions.take(),
            schema_version,
            sc_spec,
            &mut tables,
        )
        .unwrap();
        put_kv_pairs(client, kv_pairs);
        schema_version += 1;
    };

    // Init schema data
    update_schema(&mut client, StorageClassSpec::default());

    // Create WriteCF ln files
    client.put_kv(
        0..50,
        |i: usize| gen_row_key(keyspace_id, backend_table_id, i),
        |i: usize| gen_row_val(&ctx, i),
    );
    client.put_kv(
        50..100,
        |i: usize| gen_row_key(keyspace_id, backend_table_id, i),
        |i: usize| gen_row_val(&ctx, i),
    );
    client.verify_data_with_ref_store();
    must_wait(
        || {
            let all_id_vers = kvengine.get_all_shard_id_vers();
            for id_ver in all_id_vers {
                if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                    let stats = shard.get_stats();
                    if stats.cfs[0].levels[0].num_tables > 0 {
                        return true;
                    }
                }
            }
            false
        },
        10,
        || "failed to get ln files".to_string(),
    );

    // STANDARD maps to the default local storage class.
    must_wait(
        || {
            let status = dfs.get_runtime().block_on(request_storage_class_status(
                &status_addr,
                keyspace_id,
                backend_table_id,
                "STANDARD",
            ));
            status.total > 0 && status.ready == status.total
        },
        10,
        || "storage class status did not become STANDARD".to_string(),
    );

    // Change Local files to IA
    let sc_spec_ia = StorageClassSpec::from(StorageClass::Ia);
    update_schema(&mut client, sc_spec_ia.clone());
    let mut num_ia_files = 0usize;
    must_wait(
        || {
            let all_id_vers = kvengine.get_all_shard_id_vers();
            for id_ver in all_id_vers {
                if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                    if shard.get_schema_file().is_some()
                        && shard.storage_class_spec_equals(&sc_spec_ia)
                    {
                        assert!(table_bound.contains_bound(shard.range.data_bound()));
                        let (_, sst_ia_file_ids) = shard.get_local_sst_files();
                        assert!(!sst_ia_file_ids.is_empty());
                        num_ia_files = sst_ia_file_ids.len();
                        return true;
                    }
                }
            }
            false
        },
        10,
        || "failed to reload ia file".to_string(),
    );
    let ia_status = dfs.get_runtime().block_on(request_storage_class_status(
        &status_addr,
        keyspace_id,
        backend_table_id,
        "IA",
    ));
    assert!(ia_status.total > 0);
    assert_eq!(ia_status.ready, ia_status.total);
    let standard_status = dfs.get_runtime().block_on(request_storage_class_status(
        &status_addr,
        keyspace_id,
        backend_table_id,
        "STANDARD",
    ));
    assert_eq!(standard_status.ready, 0);
    assert_eq!(standard_status.total, ia_status.total);
    let invalid_target_status = dfs
        .get_runtime()
        .block_on(request_storage_class_status_response(
            &status_addr,
            keyspace_id,
            backend_table_id,
            Some("ARCHIVE"),
        ))
        .status();
    assert_eq!(invalid_target_status, StatusCode::BAD_REQUEST);
    let missing_target_status = dfs
        .get_runtime()
        .block_on(request_storage_class_status_response(
            &status_addr,
            keyspace_id,
            backend_table_id,
            None,
        ))
        .status();
    assert_eq!(missing_target_status, StatusCode::BAD_REQUEST);

    must_wait(
        || {
            let (expected_sst_files, expected_sst_ia_files) =
                collect_local_sst_files_by_type(&kvengine);
            let sst_files = dfs
                .get_runtime()
                .block_on(request_files_with_type(&status_addr, "sst"));
            let sst_ia_files = dfs
                .get_runtime()
                .block_on(request_files_with_type(&status_addr, "sst_ia"));
            sst_files.len() == 1
                && sst_files.get("sst") == Some(&expected_sst_files)
                && sst_ia_files.len() == 1
                && sst_ia_files.get("sst_ia") == Some(&expected_sst_ia_files)
        },
        10,
        || "files_with_type does not match engine files".to_string(),
    );
    client.verify_data_with_ref_store();

    // Create incremental IA files
    client.put_kv(
        100..150,
        |i: usize| gen_row_key(keyspace_id, backend_table_id, i),
        |i: usize| gen_row_val(&ctx, i),
    );
    client.put_kv(
        150..200,
        |i: usize| gen_row_key(keyspace_id, backend_table_id, i),
        |i: usize| gen_row_val(&ctx, i),
    );
    cluster.flush_memtable(table_region_id).unwrap();
    assert!(cluster.wait_for_memtable_flushed(table_region_id, Duration::from_secs(10)));
    must_wait(
        || {
            let all_id_vers = kvengine.get_all_shard_id_vers();
            for id_ver in all_id_vers {
                if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                    if shard.get_schema_file().is_some()
                        && shard.storage_class_spec_equals(&sc_spec_ia)
                    {
                        assert!(table_bound.contains_bound(shard.range.data_bound()));
                        let (_, sst_ia_file_ids) = shard.get_local_sst_files();
                        assert!(!sst_ia_file_ids.is_empty());
                        if sst_ia_file_ids.len() > num_ia_files {
                            return true;
                        }
                    }
                }
            }
            false
        },
        10,
        || "failed to load new ia file".to_string(),
    );
    client.verify_data_with_ref_store();

    // Change IA files to Local files
    update_schema(&mut client, StorageClass::Standard.into());
    must_wait(
        || {
            let all_id_vers = kvengine.get_all_shard_id_vers();
            for id_ver in all_id_vers {
                if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                    if id_ver.id == table_region_id
                        && (shard.is_compacting() || shard.has_compaction_priority())
                    {
                        return false;
                    }
                    if shard.get_schema_file().is_some()
                        || shard.get_property(STORAGE_CLASS_KEY).is_some()
                    {
                        return false;
                    }
                    let (_, sst_ia_file_ids) = shard.get_local_sst_files();
                    assert!(sst_ia_file_ids.is_empty());
                }
            }
            true
        },
        10,
        || "failed to reload local file or remove storage class property".to_string(),
    );
    let standard_status = dfs.get_runtime().block_on(request_storage_class_status(
        &status_addr,
        keyspace_id,
        backend_table_id,
        "STANDARD",
    ));
    assert!(standard_status.total > 0);
    assert_eq!(standard_status.ready, standard_status.total);
    client.verify_data_with_ref_store();

    // Change to auto IA (unspecified -> IA)
    {
        let sc_spec_auto = StorageClassSpec {
            default_tier: StorageClass::Unspecified,
            transit_rules: vec![StorageClassTransitRule {
                tier: StorageClass::Ia,
                transit_after: Duration::from_secs(1),
            }],
        };
        update_schema(&mut client, sc_spec_auto.clone());
        must_wait(
            || {
                let all_id_vers = kvengine.get_all_shard_id_vers();
                for id_ver in all_id_vers {
                    if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                        if shard.get_schema_file().is_some()
                            && shard.storage_class_spec_equals(&sc_spec_auto)
                        {
                            assert!(table_bound.contains_bound(shard.range.data_bound()));
                            return true;
                        }
                    }
                }
                false
            },
            10,
            || "failed to update schema to auto IA".to_string(),
        );
        client.verify_data_with_ref_store();
    }
    let (_, auto_ia_file_ids_origin) = collect_local_sst_files_by_type(&kvengine);

    // Verify unspecified -> IA.
    {
        client.put_kv(
            200..300,
            |i: usize| gen_row_key(keyspace_id, backend_table_id, i),
            |i: usize| gen_row_val(&ctx, i),
        );
        cluster.flush_memtable(table_region_id).unwrap();
        assert!(cluster.wait_for_memtable_flushed(table_region_id, Duration::from_secs(10)));
        must_wait(
            || {
                let (_, auto_ia_file_ids) = collect_local_sst_files_by_type(&kvengine);
                !auto_ia_file_ids.is_subset(&auto_ia_file_ids_origin)
            },
            10,
            || "wait for new IA file timeout".to_string(),
        );
        client.verify_data_with_ref_store();
    }

    cluster.stop();
    oss.shutdown();
}

#[rstest]
#[case::partitioned(true)]
fn test_region_split_and_merge_with_storage_class(#[case] partitioned: bool) {
    let node_id = alloc_node_id();
    let (_temp_dir, mut oss, dfs_config) =
        prepare_dfs("test_region_split_and_merge_with_storage_class");
    let security_config = new_security_config();
    let pd_wrapper = PdWrapper::new_test(1, &security_config, None);
    let mut cluster = ServerClusterBuilder::new(vec![node_id], |_, conf: &mut TikvConfig| {
        conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(1);
        conf.rocksdb.writecf.block_size = ReadableSize(512);
        conf.rocksdb.writecf.target_file_size_base = ReadableSize::kb(2);
        conf.coprocessor.region_split_size = ReadableSize::kb(512);
        conf.coprocessor.region_bucket_size = ReadableSize::kb(64);
        conf.storage.block_cache.capacity = Some(ReadableSize::kb(16));
        conf.dfs = dfs_config.clone();
        conf.kvengine.ia.segment_size = SEGMENT_SIZE;
        conf.kvengine.ia.freq_update_interval = ReadableDuration(FREQ_UPDATE_INTERVAL);
        conf.kvengine.ia.mem_cap = AbsoluteOrPercentSize::Abs(ReadableSize::kb(64));
        conf.kvengine.ia.disk_cap = AbsoluteOrPercentSize::Abs(ReadableSize::kb(1024));
        conf.kvengine.block_cache_type = BlockCacheType::Quick;
        conf.security = security_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_schema_manager(alloc_node_id());
    let dfs = cluster.get_dfs().unwrap();
    let keyspace_id = 9;
    let db_id = 5;
    let table_count = 4;
    let mut table_ids = dfs
        .get_runtime()
        .block_on(create_keyspace_region_without_splitting_table(
            &mut cluster,
            keyspace_id,
            table_count,
        ));
    table_ids.sort();
    let kvengine = cluster.get_kvengine(node_id);
    assert!(kvengine.is_ia_enabled());
    let mut client = cluster.new_client();
    let pd_client = cluster.get_pd_client();
    let ctx = Mutex::new(EvalContext::default());

    let mut tables = None;
    let mut partitions = partitioned.then(|| table_ids.clone());
    let mut schema_version = 10;
    let mut update_schema =
        |client: &mut ClusterClient, sc_spec: StorageClassSpec, backend_table_id: i64| {
            let (table_id, partition_id) = if partitioned {
                (65535, Some(backend_table_id))
            } else {
                (backend_table_id, None)
            };
            let kv_pairs = generate_storage_class_schema_data_for_test(
                keyspace_id,
                db_id,
                table_id,
                partition_id,
                partitions.take(),
                schema_version,
                sc_spec,
                &mut tables,
            )
            .unwrap();
            put_kv_pairs(client, kv_pairs);
            schema_version += 1;
        };

    let check_ia_shard_range = |range: DataBound<'_>| {
        let shard_inner_start = range.lower_bound.to_vec();
        let shard_table_id = decode_table_id(&shard_inner_start).unwrap();
        let table_start_key = encode_table_prefix_key(shard_table_id);
        let table_end_key = encode_table_prefix_key(shard_table_id + 1);
        let table_bound = DataBound::new(table_start_key.as_ref(), table_end_key.as_ref(), false);
        assert!(table_bound.contains_bound(range));
    };

    // Init schema data
    for &table_id in &table_ids {
        update_schema(&mut client, StorageClassSpec::default(), table_id);
    }

    // Write data for all tables
    for &table_id in &table_ids {
        client.put_kv(
            0..50,
            |i: usize| gen_row_key(keyspace_id, table_id, i),
            |i: usize| gen_row_val(&ctx, i),
        );
        client.put_kv(
            50..100,
            |i: usize| gen_row_key(keyspace_id, table_id, i),
            |i: usize| gen_row_val(&ctx, i),
        );
        client.verify_data_with_ref_store();
    }

    // Keep both different-tier and same-tier table boundaries in the matrix.
    // ( t0 Local / t1 IA / t2 IA / t3 Local )
    let ia_table_indexes = [1, 2];
    let mut last_table_id = None;
    let mut table_start_keys = Vec::with_capacity(table_ids.len());
    for &table_id in &table_ids {
        if let Some(last_table_id) = last_table_id {
            assert_eq!(last_table_id + 1, table_id);
        }
        last_table_id = Some(table_id);
        table_start_keys.push(gen_row_key(keyspace_id, table_id, 0))
    }
    let mut ia_file_tables = Vec::with_capacity(ia_table_indexes.len());
    let mut ia_table_start_keys = Vec::with_capacity(ia_table_indexes.len());
    for &ia_table_idx in &ia_table_indexes {
        ia_file_tables.push(table_ids[ia_table_idx]);
        ia_table_start_keys.push(&table_start_keys[ia_table_idx]);
    }
    let sc_spec_ia = StorageClassSpec::from(StorageClass::Ia);
    for &table_id in &ia_file_tables {
        update_schema(&mut client, sc_spec_ia.clone(), table_id);
    }
    must_wait(
        || {
            let mut shard_with_ia_count = 0;
            for &id_ver in &kvengine.get_all_shard_id_vers() {
                let shard = kvengine.get_shard(id_ver.id).unwrap();
                if shard.get_schema_file().is_some() && shard.storage_class_spec_equals(&sc_spec_ia)
                {
                    check_ia_shard_range(shard.range.data_bound());
                    let (_, sst_ia_file_ids) = shard.get_local_sst_files();
                    if sst_ia_file_ids.is_empty() {
                        return false;
                    }
                    shard_with_ia_count += 1;
                }
            }
            shard_with_ia_count == ia_file_tables.len()
        },
        20,
        || "failed to wait storage class ia".to_string(),
    );
    client.verify_data_with_ref_store();

    // Test regions that cannot be merged
    // ( t0 Local / t1 IA / t2 IA / t3 Local )
    // None of the two adjacent tables can be merged.
    let regions_count = pd_client.get_regions_number();
    for idx in 0..table_ids.len() - 1 {
        client.try_merge(&table_start_keys[idx], &table_start_keys[idx + 1]);
    }
    std::thread::sleep(Duration::from_millis(500));
    assert_eq!(pd_client.get_regions_number(), regions_count);
    for &table_start in &ia_table_start_keys {
        let shard_id = pd_client
            .get_region(&encode_bytes(table_start))
            .unwrap()
            .get_id();
        let shard = kvengine.get_shard(shard_id).unwrap();
        assert!(shard.get_schema_file().is_some() && shard.storage_class_spec_equals(&sc_spec_ia));
        check_ia_shard_range(shard.range.data_bound());
        let (_, sst_ia_file_ids) = shard.get_local_sst_files();
        assert!(!sst_ia_file_ids.is_empty());
    }
    client.verify_data_with_ref_store();

    // Split IA region
    let mut split_keys = Vec::with_capacity(ia_file_tables.len());
    for &table_id in &ia_file_tables {
        let regions_count = pd_client.get_regions_number();
        let split_key = gen_row_key(keyspace_id, table_id, 100);
        client.split(&split_key);
        split_keys.push(split_key);
        cluster.wait_pd_region_count(regions_count + 1);
        client.verify_data_with_ref_store();
    }

    // Create IA files for new IA regions
    for &table_id in &ia_file_tables {
        client.put_kv(
            100..150,
            |i: usize| gen_row_key(keyspace_id, table_id, i),
            |i: usize| gen_row_val(&ctx, i),
        );
        client.put_kv(
            150..200,
            |i: usize| gen_row_key(keyspace_id, table_id, i),
            |i: usize| gen_row_val(&ctx, i),
        );
    }
    must_wait(
        || {
            let all_id_vers = kvengine.get_all_shard_id_vers();
            let mut shard_with_ia_count = 0;
            for id_ver in all_id_vers {
                if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                    if shard.get_schema_file().is_some()
                        && shard.storage_class_spec_equals(&sc_spec_ia)
                    {
                        check_ia_shard_range(shard.range.data_bound());
                        let (_, sst_ia_file_ids) = shard.get_local_sst_files();
                        if sst_ia_file_ids.is_empty() {
                            return false;
                        }
                        shard_with_ia_count += 1;
                    }
                }
            }
            shard_with_ia_count == ia_file_tables.len() * 2
        },
        10,
        || "failed to load new ia file for new ia region".to_string(),
    );
    client.verify_data_with_ref_store();

    // Merge two adjacent IA regions of the same table.
    for idx in 0..ia_file_tables.len() {
        let table_id = ia_file_tables[idx];
        let split_key = &split_keys[idx];
        let regions_count = pd_client.get_regions_number();
        client.try_merge(&gen_row_key(keyspace_id, table_id, 0), split_key);
        cluster.wait_pd_region_count(regions_count - 1);
    }
    must_wait(
        || {
            let all_id_vers = kvengine.get_all_shard_id_vers();
            let mut shard_with_ia_count = 0;
            for id_ver in all_id_vers {
                if let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) {
                    if shard.get_schema_file().is_some()
                        && shard.storage_class_spec_equals(&sc_spec_ia)
                    {
                        check_ia_shard_range(shard.range.data_bound());
                        let (_, sst_ia_file_ids) = shard.get_local_sst_files();
                        if sst_ia_file_ids.is_empty() {
                            return false;
                        }
                        shard_with_ia_count += 1;
                    }
                }
            }
            shard_with_ia_count == ia_file_tables.len()
        },
        10,
        || "failed to merge two adjacent IA regions of the same table".to_string(),
    );
    client.verify_data_with_ref_store();

    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_global_index_storage_class_follows_table_level_with_schema_manager() {
    let node_id = alloc_node_id();
    let (_temp_dir, mut oss, dfs_config) =
        prepare_dfs("test_global_index_storage_class_follows_table_level_with_schema_manager");
    let security_config = new_security_config();
    let pd_wrapper = PdWrapper::new_test(1, &security_config, None);
    let mut cluster = ServerClusterBuilder::new(vec![node_id], |_, conf: &mut TikvConfig| {
        conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(1);
        conf.rocksdb.writecf.block_size = ReadableSize(512);
        conf.rocksdb.writecf.target_file_size_base = ReadableSize::kb(2);
        conf.coprocessor.region_split_size = ReadableSize::kb(512);
        conf.coprocessor.region_bucket_size = ReadableSize::kb(64);
        conf.storage.block_cache.capacity = Some(ReadableSize::kb(16));
        conf.dfs = dfs_config.clone();
        conf.kvengine.ia.segment_size = SEGMENT_SIZE;
        conf.kvengine.ia.freq_update_interval = ReadableDuration(FREQ_UPDATE_INTERVAL);
        conf.kvengine.ia.mem_cap = AbsoluteOrPercentSize::Abs(ReadableSize::kb(64));
        conf.kvengine.ia.disk_cap = AbsoluteOrPercentSize::Abs(ReadableSize::kb(1024));
        conf.kvengine.block_cache_type = BlockCacheType::Quick;
        conf.security = security_config.clone();
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_schema_manager(alloc_node_id());
    let dfs = cluster.get_dfs().unwrap();
    let keyspace_id = 9;
    let db_id = 5;
    let mut table_ids = dfs
        .get_runtime()
        .block_on(create_keyspace_region_without_splitting_table(
            &mut cluster,
            keyspace_id,
            2,
        ));
    table_ids.sort();
    assert_eq!(table_ids.len(), 2);
    // Table IDs are allocated from a process-global test allocator and may not
    // be adjacent when this test runs with other cloud_engine tests.
    let partition_id = table_ids[0];
    let parent_table_id = table_ids[1];

    let kvengine = cluster.get_kvengine(node_id);
    assert!(kvengine.is_ia_enabled());
    let mut client = cluster.new_client();
    let pd_client = cluster.get_pd_client();
    let ctx = Mutex::new(EvalContext::default());

    let mut tables = None;
    let mut partitions = Some(vec![partition_id]);
    let mut schema_version = 10;
    let mut update_schema = |client: &mut ClusterClient, sc_spec: StorageClassSpec| {
        let kv_pairs = generate_storage_class_schema_data_for_test(
            keyspace_id,
            db_id,
            parent_table_id,
            None,
            partitions.take(),
            schema_version,
            sc_spec,
            &mut tables,
        )
        .unwrap();
        put_kv_pairs(client, kv_pairs);
        schema_version += 1;
    };

    update_schema(&mut client, StorageClassSpec::default());

    // Parent table ID is used by global index data; partition ID is used by row
    // data.
    client.put_kv(
        0..100,
        |i: usize| gen_global_index_key(keyspace_id, parent_table_id, i),
        |i: usize| gen_row_val(&ctx, i),
    );
    client.put_kv(
        0..100,
        |i: usize| gen_row_key(keyspace_id, partition_id, i),
        |i: usize| gen_row_val(&ctx, i),
    );
    client.verify_data_with_ref_store();
    must_wait(
        || {
            kvengine
                .get_all_shard_id_vers()
                .into_iter()
                .filter_map(|id_ver| kvengine.get_shard_with_ver(id_ver.id, id_ver.ver).ok())
                .any(|shard| shard.get_stats().cfs[0].levels[0].num_tables > 0)
        },
        10,
        || "failed to get ln files for global index storage class test".to_string(),
    );

    let sc_spec_ia = StorageClassSpec::from(StorageClass::Ia);
    update_schema(&mut client, sc_spec_ia.clone());

    let parent_start_key = encode_table_prefix_key(parent_table_id);
    let parent_end_key = encode_table_prefix_key(parent_table_id + 1);
    let parent_bound = DataBound::new(parent_start_key.as_ref(), parent_end_key.as_ref(), false);
    let global_index_key = gen_global_index_key(keyspace_id, parent_table_id, 0);
    must_wait(
        || {
            let shard_id = pd_client
                .get_region(&encode_bytes(&global_index_key))
                .unwrap()
                .get_id();
            let shard = kvengine.get_shard(shard_id).unwrap();
            if shard.get_schema_file().is_none() || !shard.storage_class_spec_equals(&sc_spec_ia) {
                return false;
            }
            assert!(parent_bound.contains_bound(shard.range.data_bound()));
            let (_, sst_ia_file_ids) = shard.get_local_sst_files();
            !sst_ia_file_ids.is_empty()
        },
        20,
        || "failed to apply table-level IA to global index shard".to_string(),
    );

    let partition_key = gen_row_key(keyspace_id, partition_id, 0);
    let partition_shard_id = pd_client
        .get_region(&encode_bytes(&partition_key))
        .unwrap()
        .get_id();
    let partition_shard = kvengine.get_shard(partition_shard_id).unwrap();
    assert!(partition_shard.storage_class_spec_equals(&StorageClassSpec::default()));
    let (_, partition_ia_file_ids) = partition_shard.get_local_sst_files();
    assert!(partition_ia_file_ids.is_empty());
    client.verify_data_with_ref_store();

    cluster.stop();
    oss.shutdown();
}

fn put_kv_pairs(client: &mut ClusterClient, kv_pairs: HashMap<Vec<u8>, Vec<u8>>) {
    let mut muts: Vec<Mutation> = Vec::with_capacity(kv_pairs.len());
    for (key, value) in kv_pairs {
        let mut m = Mutation::default();
        m.key = Bytes::from(key);
        m.value = Bytes::from(value);
        m.op = kvrpcpb::Op::Put;
        muts.push(m);
    }
    client.kv_mutate(muts).unwrap();
}

fn gen_global_index_key(keyspace_id: u32, table_id: i64, i: usize) -> Vec<u8> {
    let mut user_key = Vec::new();
    user_key.encode_i64(1).unwrap();
    user_key.encode_i64(i as i64).unwrap();
    make_index_key(keyspace_id, table_id, &user_key)
}

fn collect_local_sst_files_by_type(kvengine: &kvengine::Engine) -> (HashSet<u64>, HashSet<u64>) {
    let mut local_files = HashSet::new();
    let mut ia_files = HashSet::new();
    for id_ver in kvengine.get_all_shard_id_vers() {
        let Ok(shard) = kvengine.get_shard_with_ver(id_ver.id, id_ver.ver) else {
            continue;
        };
        let (shard_local_files, shard_ia_files) = shard.get_local_sst_files();
        local_files.extend(shard_local_files);
        ia_files.extend(shard_ia_files);
    }
    (local_files, ia_files)
}

async fn request_files_with_type(status_addr: &str, filter: &str) -> HashMap<String, HashSet<u64>> {
    let uri = format!("http://{status_addr}/kvengine/files_with_type?filter={filter}")
        .parse::<hyper::Uri>()
        .unwrap();
    let req = Request::get(uri).body(Body::empty()).unwrap();
    let resp = hyper::Client::new().request(req).await.unwrap();
    assert!(resp.status().is_success());
    let body = hyper::body::to_bytes(resp.into_body()).await.unwrap();
    serde_json::from_slice(&body).unwrap()
}

async fn request_storage_class_status(
    status_addr: &str,
    keyspace_id: u32,
    table_id: i64,
    target: &str,
) -> StorageClassStatusResp {
    let resp =
        request_storage_class_status_response(status_addr, keyspace_id, table_id, Some(target))
            .await;
    assert_eq!(resp.status(), StatusCode::OK);
    let body = hyper::body::to_bytes(resp.into_body()).await.unwrap();
    serde_json::from_slice(&body).unwrap()
}

async fn request_storage_class_status_response(
    status_addr: &str,
    keyspace_id: u32,
    table_id: i64,
    target: Option<&str>,
) -> hyper::Response<Body> {
    let target = target
        .map(|target| format!("&target={target}"))
        .unwrap_or_default();
    let uri = format!(
        "http://{status_addr}/kvengine/storage_class_status?keyspace_id={keyspace_id}&table_id={table_id}{target}"
    )
    .parse::<hyper::Uri>()
    .unwrap();
    let req = Request::get(uri).body(Body::empty()).unwrap();
    hyper::Client::new().request(req).await.unwrap()
}

async fn create_keyspace_region_without_splitting_table(
    cluster: &mut ServerCluster,
    keyspace_id: u32,
    table_count: usize,
) -> Vec<i64> /* table_ids */
{
    cluster
        .keyspace_manager()
        .create_single_keyspace(
            keyspace_id,
            format!("ks{}", keyspace_id),
            &CreateKeyspaceOptions {
                table_count,
                ..Default::default()
            },
            false,
        )
        .await;
    let keyspace_split_keys = get_keyspace_split_keys(keyspace_id);
    cluster
        .get_pd_client()
        .split_regions_with_retry(keyspace_split_keys, Duration::from_secs(10))
        .await
        .unwrap();
    let table_ids = cluster
        .keyspace_manager()
        .get_keyspace_meta(keyspace_id)
        .unwrap()
        .get_all_available_tables();
    table_ids
}