use std::cmp::Ordering;

use anyhow::{Result, anyhow};
use hex;
use log::{debug, info, warn};
use storage_backend::{backend::BackendOptions, storage::gcs_storage};
use tici_common::{codec::table_end_key, config::S3Config, index::IndexKey, keyspace::NULLSPACEID};
use tici_grpc::{
    ErrorCode, FinishImportPartitionBuildRequest, GetIndexProgressResponse,
    PreSplitImportIndexResult, PreSplitImportShardMeta, PreSplitImportShardsRequest,
    PreSplitImportShardsResponse, WriterProgressState, prelude_import::*,
};
use tici_import::{ImportTaskParams, ImportTaskSummary, build_sharding_key_from_import_key};
use url::form_urlencoded;

use crate::meta::{
    MetaServiceCore,
    metrics::{
        IMPORT_INDEX_UPLOAD_COUNTER, IMPORT_PARTITION_BUILD_COUNTER,
        IMPORT_PARTITION_UPLOAD_COUNTER,
    },
    service::ImportJobMeta,
    writer::{ExplicitSplitReason, ShardTaskHeader},
};

// Start slightly above the observed imported index KV size so pre-split stays
// on the safe side. This can be tuned later with more import/build samples.
const IMPORT_PRESPLIT_SIZE_RATIO: f64 = 1.1;

fn derive_writer_progress_state(
    job_status: &str,
    total_tasks: i64,
    running_tasks: i64,
    finished_tasks: i64,
    failed_tasks: i64,
) -> WriterProgressState {
    if failed_tasks > 0 || job_status == "failed" {
        WriterProgressState::Failed
    } else if job_status == "finished" || (total_tasks > 0 && finished_tasks == total_tasks) {
        WriterProgressState::Completed
    } else if running_tasks > 0 || job_status == "running" {
        WriterProgressState::Running
    } else {
        WriterProgressState::Pending
    }
}

fn completed_response_for_missing_import_job() -> GetIndexProgressResponse {
    GetIndexProgressResponse {
        status: ErrorCode::Success.into(),
        error_message: String::new(),
        // Tradeoff: treat missing import job as completed so empty-table index
        // creation does not block on import workflow. This may be a false-positive
        // if caller skips/forgets import APIs; caller-side contract must guarantee
        // import jobs are created for non-empty table imports.
        state: WriterProgressState::Completed.into(),
    }
}

fn build_import_storage_prefix(
    cfg: &S3Config,
    keyspace_id: u32,
    table_id: i64,
    index_ids: &[i64],
    job_id: u64,
) -> Result<String> {
    let input_key =
        tici_common::s3_keys::import_input_prefix(keyspace_id, table_id, index_ids, job_id);
    let url = if gcs_storage::is_gcs_endpoint(&cfg.endpoint) {
        tici_common::s3_keys::to_full_path("gs", &cfg.bucket, &cfg.prefix, &input_key)
    } else {
        tici_common::s3_keys::to_full_path("s3", &cfg.bucket, &cfg.prefix, &input_key)
    };
    let options = BackendOptions::from_config(cfg)?;
    let mut q = "".to_string();
    for (k, v) in options.to_query_iter() {
        q = format!(
            "{}{}{}={}",
            q,
            if q.is_empty() { "" } else { "&" },
            k,
            form_urlencoded::byte_serialize(v.as_bytes()).collect::<String>()
        );
    }
    Ok(format!(
        "{}{}{}",
        url,
        if q.is_empty() { "" } else { "?" },
        q
    ))
}

#[derive(Clone, Copy, Debug)]
struct PresplitGroupSizeEstimator {
    size_ratio: f64,
}

impl PresplitGroupSizeEstimator {
    /// Builds an estimator for one merged meta group.
    ///
    /// `meta_groups.total_kv_size` already reflects the KV produced by
    /// `process_kv_pair`, which is an index KV rather than the original table
    /// row. That means the payload has already been projected to the columns
    /// that TiCI will materialize for the target index, so pre-split only needs
    /// a tunable ratio on top of the observed group size.
    ///
    /// Keep the ratio behind an abstraction so we can adjust it later if the
    /// imported index KV size is found to systematically under/over-estimate
    /// the final shard size.
    fn new(size_ratio: f64) -> Self {
        Self { size_ratio }
    }

