use std::{
    sync::Arc,
    time::{Duration, SystemTime},
};

use anyhow::{Context, Result, anyhow, bail};
use log::{error, info};
use tici_common::{cdc_file_name::CDCFileName, index::IndexKey};
use tici_grpc::{
    AppendBaseFragMetaRequest, AppendDeltaFragMetaRequest, AppendFragMetaResponse, ErrorCode,
    FragMeta, ShardProgress,
};
use tokio::{
    spawn,
    sync::mpsc::{Receiver, Sender},
    task::{self, JoinHandle},
};

use crate::{
    metadata::metadata_store::MetadataStore,
    utils::{
        meta_client::MetaClient,
        meta_client_pool::MetaClientPool,
        metrics::{WRITER_FRAG_OPERATION_COUNTER, WRITER_FRAG_OPERATION_LATENCY},
        s3_frag,
    },
    writer::{frag_batch_uploader::FragBatchUploader, frag_writer::FragWriter},
};

const MAX_RETRY_DELAY_SECS: u64 = 60;

#[derive(Debug)]
pub struct SubmitInfo {
    is_base: bool,
    frag_meta: Option<FragMeta>,
    last_cdc_file: Option<CDCFileName>,
    frag_path: String,
}

impl SubmitInfo {
    pub fn for_update_progress(last_cdc_file: CDCFileName, frag_path: String) -> Self {
        Self {
            is_base: false,
            frag_meta: None,
            last_cdc_file: Some(last_cdc_file),
            frag_path,
        }
    }
}

#[allow(clippy::large_enum_variant)]
pub enum FragUploadTask {
    Frag(FragWriter),
    Progress(SubmitInfo),
}

#[allow(clippy::large_enum_variant)]
enum FragSubmitTask {
    Frag(JoinHandle<Result<SubmitInfo>>),
    Progress(SubmitInfo),
}

// Receive -> Concurrent Upload -> Ordered Submit
// FragUploader orchestrates two asynchronous workflows:
// - recv_and_upload: Spawns a background job to persist the FragWriter's data
//   to S3 upon receiving the object.
// - wait_and_submit: Waits the front task in the uploading-queue to finish and
//   submits it's metadata to MetaService.
// Uploading multiple FragWriters to S3 can be performed concurrently, but
// submissions to MetaService must maintain strict ordering.
pub struct FragUploader {
    tracing_id: String,
    recv_and_upload_handle: JoinHandle<Result<()>>,
    wait_and_submit_handle: JoinHandle<Result<()>>,
}