    /// Estimates how many bytes one merged meta group should contribute to the
    /// pre-split accumulation.
    fn estimate_group_size(self, group: &PreSplitImportShardMeta) -> u64 {
        let estimated_size = (group.total_kv_size as f64) * self.size_ratio;
        estimated_size.ceil().clamp(0.0, u64::MAX as f64) as u64
    }
}

fn compare_meta_groups(a: &PreSplitImportShardMeta, b: &PreSplitImportShardMeta) -> Ordering {
    a.start_key
        .cmp(&b.start_key)
        .then_with(|| a.end_key.cmp(&b.end_key))
        .then_with(|| a.ele_id.cmp(&b.ele_id))
}

fn derive_presplit_split_source_keys<F>(
    meta_groups: &[PreSplitImportShardMeta],
    start_key: &[u8],
    end_key: &[u8],
    max_shard_size: u64,
    estimate_group_size: F,
) -> Vec<Vec<u8>>
where
    F: FnMut(&PreSplitImportShardMeta) -> u64,
{
    if meta_groups.is_empty() || max_shard_size == 0 {
        return vec![];
    }

    let mut meta_groups = meta_groups.to_vec();
    meta_groups.sort_by(compare_meta_groups);

    let mut split_keys = Vec::new();
    let mut accumulated_size = 0_u64;
    let mut estimate_group_size = estimate_group_size;
    for (group_idx, group) in meta_groups.iter().enumerate() {
        let estimated_size = estimate_group_size(group);
        if estimated_size == 0 {
            continue;
        }
        accumulated_size = accumulated_size.saturating_add(estimated_size);
        if accumulated_size < max_shard_size {
            continue;
        }
        if let Some(next_group) = meta_groups.get(group_idx + 1) {
            if next_group.start_key.as_slice() > start_key
                && next_group.start_key.as_slice() < end_key
                && split_keys.last() != Some(&next_group.start_key)
            {
                split_keys.push(next_group.start_key.clone());
            }
        }
        accumulated_size = 0;
    }
    split_keys
}

#[derive(Debug, PartialEq, Eq)]
struct PresplitKeyConversionFailure {
    key: Vec<u8>,
    error: String,
}

fn convert_presplit_keys_best_effort<F>(
    source_keys: &[Vec<u8>],
    mut convert: F,
) -> (Vec<Vec<u8>>, Vec<PresplitKeyConversionFailure>)
where
    F: FnMut(&[u8]) -> Result<Vec<u8>>,
{
    let mut converted = Vec::with_capacity(source_keys.len());
    let mut failures = Vec::new();
    for key in source_keys {
        match convert(key) {
            Ok(converted_key) => converted.push(converted_key),
            Err(err) => failures.push(PresplitKeyConversionFailure {
                key: key.clone(),
                error: err.to_string(),
            }),
        }
    }
    (converted, failures)
}

fn presplit_failure_details(failures: &[PresplitKeyConversionFailure]) -> Vec<String> {
    failures
        .iter()
        .map(|failure| format!("{} ({})", hex::encode(&failure.key), failure.error))
        .collect()
}