impl FragUploader {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        tracing_id: String,
        index_key: &IndexKey,
        shard_id: u64,
        worker_addr: String,
        meta_client_pool: Arc<dyn MetaClientPool>,
        frag_batch_uploader: Arc<FragBatchUploader>,
        s3_bucket: String,
        s3_upload_prefix: String,
        frag_upload_task_receiver: Receiver<FragUploadTask>,
        remove_frag_if_uploaded: bool,
        metadata_store: Arc<MetadataStore>,
    ) -> Self {
        let tracing_id_clone = tracing_id.clone();
        let (submit_task_sender, submit_task_receiver) = tokio::sync::mpsc::channel(100);
        let recv_and_upload_handle = spawn(async move {
            Self::recv_and_upload(
                tracing_id_clone,
                frag_upload_task_receiver,
                frag_batch_uploader,
                s3_bucket,
                s3_upload_prefix,
                submit_task_sender,
                remove_frag_if_uploaded,
            )
            .await
        });

        let index_key_clone = index_key.clone();
        let tracing_id_clone = tracing_id.clone();
        let wait_and_submit_handle = spawn(async move {
            Self::wait_and_submit(
                tracing_id_clone,
                &index_key_clone,
                shard_id,
                worker_addr,
                meta_client_pool,
                submit_task_receiver,
                metadata_store,
            )
            .await
        });

        Self {
            tracing_id,
            recv_and_upload_handle,
            wait_and_submit_handle,
        }
    }

    pub async fn shutdown(self) {
        let tracing_id = self.tracing_id.clone();
        self.recv_and_upload_handle.abort();
        self.wait_and_submit_handle.abort();
        // Using abort to exit inevitably causes wait_finished to return a join error.
        // TODO: Optimize the error log.
        let _ = self.wait_finished().await.inspect_err(|e| {
            info!("wait_finished: {}, err={:?}", tracing_id, e);
        });
    }

    pub async fn wait_finished(self) -> Result<()> {
        let tracing_id = self.tracing_id.clone();
        info!("waitting recv_and_upload: {}", tracing_id);
        self.recv_and_upload_handle
            .await
            .with_context(|| format!("recv_and_upload {}", tracing_id))?
            .with_context(|| format!("recv_and_upload {}", tracing_id))?;
        info!("waitting wait_and_submit: {}", tracing_id);
        self.wait_and_submit_handle
            .await
            .with_context(|| format!("wait_and_submit {}", tracing_id))?
            .with_context(|| format!("wait_and_submit {}", tracing_id))?;
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    async fn recv_and_upload(
        tracing_id: String,
        mut frag_upload_task_receiver: Receiver<FragUploadTask>,
        frag_batch_uploader: Arc<FragBatchUploader>,
        s3_bucket: String,
        s3_upload_prefix: String,
        submit_task_sender: Sender<FragSubmitTask>,
        remove_frag_if_uploaded: bool,
    ) -> Result<()> {
        while let Some(frag_upload_task) = frag_upload_task_receiver.recv().await {
            match frag_upload_task {
                FragUploadTask::Frag(frag_writer) => {
                    let tracing_id_for_upload = tracing_id.clone();
                    let frag_batch_uploader = frag_batch_uploader.clone();
                    let s3_bucket = s3_bucket.clone();
                    let s3_upload_prefix = s3_upload_prefix.clone();
                    let upload_frag_handle = spawn(async move {
                        Self::upload_frag(
                            tracing_id_for_upload,
                            &frag_batch_uploader,
                            &s3_bucket,
                            &s3_upload_prefix,
                            frag_writer,
                            remove_frag_if_uploaded,
                        )
                        .await
                    });
                    // Send task to ordered submitter
                    submit_task_sender
                        .send(FragSubmitTask::Frag(upload_frag_handle))
                        .await
                        .inspect_err(|e| {
                            error!(
                                "send_submit_task_failed, {}, task_type=frag, err={:?}",
                                tracing_id, e
                            )
                        })?;
                }
                FragUploadTask::Progress(submit_info) => {
                    // Send progress to ordered submitter
                    submit_task_sender
                        .send(FragSubmitTask::Progress(submit_info))
                        .await
                        .inspect_err(|e| {
                            error!(
                                "send_submit_task_failed, {}, task_type=progress, err={:?}",
                                tracing_id, e
                            )
                        })?;
                }
            }
        }
        Ok(())
    }

    async fn upload_frag(
        tracing_id: String,
        frag_batch_uploader: &FragBatchUploader,
        s3_bucket: &str,
        s3_upload_prefix: &str,
        mut frag_writer: FragWriter,
        remove_frag_if_uploaded: bool,
    ) -> Result<SubmitInfo> {
        let mut submit_info = {
            SubmitInfo {
                is_base: frag_writer.is_base(),
                frag_meta: None,
                last_cdc_file: frag_writer.last_cdc_file().clone(),
                frag_path: frag_writer.frag_path(),
            }
        };
        let local_data_path = frag_writer.frag_base_path();
        let frag_res = task::spawn_blocking(move || {
            // The finalize operation mainly performs actions like writing local files and
            // waiting for threads to complete. If finalize failed, the
            // index_writer inside frag_writer has already been move and cannot be retried.
            let res = frag_writer.finalize();
            if res.is_err() {
                WRITER_FRAG_OPERATION_COUNTER
                    .with_label_values(&["flush_fail"])
                    .inc();
            } else {
                WRITER_FRAG_OPERATION_COUNTER
                    .with_label_values(&["flush_success"])
                    .inc();
            }
            res
        })
        .await??;

        // In theory, uploading to S3 shouldn't consistently fail. Unless there are some
        // configuration issues, such as permission problems, which would require manual
        // intervention to investigate.
        let mut retries = 0;
        let upload_start = SystemTime::now();
        let frag_meta = loop {
            let upload_result = s3_frag::upload_frag(
                frag_batch_uploader,
                s3_bucket,
                s3_upload_prefix,
                &local_data_path,
                &submit_info.frag_path,
                frag_res.handle_stat.clone(),
                &frag_res.seg_ids,
            )
            .await;

            match upload_result {
                Ok(frag_meta) => {
                    if remove_frag_if_uploaded {
                        if let Err(e) = std::fs::remove_dir_all(&submit_info.frag_path) {
                            error!(
                                "remove_local_frag_failed, {}, frag_path={}, err={:?}",
                                tracing_id, submit_info.frag_path, e
                            ); /* Ignore error */
                        }
                    }
                    break frag_meta;
                }
                Err(e) => {
                    WRITER_FRAG_OPERATION_COUNTER
                        .with_label_values(&["upload_fail"])
                        .inc();
                    retries += 1;
                    let delay_secs = retries.min(MAX_RETRY_DELAY_SECS);
                    error!(
                        "upload_frag_retry, {}, frag_path={}, retry_count={}, retry_delay_s={}, err={:?}",
                        tracing_id, submit_info.frag_path, retries, delay_secs, e
                    );
                    tokio::time::sleep(Duration::from_secs(delay_secs)).await;
                }
            }
        };
        let upload_elapsed_ms = upload_start.elapsed()?.as_millis();
        WRITER_FRAG_OPERATION_LATENCY
            .with_label_values(&["upload"])
            .observe(upload_start.elapsed()?.as_secs_f64());
        info!(
            "upload_frag_success, {}, frag_path={}, finalize_elapsed_ms={}, upload_elapsed_ms={}",
            tracing_id, submit_info.frag_path, frag_res.finalize_elapsed_ms, upload_elapsed_ms
        );
        submit_info.frag_meta = Some(frag_meta);
        Ok(submit_info)
    }

    async fn wait_and_submit(
        tracing_id: String,
        index_key: &IndexKey,
        shard_id: u64,
        worker_addr: String,
        meta_client_pool: Arc<dyn MetaClientPool>,
        mut submit_task_receiver: Receiver<FragSubmitTask>,
        metadata_store: Arc<MetadataStore>,
    ) -> Result<()> {
        while let Some(submit_task) = submit_task_receiver.recv().await {
            Self::wait_and_submit_front_task(
                tracing_id.clone(),
                index_key,
                shard_id,
                worker_addr.clone(),
                &meta_client_pool,
                submit_task,
                metadata_store.clone(),
            )
            .await?;
        }
        Ok(())
    }

    async fn wait_and_submit_front_task(
        tracing_id: String,
        index_key: &IndexKey,
        shard_id: u64,
        worker_addr: String,
        meta_client_pool: &Arc<dyn MetaClientPool>,
        submit_task: FragSubmitTask,
        metadata_store: Arc<MetadataStore>,
    ) -> Result<()> {
        let submit_info = match submit_task {
            FragSubmitTask::Frag(upload_task_handle) => upload_task_handle
                .await
                .inspect_err(|e| error!("upload_task_join_failed, {}, err={:?}", tracing_id, e))?
                .inspect_err(|e| error!("upload_task_failed, {}, err={:?}", tracing_id, e))?,
            FragSubmitTask::Progress(submit_info) => submit_info,
        };

        let start = SystemTime::now();
        let mut retries = 0;
        info!(
            "submit_frag_start, {}, frag_path={}",
            tracing_id, submit_info.frag_path
        );
        loop {
            match Self::submit_frag(
                meta_client_pool,
                index_key,
                shard_id,
                &worker_addr,
                &submit_info,
            )
            .await
            {
                Ok(_) => {
                    let elapsed = start.elapsed()?.as_millis();
                    WRITER_FRAG_OPERATION_LATENCY
                        .with_label_values(&["submit"])
                        .observe(start.elapsed()?.as_secs_f64());
                    info!(
                        "submit_frag_success, {}, submit_info={:?}, elapsed_ms={}",
                        tracing_id, submit_info, elapsed
                    );
                    if let Some(last_cdc_file) = submit_info.last_cdc_file.clone() {
                        // Update last cdc file in metadata store asynchronously.
                        // Ignore error.
                        let _ = metadata_store
                            .update_last_cdc_file(shard_id, last_cdc_file)
                            .await;
                    }
                    return Ok(());
                }
                Err(e) => {
                    // TODO: Some errors are irrecoverable, such as when a shard has already been
                    // migrated away. In such cases, we need to discard requests
                    // for that shard.
                    WRITER_FRAG_OPERATION_COUNTER
                        .with_label_values(&["submit_fail"])
                        .inc();
                    retries += 1;
                    let delay_secs = retries.min(MAX_RETRY_DELAY_SECS);
                    error!(
                        "submit_frag_retry, {}, frag_path={}, retry_count={}, retry_delay_s={}, err={:?}",
                        tracing_id, submit_info.frag_path, retries, delay_secs, e
                    );
                    tokio::time::sleep(Duration::from_secs(delay_secs)).await;
                }
            }
        }
    }

    async fn submit_base_frag(
        mut client: Box<dyn MetaClient>,
        index_key: &IndexKey,
        submit_info: &SubmitInfo,
    ) -> Result<AppendFragMetaResponse> {
        let frag_meta = submit_info
            .frag_meta
            .clone()
            .ok_or_else(|| anyhow!("submit_base_frag: frag_meta should not be None"))?;
        let req = AppendBaseFragMetaRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            frag_metas: vec![frag_meta],
        };
        let response = client
            .append_base_frag(req)
            .await
            .with_context(|| "Failed to append base frag meta")?;
        Ok(response)
    }

    async fn submit_delta_frag(
        mut client: Box<dyn MetaClient>,
        index_key: &IndexKey,
        shard_id: u64,
        worker_addr: &str,
        submit_info: &SubmitInfo,
    ) -> Result<AppendFragMetaResponse> {
        if shard_id == 0 {
            bail!("shard_id must be greater than 0 for delta frags");
        }
        let frag_metas = if let Some(frag_meta) = submit_info.frag_meta.clone() {
            vec![frag_meta]
        } else {
            vec![]
        };
        let req = AppendDeltaFragMetaRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            shard_id,
            worker_node_addr: worker_addr.to_string(),
            frag_metas,
            progress: submit_info
                .last_cdc_file
                .clone()
                .map(|last_cdc_file| ShardProgress {
                    last_cdc_file: last_cdc_file.to_string(),
                }),
        };
        let response = client
            .append_delta_frag(req)
            .await
            .with_context(|| "Failed to append delta frag meta")?;
        Ok(response)
    }

    async fn submit_frag(
        meta_client_pool: &Arc<dyn MetaClientPool>,
        index_key: &IndexKey,
        shard_id: u64,
        worker_addr: &str,
        submit_info: &SubmitInfo,
    ) -> Result<()> {
        let client = meta_client_pool.get_meta_client().await?;
        let res = if submit_info.is_base {
            Self::submit_base_frag(client, index_key, submit_info).await
        } else {
            Self::submit_delta_frag(client, index_key, shard_id, worker_addr, submit_info).await
        };
        match res {
            Ok(rsp) => {
                if rsp.status == ErrorCode::Success.to_i32() {
                    Ok(())
                } else {
                    bail!("submit_frag: status={}", rsp.status);
                }
            }
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::{SystemTime, UNIX_EPOCH};

    use tempfile::TempDir;
    use tici_grpc::KeyRange;
    use tici_test_utils::test_utils::init_logger_for_ut;
    use tokio::sync::mpsc;

    use super::*;
    use crate::{
        config::FragWriterConfig,
        test_utils::{
            TEST_INDEX_META_OVERLAP, TEST_S3_BUCKET, TEST_SHARD_META, is_sorted_frag,
            test_get_frag_batch_uploader, test_new_frag_path_generator, test_new_metadata_store,
        },
        utils::{
            self, frag_multi_path, key::ShardingKey, meta_client_pool::MockMetaClientPool,
            s3_frag::list_all_files_in_relative, tantivy_doc::RawDoc,
        },
        writer::{cdc_log::OperationType, cdc_log_frag_writer::CDCLogFragWriter},
    };

    #[tokio::test]
    async fn test_run_channel_closed() {
        let (tx, rx) = mpsc::channel(1);
        drop(tx); // Close the channel

        let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap();
        let temp_dir = TempDir::new().unwrap();
        let temp_path = temp_dir.path().to_str().unwrap();
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let uploader = FragUploader::new(
            "test-tracing-id".to_string(),
            &IndexKey::new_nullspace(1, 2),
            3,
            "test-worker-addr".to_string(),
            meta_client_pool,
            frag_batch_uploader,
            TEST_S3_BUCKET.to_string(),
            "test-prefix".to_string(),
            rx,
            true,
            test_new_metadata_store(temp_path).unwrap(),
        );
        uploader.wait_finished().await.unwrap();
    }

    #[tokio::test]
    async fn test_upload_frag_success() {
        init_logger_for_ut();
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        // Insert documents
        let raw_doc1 = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (2, Some("This is the first document.".to_string())),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };
        let raw_doc2 = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 2,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document Two".to_string())),
                (2, Some("This is the second document.".to_string())),
                (3, Some("pk_value3".to_string())),
                (4, Some("pk_value4".to_string())),
            ],
            is_local_timezone: true,
        };

        // Update document
        let raw_doc3 = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 3,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (
                    2,
                    Some("This is the first document that has been updated.".to_string()),
                ),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };

        // Delete document
        let raw_doc4 = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 4,
            op_type: OperationType::Delete,
            cols: vec![],
            is_local_timezone: true,
        };

        for _ in 0..100 {
            cdc_log_frag_writer.add_raw_doc(raw_doc1.clone()).unwrap();
            cdc_log_frag_writer.add_raw_doc(raw_doc2.clone()).unwrap();
            cdc_log_frag_writer.add_raw_doc(raw_doc3.clone()).unwrap();
            cdc_log_frag_writer.add_raw_doc(raw_doc4.clone()).unwrap();
        }
        let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap();
        let backend_storage = frag_batch_uploader.backend_storage();
        let s3_upload_prefix = format!(
            "test_frag_uploader_{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        );
        let meta_client_pool = Arc::new(MockMetaClientPool::new());

        // Commit changes and upload frag
        let frag_path = {
            let frag_writer = cdc_log_frag_writer.take_frag_writer().unwrap();
            let frag_path = frag_writer.frag_path();
            let (tx, rx) = mpsc::channel(1);
            let frag_uploader = FragUploader::new(
                "test-tracing-id".to_string(),
                &IndexKey::new_nullspace(1, 2),
                3,
                "test-worker-addr".to_string(),
                meta_client_pool.clone(),
                frag_batch_uploader,
                TEST_S3_BUCKET.to_string(),
                s3_upload_prefix.clone(),
                rx,
                false,
                test_new_metadata_store(local_data_path).unwrap(),
            );
            tx.send(FragUploadTask::Frag(frag_writer)).await.unwrap();
            drop(tx);
            frag_uploader.wait_finished().await.unwrap();
            frag_path
        };

        let mut mock_client = meta_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, _seq) = mock_client
            .get_frags_by_keyrange(3, &IndexKey::new_nullspace(0, 0), KeyRange::default())
            .await
            .unwrap();
        assert_eq!(frag_metas.len(), 1);

        let download_frag_path = utils::s3_frag::download_frag(
            backend_storage.as_ref(),
            TEST_S3_BUCKET,
            &frag_metas[0],
            &format!("{}/download", local_data_path),
        )
        .await
        .unwrap();

        let mut files = list_all_files_in_relative(&frag_path).unwrap();
        files.sort();

        let mut download_files = list_all_files_in_relative(&download_frag_path).unwrap();
        download_files.sort();

        assert_eq!(files, download_files);

        assert!(is_sorted_frag(&frag_path, &index_meta));
        assert!(is_sorted_frag(&download_frag_path, &index_meta));
    }

    #[tokio::test]
    async fn test_upload_many_frags() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        // Add a document
        let raw_doc1 = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (2, Some("This is the first document.".to_string())),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };
        let raw_doc2 = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 2,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document Two".to_string())),
                (2, Some("This is the second document.".to_string())),
                (3, Some("pk_value3".to_string())),
                (4, Some("pk_value4".to_string())),
            ],
            is_local_timezone: true,
        };

        let (tx, rx) = mpsc::channel(10);
        let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap();
        let backend_storage = frag_batch_uploader.backend_storage();
        let s3_upload_prefix = format!(
            "test_frag_uploader_{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        );
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let frag_uploader = FragUploader::new(
            "test-tracing-id".to_string(),
            &IndexKey::new_nullspace(1, 2),
            3,
            "test-worker-addr".to_string(),
            meta_client_pool,
            frag_batch_uploader.clone(),
            TEST_S3_BUCKET.to_string(),
            s3_upload_prefix.clone(),
            rx,
            false,
            test_new_metadata_store(local_data_path).unwrap(),
        );

        let mut frag_paths = Vec::new();
        for _ in 0..100 {
            cdc_log_frag_writer.add_raw_doc(raw_doc1.clone()).unwrap();
            cdc_log_frag_writer.add_raw_doc(raw_doc2.clone()).unwrap();
            let old_frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap();
            frag_paths.push(old_frag_writer.frag_path());
            tx.send(FragUploadTask::Frag(old_frag_writer))
                .await
                .unwrap();
        }
        drop(tx);

        frag_uploader.wait_finished().await.unwrap();

        for frag_path in &frag_paths {
            assert!(is_sorted_frag(frag_path, &index_meta));
        }

        let mut expected_uploaded_files = Vec::new();
        for frag_path in frag_paths {
            expected_uploaded_files.push(frag_multi_path(&frag_path, 0));
        }
        expected_uploaded_files.sort();

        let mut keys = backend_storage
            .list(TEST_S3_BUCKET, &s3_upload_prefix, None)
            .await
            .unwrap();
        keys.sort();

        assert_eq!(expected_uploaded_files.len(), keys.len());
        for i in 0..expected_uploaded_files.len() {
            let file = expected_uploaded_files[i]
                .strip_prefix(local_data_path)
                .unwrap();
            let key = keys[i].strip_prefix(&s3_upload_prefix).unwrap();
            assert_eq!(file, key);
        }
    }

    #[tokio::test]
    async fn test_update_progress() {
        init_logger_for_ut();
        let (tx, rx) = mpsc::channel(10);
        let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap();
        let s3_upload_prefix = format!(
            "test_frag_uploader_{}",
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        );
        let shard_id = 3;
        let temp_dir = TempDir::new().unwrap();
        let temp_path = temp_dir.path().to_str().unwrap();
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let frag_uploader = FragUploader::new(
            "test-tracing-id".to_string(),
            &IndexKey::new_nullspace(1, 2),
            shard_id,
            "test-worker-addr".to_string(),
            meta_client_pool.clone(),
            frag_batch_uploader,
            TEST_S3_BUCKET.to_string(),
            s3_upload_prefix.clone(),
            rx,
            false,
            test_new_metadata_store(temp_path).unwrap(),
        );

        let last_cdc_file = CDCFileName::parse(
            "tici_default_prefix/cdc/123/461869219056451591/CDC00000000000000000001.json",
        )
        .unwrap();
        let frag_path = "frag1".to_string();
        tx.send(FragUploadTask::Progress(SubmitInfo::for_update_progress(
            last_cdc_file.clone(),
            frag_path,
        )))
        .await
        .unwrap();
        drop(tx);
        frag_uploader.wait_finished().await.unwrap();
        let mut mock_client = meta_client_pool.get_meta_client().await.unwrap();
        let progress = mock_client
            .get_shard_progress(shard_id)
            .await
            .unwrap()
            .unwrap();
        assert_eq!(progress, last_cdc_file.to_string());
    }
}