impl MetaServiceCore {
    pub async fn get_import_storage_prefix(
        &self,
        req: &GetImportStoragePrefixRequest,
    ) -> Result<GetImportStoragePrefixResponse> {
        if req.table_id == 0 || req.index_ids.is_empty() {
            return Ok(GetImportStoragePrefixResponse {
                status: tici_grpc::ErrorCode::InvalidArguments.into(),
                error_message: "failed to start import index, table_id is 0 or index_ids is empty"
                    .to_string(),
                ..Default::default()
            });
        }

        let job_id = self
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: req.tidb_task_id.clone(),
                keyspace_id: req.keyspace_id,
                table_id: req.table_id,
                index_ids: req.index_ids.clone(),
                job_id: 0, // leave it to meta store to assign
            })
            .await?;

        // TODO: refine the logic of generating storage prefix by different storage
        // backends.
        let s3_config = self.config.load().s3.clone();
        let input_prefix = build_import_storage_prefix(
            &s3_config,
            req.keyspace_id,
            req.table_id,
            &req.index_ids,
            job_id,
        )?;

        Ok(GetImportStoragePrefixResponse {
            status: tici_grpc::ErrorCode::Success.into(),
            error_message: String::new(),
            job_id,
            storage_uri: input_prefix,
        })
    }

    pub async fn presplit_import_shards(
        &self,
        req: &PreSplitImportShardsRequest,
    ) -> Result<PreSplitImportShardsResponse> {
        if req.table_id == 0 || req.index_ids.is_empty() {
            return Ok(PreSplitImportShardsResponse {
                status: ErrorCode::InvalidArguments.into(),
                error_message:
                    "failed to presplit import shards, table_id is 0 or index_ids is empty"
                        .to_string(),
                ..Default::default()
            });
        }
        if req.start_key.is_empty() || req.end_key.is_empty() || req.start_key >= req.end_key {
            return Ok(PreSplitImportShardsResponse {
                status: ErrorCode::InvalidArguments.into(),
                error_message: "failed to presplit import shards, key range is invalid".to_string(),
                ..Default::default()
            });
        }

        self.meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: req.tidb_task_id.clone(),
                keyspace_id: req.keyspace_id,
                table_id: req.table_id,
                index_ids: req.index_ids.clone(),
                job_id: 0,
            })
            .await?;

        let size_estimator = PresplitGroupSizeEstimator::new(IMPORT_PRESPLIT_SIZE_RATIO);
        let split_source_keys = derive_presplit_split_source_keys(
            &req.meta_groups,
            &req.start_key,
            &req.end_key,
            self.config.load().shard.max_size.0,
            |group| size_estimator.estimate_group_size(group),
        );
        let mut index_keys = Vec::new();
        for index_id in &req.index_ids {
            let mut partition_index_keys = self
                .shard_task_manager
                .collect_partition_index_keys(req.keyspace_id, req.table_id, *index_id)
                .await;
            if partition_index_keys.is_empty() {
                partition_index_keys.push(IndexKey::new(req.keyspace_id, req.table_id, *index_id));
            }
            index_keys.extend(partition_index_keys);
        }
        index_keys.sort_by_key(|index_key| {
            (
                index_key.keyspace_id,
                index_key.table_id,
                index_key.index_id,
            )
        });
        index_keys.dedup();
        self.shard_task_manager
            .set_import_topology_frozen(req.keyspace_id, req.table_id, &req.index_ids, true)
            .await;

        let mut index_results = Vec::with_capacity(index_keys.len());
        for index_key in index_keys {
            let index = self.shard_task_manager.get_index(&index_key).await?;
            let request_start_key = match build_sharding_key_from_import_key(
                req.keyspace_id,
                &index.table_info,
                &index.index_info,
                &req.start_key,
            ) {
                Ok(key) => key,
                Err(err) => {
                    warn!(
                        "presplit_import_shards: skip index due to invalid start key, keyspace_id={} table_id={} index_id={} key={} err={:#}",
                        req.keyspace_id,
                        index_key.table_id,
                        index_key.index_id,
                        hex::encode(&req.start_key),
                        err
                    );
                    info!(
                        "presplit_import_shards: index summary keyspace_id={} table_id={} index_id={} candidate_split_keys={} accepted_split_keys=0 ignored_split_keys=0 unique_split_keys=0 split_count=0 skipped=true",
                        req.keyspace_id,
                        index_key.table_id,
                        index_key.index_id,
                        split_source_keys.len(),
                    );
                    index_results.push(PreSplitImportIndexResult {
                        table_id: index_key.table_id,
                        index_id: index_key.index_id,
                        split_count: 0,
                        shard_count_after_split: index.shards_len().await as u32,
                    });
                    continue;
                }
            };
            let (request_end_key, end_key_fallback) = match build_sharding_key_from_import_key(
                req.keyspace_id,
                &index.table_info,
                &index.index_info,
                &req.end_key,
            ) {
                Ok(key) => (key, false),
                Err(err) => {
                    let fallback_end_key = table_end_key(&index_key);
                    warn!(
                        "presplit_import_shards: fallback to table end key due to invalid end key, keyspace_id={} table_id={} index_id={} key={} fallback_end_key={} err={:#}",
                        req.keyspace_id,
                        index_key.table_id,
                        index_key.index_id,
                        hex::encode(&req.end_key),
                        hex::encode(&fallback_end_key),
                        err
                    );
                    (fallback_end_key, true)
                }
            };
            let (mut split_keys, ignored_split_keys) =
                convert_presplit_keys_best_effort(&split_source_keys, |source_key| {
                    build_sharding_key_from_import_key(
                        req.keyspace_id,
                        &index.table_info,
                        &index.index_info,
                        source_key,
                    )
                });
            if !ignored_split_keys.is_empty() {
                warn!(
                    "presplit_import_shards: ignore split source keys, keyspace_id={} table_id={} index_id={} accepted={} ignored={} failed_keys={:?}",
                    req.keyspace_id,
                    index_key.table_id,
                    index_key.index_id,
                    split_keys.len(),
                    ignored_split_keys.len(),
                    presplit_failure_details(&ignored_split_keys)
                );
            }
            let converted_split_key_count = split_keys.len();
            split_keys.sort();
            split_keys.dedup();
            info!(
                "presplit_import_shards: converted split source keys, keyspace_id={} table_id={} index_id={} accepted={} ignored={} unique={}",
                req.keyspace_id,
                index_key.table_id,
                index_key.index_id,
                converted_split_key_count,
                ignored_split_keys.len(),
                split_keys.len()
            );
            let overlapping_shards = index
                .find_overlapping_shards(
                    &[tici_grpc::KeyRange {
                        start_key: request_start_key,
                        end_key: request_end_key,
                    }],
                    false,
                )
                .await;
            let mut split_rxs = Vec::new();
            for shard in overlapping_shards {
                let shard = shard.read().await.clone();
                let local_split_keys = split_keys
                    .iter()
                    .filter(|split_key| {
                        split_key.as_slice() > shard.range.encoded_start().as_slice()
                            && split_key.as_slice() < shard.range.encoded_end().as_slice()
                    })
                    .cloned()
                    .collect::<Vec<_>>();
                if local_split_keys.is_empty() {
                    continue;
                }
                split_rxs.push(self.shard_task_manager.schedule_explicit_split(
                    ShardTaskHeader::from(&shard),
                    index.clone(),
                    local_split_keys,
                    ExplicitSplitReason::PresplitImport,
                )?);
            }
            let mut split_count = 0_u32;
            for rx in split_rxs {
                let result = rx.await.map_err(|e| {
                    anyhow!(
                        "presplit_import_shards: explicit split task canceled, index_key={:?}, err={:?}",
                        index_key,
                        e
                    )
                })??;
                split_count = split_count.saturating_add(result.split_count);
            }
            info!(
                "presplit_import_shards: index summary keyspace_id={} table_id={} index_id={} candidate_split_keys={} accepted_split_keys={} ignored_split_keys={} unique_split_keys={} split_count={} skipped=false end_key_fallback={}",
                req.keyspace_id,
                index_key.table_id,
                index_key.index_id,
                split_source_keys.len(),
                converted_split_key_count,
                ignored_split_keys.len(),
                split_keys.len(),
                split_count,
                end_key_fallback,
            );
            index_results.push(PreSplitImportIndexResult {
                table_id: index_key.table_id,
                index_id: index_key.index_id,
                split_count,
                shard_count_after_split: index.shards_len().await as u32,
            });
        }

        Ok(PreSplitImportShardsResponse {
            status: ErrorCode::Success.into(),
            error_message: String::new(),
            // Top-level split keys are the source import keys so the field stays
            // index-agnostic even when one request targets multiple indexes.
            split_keys: split_source_keys,
            index_results,
        })
    }

    pub async fn finish_partition_upload(
        &self,
        req: &FinishImportPartitionUploadRequest,
    ) -> Result<FinishImportResponse> {
        if req.key_range.is_none() {
            IMPORT_PARTITION_UPLOAD_COUNTER
                .with_label_values(&["invalid_key_range"])
                .inc();
            return Ok(FinishImportResponse {
                status: tici_grpc::ErrorCode::InvalidArguments.into(),
                error_message: "failed to mark partition upload finished, key_range is None"
                    .to_string(),
            });
        }
        let key_range = req
            .key_range
            .as_ref()
            .ok_or(anyhow!("keyrange is empty, {:?}", req))?;
        if key_range.start_key.is_empty() && key_range.end_key.is_empty() {
            IMPORT_PARTITION_UPLOAD_COUNTER
                .with_label_values(&["empty_key_range"])
                .inc();
            return Ok(FinishImportResponse {
                status: tici_grpc::ErrorCode::InvalidArguments.into(),
                error_message: "failed to mark partition upload finished, key_range is empty"
                    .to_string(),
            });
        }
        if req.storage_uri.is_empty() {
            IMPORT_PARTITION_UPLOAD_COUNTER
                .with_label_values(&["empty_storage_uri"])
                .inc();
            return Ok(FinishImportResponse {
                status: tici_grpc::ErrorCode::InvalidArguments.into(),
                error_message: "failed to mark partition upload finished, storage_uri is empty"
                    .to_string(),
            });
        }

        if let Err(e) = self
            .meta_store
            .add_import_task(
                &req.tidb_task_id,
                ImportTaskParams::new(
                    hex::encode(&key_range.start_key),
                    hex::encode(&key_range.end_key),
                    req.storage_uri.clone(),
                )
                .with_index_id(req.index_id),
            )
            .await
        {
            IMPORT_PARTITION_UPLOAD_COUNTER
                .with_label_values(&["store_error"])
                .inc();
            return Err(e);
        }
        IMPORT_PARTITION_UPLOAD_COUNTER
            .with_label_values(&["success"])
            .inc();
        Ok(FinishImportResponse {
            status: tici_grpc::ErrorCode::Success.into(),
            error_message: String::new(),
        })
    }

    pub async fn finish_partition_build(
        &self,
        req: &FinishImportPartitionBuildRequest,
    ) -> Result<FinishImportResponse> {
        match req.status {
            0 => {
                let summary = ImportTaskSummary {
                    rows_processed: req.rows_processed,
                    bytes_processed: req.bytes_processed,
                    input_file_size: req.input_file_size,
                };
                let updated = match self
                    .meta_store
                    .import_task_build_success(req.task_id, &req.owner_addr, req.owner_seq, summary)
                    .await
                {
                    Ok(updated) => updated,
                    Err(e) => {
                        IMPORT_PARTITION_BUILD_COUNTER
                            .with_label_values(&["store_error"])
                            .inc();
                        return Err(e);
                    }
                };
                if updated {
                    let _ = self
                        .worker_node_mgr
                        .on_import_task_finished(&req.owner_addr, req.task_id)
                        .await;
                }
                IMPORT_PARTITION_BUILD_COUNTER
                    .with_label_values(&["success"])
                    .inc();
            }
            _ => {
                IMPORT_PARTITION_BUILD_COUNTER
                    .with_label_values(&["failed"])
                    .inc();
                let updated = self
                    .meta_store
                    .import_task_build_fail(
                        req.task_id,
                        &req.owner_addr,
                        req.owner_seq,
                        &req.error_message,
                    )
                    .await?;
                if updated {
                    let _ = self
                        .worker_node_mgr
                        .on_import_task_finished(&req.owner_addr, req.task_id)
                        .await;
                }
            }
        }
        Ok(FinishImportResponse {
            status: tici_grpc::ErrorCode::Success.into(),
            error_message: String::new(),
        })
    }

    pub async fn mark_table_upload_finished(
        &self,
        req: &FinishImportIndexUploadRequest,
    ) -> Result<FinishImportResponse> {
        if let Err(e) = self
            .meta_store
            .finish_import_job_upload(&req.tidb_task_id)
            .await
        {
            IMPORT_INDEX_UPLOAD_COUNTER
                .with_label_values(&["store_error"])
                .inc();
            return Err(e);
        }
        IMPORT_INDEX_UPLOAD_COUNTER
            .with_label_values(&["success"])
            .inc();
        Ok(FinishImportResponse {
            status: tici_grpc::ErrorCode::Success.into(),
            error_message: String::new(),
        })
    }

    pub async fn get_index_progress(
        &self,
        req: &tici_grpc::GetIndexProgressRequest,
    ) -> Result<tici_grpc::GetIndexProgressResponse> {
        debug!("GetIndexProgress called, request={:?}", req);
        if req.table_id == 0 || req.index_id == 0 {
            return Ok(GetIndexProgressResponse {
                status: ErrorCode::InvalidArguments.into(),
                error_message: "table_id or index_id is 0".to_string(),
                state: WriterProgressState::Error.into(),
            });
        }

        let keyspace_id = if req.keyspace_id == 0 {
            NULLSPACEID
        } else {
            req.keyspace_id
        };
        let progress = self
            .meta_store
            .get_import_index_progress(keyspace_id, req.table_id, req.index_id)
            .await?;

        let Some(progress) = progress else {
            return Ok(completed_response_for_missing_import_job());
        };

        let total_tasks = progress.task_count;
        let running_tasks = progress.running_task_count;
        let finished_tasks = progress.finished_task_count;
        let failed_tasks = progress.failed_task_count;
        let state = derive_writer_progress_state(
            &progress.job_status,
            total_tasks,
            running_tasks,
            finished_tasks,
            failed_tasks,
        );

        Ok(GetIndexProgressResponse {
            status: ErrorCode::Success.into(),
            error_message: String::new(),
            state: state.into(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_derive_writer_progress_state_pending_for_init_with_only_init_tasks() {
        let state = derive_writer_progress_state("init", 2, 0, 0, 0);
        assert_eq!(state, WriterProgressState::Pending);
    }

    #[test]
    fn test_derive_presplit_split_source_keys_accumulates_by_group_size() {
        let request = PreSplitImportShardsRequest {
            table_id: 1,
            index_ids: vec![2],
            start_key: b"a".to_vec(),
            end_key: b"z".to_vec(),
            meta_groups: vec![
                PreSplitImportShardMeta {
                    ele_id: 1,
                    start_key: b"a".to_vec(),
                    end_key: b"g".to_vec(),
                    total_kv_size: 6,
                    ..Default::default()
                },
                PreSplitImportShardMeta {
                    ele_id: 2,
                    start_key: b"g".to_vec(),
                    end_key: b"m".to_vec(),
                    total_kv_size: 5,
                    ..Default::default()
                },
                PreSplitImportShardMeta {
                    ele_id: 3,
                    start_key: b"m".to_vec(),
                    end_key: b"t".to_vec(),
                    total_kv_size: 4,
                    ..Default::default()
                },
                PreSplitImportShardMeta {
                    ele_id: 4,
                    start_key: b"t".to_vec(),
                    end_key: b"y".to_vec(),
                    total_kv_size: 8,
                    ..Default::default()
                },
                PreSplitImportShardMeta {
                    ele_id: 5,
                    start_key: b"y".to_vec(),
                    end_key: b"z".to_vec(),
                    total_kv_size: 1,
                    ..Default::default()
                },
            ],
            ..Default::default()
        };

        let split_keys = derive_presplit_split_source_keys(
            &request.meta_groups,
            &request.start_key,
            &request.end_key,
            10,
            |group| group.total_kv_size,
        );
        assert_eq!(split_keys, vec![b"m".to_vec(), b"y".to_vec()]);
    }

    #[test]
    fn test_derive_presplit_split_source_keys_filters_range_boundaries() {
        let request = PreSplitImportShardsRequest {
            table_id: 1,
            index_ids: vec![2],
            start_key: b"a".to_vec(),
            end_key: b"z".to_vec(),
            meta_groups: vec![
                PreSplitImportShardMeta {
                    ele_id: 1,
                    start_key: b"a".to_vec(),
                    end_key: b"a".to_vec(),
                    total_kv_size: 10,
                    ..Default::default()
                },
                PreSplitImportShardMeta {
                    ele_id: 2,
                    start_key: b"a".to_vec(),
                    end_key: b"m".to_vec(),
                    total_kv_size: 10,
                    ..Default::default()
                },
                PreSplitImportShardMeta {
                    ele_id: 3,
                    start_key: b"m".to_vec(),
                    end_key: b"z".to_vec(),
                    total_kv_size: 10,
                    ..Default::default()
                },
            ],
            ..Default::default()
        };

        let split_keys = derive_presplit_split_source_keys(
            &request.meta_groups,
            &request.start_key,
            &request.end_key,
            10,
            |group| group.total_kv_size,
        );
        assert_eq!(split_keys, vec![b"m".to_vec()]);
    }

    #[test]
    fn test_presplit_group_size_estimator_applies_ratio_to_index_kv_size() {
        let estimator = PresplitGroupSizeEstimator::new(0.5);
        let estimated = estimator.estimate_group_size(&PreSplitImportShardMeta {
            total_kv_size: 11,
            ..Default::default()
        });

        assert_eq!(estimated, 6);
    }

    #[test]
    fn test_convert_presplit_keys_best_effort_collects_failures() {
        let source_keys = vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()];
        let (converted, failures) = convert_presplit_keys_best_effort(&source_keys, |key| {
            if key == b"b" {
                Err(anyhow!("bad key"))
            } else {
                Ok([key, b"_ok"].concat())
            }
        });

        assert_eq!(converted, vec![b"a_ok".to_vec(), b"c_ok".to_vec()]);
        assert_eq!(
            failures,
            vec![PresplitKeyConversionFailure {
                key: b"b".to_vec(),
                error: "bad key".to_string(),
            }]
        );
    }

    #[test]
    fn test_derive_writer_progress_state_running_when_running_tasks_exist() {
        let state = derive_writer_progress_state("init", 2, 1, 0, 0);
        assert_eq!(state, WriterProgressState::Running);
    }

    #[test]
    fn test_derive_writer_progress_state_completed_for_finished_job() {
        let state = derive_writer_progress_state("finished", 2, 0, 0, 0);
        assert_eq!(state, WriterProgressState::Completed);
    }

    #[test]
    fn test_missing_import_job_treated_as_completed() {
        let resp = completed_response_for_missing_import_job();
        assert_eq!(resp.status, ErrorCode::Success as i32);
        assert_eq!(resp.state, WriterProgressState::Completed as i32);
        assert_eq!(resp.error_message, "");
    }

    #[test]
    fn test_build_import_storage_prefix() {
        struct TestCase {
            name: &'static str,
            s3_config: S3Config,
            keyspace_id: u32,
            table_id: i64,
            index_ids: Vec<i64>,
            job_id: u64,
            expected: &'static str,
        }

        let cases = vec![
            TestCase {
                name: "gcs endpoint",
                s3_config: S3Config {
                    endpoint: "https://storage.googleapis.com".to_string(),
                    bucket: "bucket-a".to_string(),
                    prefix: "prefix-a".to_string(),
                    ..Default::default()
                },
                keyspace_id: 7,
                table_id: 11,
                index_ids: vec![13],
                job_id: 17,
                expected: "gs://bucket-a/prefix-a/ks_7/t_11/i_13/import/17",
            },
            TestCase {
                name: "s3 endpoint without credentials",
                s3_config: S3Config {
                    endpoint: "https://s3.example.com/custom/path".to_string(),
                    bucket: "bucket-b".to_string(),
                    prefix: "prefix-b".to_string(),
                    use_path_style: false,
                    ..Default::default()
                },
                keyspace_id: 3,
                table_id: 5,
                index_ids: vec![9],
                job_id: 21,
                expected: "s3://bucket-b/prefix-b/ks_3/t_5/i_9/import/21?endpoint=https%3A%2F%2Fs3.example.com%2Fcustom%2Fpath&region=us-east-1&force-path-style=false",
            },
            TestCase {
                name: "s3 endpoint with credentials",
                s3_config: S3Config {
                    endpoint: "https://s3.us-east-1.amazonaws.com".to_string(),
                    bucket: "bucket-c".to_string(),
                    prefix: "prefix-c".to_string(),
                    access_key: "ak+/=".to_string(),
                    secret_key: "sk+/=x".to_string(),
                    use_path_style: false,
                    ..Default::default()
                },
                keyspace_id: NULLSPACEID,
                table_id: 2,
                index_ids: vec![4, 6],
                job_id: 8,
                expected: "s3://bucket-c/prefix-c/t_2/i_all/import/8?endpoint=https%3A%2F%2Fs3.us-east-1.amazonaws.com&region=us-east-1&access-key=ak%2B%2F%3D&secret-access-key=sk%2B%2F%3Dx&force-path-style=false",
            },
        ];

        for case in cases {
            let input_prefix = build_import_storage_prefix(
                &case.s3_config,
                case.keyspace_id,
                case.table_id,
                &case.index_ids,
                case.job_id,
            )
            .unwrap();
            assert_eq!(input_prefix, case.expected, "case: {}", case.name);
        }
    }
}