use std::{
    collections::{HashMap, HashSet},
    mem,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    time::Instant,
};

use anyhow::{Context, Result, bail};
use log::info;
use parking_lot::Mutex;
use storage_backend::storage::BackendStorage;
use tici_grpc::WorkerNodeShardStatus;
use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard, mpsc};

use crate::{
    ShardWriterMeta,
    config::WorkerNodeConfig,
    metadata::metadata_store::MetadataStore,
    utils::{
        frag_path_generator::FragPathGenerator, key::ShardingKeyRange,
        meta_client_pool::MetaClientPool, metrics::begin_shard_writer_manager_metrics,
        schema::IndexMeta,
    },
    writer::{
        cdc_file_reader::CDCFileReader, cdc_file_reader_state::ShardInfo,
        cdc_list_dispatcher::CDCListDispatcher, frag_batch_uploader::FragBatchUploader,
        shard_writer::ShardWriter,
    },
};

#[derive(Hash, Eq, PartialEq, Clone)]
pub struct KeySpaceTableID {
    pub keyspace_id: u32,
    pub physical_table_id: i64,
}

impl std::fmt::Debug for KeySpaceTableID {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ks_{}/t_{}", self.keyspace_id, self.physical_table_id)
    }
}

// ShardWriterManager manages multiple ShardWriters and CDCFileReaders.
pub struct ShardWriterManager {
    config: WorkerNodeConfig,
    backend_storage: Arc<dyn BackendStorage>,
    meta_client_pool: Arc<dyn MetaClientPool>,
    metadata_store: Arc<MetadataStore>,
    frag_path_gen: Arc<FragPathGenerator>,
    frag_batch_uploader: Arc<FragBatchUploader>,
    cdc_list_dispatcher: Arc<CDCListDispatcher>,
    inner: Mutex<ShardWriterManagerInner>,
    operation_gate: ShardOperationGate,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AddShardOutcome {
    Added,
    AlreadyExists,
}

#[derive(Default)]
struct ShardWriterManagerInner {
    file_readers: HashMap<KeySpaceTableID, CDCFileReader>,
    // Shard ID -> ShardWriter
    shard_writers: HashMap<u64, ShardWriter>,
}

impl ShardWriterManager {
    pub fn new(
        config: WorkerNodeConfig,
        backend_storage: Arc<dyn BackendStorage>,
        meta_client_pool: Arc<dyn MetaClientPool>,
        metadata_store: Arc<MetadataStore>,
        frag_path_gen: Arc<FragPathGenerator>,
        frag_batch_uploader: Arc<FragBatchUploader>,
    ) -> Self {
        let cdc_list_dispatcher = Arc::new(CDCListDispatcher::new(
            backend_storage.clone(),
            config.frag_writer.cdc_list_max_concurrency,
        ));
        ShardWriterManager {
            config,
            backend_storage,
            meta_client_pool,
            metadata_store,
            frag_path_gen,
            frag_batch_uploader,
            cdc_list_dispatcher,
            inner: Mutex::new(ShardWriterManagerInner::default()),
            operation_gate: ShardOperationGate::new(),
        }
    }

    fn remove_shard_from_reader(
        inner: &mut ShardWriterManagerInner,
        keyspace_table: &KeySpaceTableID,
        shard_id: u64,
    ) -> Option<CDCFileReader> {
        if let Some(reader) = inner.file_readers.get_mut(keyspace_table) {
            reader.remove_shard(shard_id);
            if reader.is_empty() {
                return inner.file_readers.remove(keyspace_table);
            }
        }
        None
    }

    pub fn shard_statuses(&self) -> Vec<WorkerNodeShardStatus> {
        self.inner
            .lock()
            .shard_writers
            .values()
            .map(|writer| {
                let shard_header = &writer.shard_meta().shard_header;
                WorkerNodeShardStatus {
                    start_key: shard_header.start_key.clone(),
                    shard_id: shard_header.shard_id,
                }
            })
            .collect()
    }

    pub async fn add_shard(
        &self,
        shard_meta: &ShardWriterMeta,
        index_meta: &IndexMeta,
    ) -> Result<AddShardOutcome> {
        let metrics = begin_shard_writer_manager_metrics("add_shard");
        let result = async {
            let shard_id = shard_meta.shard_header.shard_id;
            let _operation = self.operation_gate.enter_shard_op(shard_id, "add_shard")?;
            self.add_shard_inner(shard_meta, index_meta).await
        }
        .await;
        metrics.mark_result(&result);
        result
    }

    async fn add_shard_inner(
        &self,
        shard_meta: &ShardWriterMeta,
        index_meta: &IndexMeta,
    ) -> Result<AddShardOutcome> {
        let start = Instant::now();
        let shard_tracing_id = shard_meta.shard_tracing_id();
        if shard_meta.index_key != *index_meta.index_key() {
            bail!(
                "Index key mismatch between shard meta and index meta: {}, {:?} != {:?}",
                shard_tracing_id,
                shard_meta.index_key,
                index_meta.index_key()
            );
        }
        let shard_id = shard_meta.shard_header.shard_id;
        if self.inner.lock().shard_writers.contains_key(&shard_id) {
            info!("add_shard_skip_already_exists, {}", shard_tracing_id);
            return Ok(AddShardOutcome::AlreadyExists);
        }

        let cdc_checkpoint = &shard_meta.cdc_checkpoint;
        let prefix = &cdc_checkpoint.table_prefix;
        let (sender, receiver) = mpsc::channel(self.config.frag_writer.cdc_sender_channel_size);
        let shard_range = ShardingKeyRange::half_open(
            shard_meta.shard_header.start_key.clone(),
            shard_meta.shard_header.end_key.clone(),
        )?;
        // sender is moved into reader_shard_info
        let reader_shard_info = ShardInfo::new(
            shard_id,
            sender,
            shard_range.clone(),
            shard_tracing_id.clone(),
            cdc_checkpoint.last_file.clone(),
            index_meta.clone(),
        )?;

        // receiver is moved into shard_writer
        let create_shard_writer = Instant::now();
        let shard_writer = ShardWriter::new(
            shard_meta,
            index_meta,
            &self.config,
            self.frag_batch_uploader.clone(),
            self.meta_client_pool.clone(),
            self.frag_path_gen.clone(),
            receiver,
            self.metadata_store.clone(),
        )
        .await
        .with_context(|| {
            format!(
                "Failed to create shard writer: {}, prefix={}",
                shard_tracing_id, prefix
            )
        })?;
        let create_shard_writer_elapsed = create_shard_writer.elapsed();

        // Add shard to the file reader. Each physical table has a unique file reader.
        let keyspace_table = KeySpaceTableID {
            keyspace_id: index_meta.index_key().keyspace_id,
            physical_table_id: index_meta.index_key().table_id,
        };
        let mut inner = self.inner.lock();
        match inner.file_readers.get_mut(&keyspace_table) {
            None => {
                let mut reader = CDCFileReader::new_with_dispatcher(
                    self.backend_storage.clone(),
                    cdc_checkpoint,
                    self.config.frag_writer.cdc_poll_interval.into(),
                    self.config.frag_writer.cdc_max_poll_interval.into(),
                    self.config.frag_writer.cdc_prefetch_window,
                    self.config
                        .frag_writer
                        .cdc_head_object_last_modified_fallback,
                    self.cdc_list_dispatcher.clone(),
                );
                reader.start();
                reader.add_shard(reader_shard_info);
                inner.file_readers.insert(keyspace_table.clone(), reader);
            }
            Some(file_reader) => {
                file_reader.add_shard(reader_shard_info);
            }
        }

        inner.shard_writers.insert(shard_id, shard_writer);

        info!(
            "add_shard_success, {}, cdc_checkpoint={:?}, range={:?}, elapsed={:?}, create_shard_writer_elapsed={:?}",
            shard_tracing_id,
            cdc_checkpoint,
            shard_range,
            start.elapsed(),
            create_shard_writer_elapsed,
        );
        Ok(AddShardOutcome::Added)
    }

    pub async fn remove_shard(&self, shard_id: u64) -> Result<()> {
        let metrics = begin_shard_writer_manager_metrics("remove_shard");
        let result = async {
            let _operation = self
                .operation_gate
                .enter_shard_op(shard_id, "remove_shard")?;
            self.remove_shard_inner(shard_id).await
        }
        .await;
        metrics.mark_result(&result);
        result
    }

    async fn remove_shard_inner(&self, shard_id: u64) -> Result<()> {
        let start = Instant::now();
        let (shard_writer, reader_to_stop, shard_tracing_id, prefix, keyspace_table) = {
            let mut inner = self.inner.lock();
            let shard_writer = match inner.shard_writers.remove(&shard_id) {
                Some(writer) => writer,
                None => {
                    info!("remove_shard_skip_not_found, shard_id={}", shard_id);
                    return Ok(());
                }
            };
            let (shard_tracing_id, prefix, keyspace_table) = {
                let shard_meta = shard_writer.shard_meta();
                (
                    shard_meta.shard_tracing_id(),
                    shard_meta.cdc_checkpoint.table_prefix.clone(),
                    KeySpaceTableID {
                        keyspace_id: shard_meta.index_key.keyspace_id,
                        physical_table_id: shard_meta.index_key.table_id,
                    },
                )
            };

            (
                shard_writer,
                Self::remove_shard_from_reader(&mut inner, &keyspace_table, shard_id),
                shard_tracing_id,
                prefix,
                keyspace_table,
            )
        };

        let shutdown_time = Instant::now();
        shard_writer.shutdown().await;
        if let Some(reader) = reader_to_stop {
            reader.stop().await;
            info!(
                "remove_cdc_file_reader, {}, prefix={}, keyspace_table={:?}",
                shard_tracing_id, prefix, keyspace_table
            );
        }
        let shutdown_elapsed = shutdown_time.elapsed();

        info!(
            "remove_shard_success, {}, prefix={}, elapsed={:?}, shutdown_elapsed={:?}",
            shard_tracing_id,
            prefix,
            start.elapsed(),
            shutdown_elapsed
        );
        Ok(())
    }

    pub async fn shutdown(&self) {
        let shutdown_guard = self.operation_gate.begin_shutdown().await;
        if shutdown_guard.already_started {
            info!("shard_writer_manager_shutdown_skip_already_started");
            return;
        }

        let (file_readers, shard_writers) = {
            let mut inner = self.inner.lock();
            (
                mem::take(&mut inner.file_readers),
                mem::take(&mut inner.shard_writers),
            )
        };

        info!(
            "shard_writer_manager_shutdown_start, reader_count={}, shard_writer_count={}",
            file_readers.len(),
            shard_writers.len()
        );

        for (keyspace_table, reader) in file_readers {
            info!(
                "stop_cdc_file_reader, keyspace_table={:?}, prefix={:?}",
                keyspace_table,
                reader.get_prefix()
            );
            reader.stop().await;
        }

        info!("shard_writer_manager_shutdown_readers_finished");

        for (id, sw) in shard_writers {
            info!("shutdown_shard_writer, shard_id={}", id);
            sw.shutdown().await;
        }

        info!("shard_writer_manager_shutdown_complete");
    }

    pub fn meta_client_pool(&self) -> Arc<dyn MetaClientPool> {
        self.meta_client_pool.clone()
    }

    pub fn metadata_store(&self) -> Arc<MetadataStore> {
        self.metadata_store.clone()
    }

    pub fn frag_path_gen(&self) -> Arc<FragPathGenerator> {
        self.frag_path_gen.clone()
    }

    pub fn frag_batch_uploader(&self) -> Arc<FragBatchUploader> {
        self.frag_batch_uploader.clone()
    }
}

// Coordinates shard add/remove operations with manager shutdown.
//
// The gate enforces two independent rules:
// - Only one add/remove operation may run for the same shard at a time.
// - Shutdown must wait for in-flight add/remove operations, and new operations
//   must fail once shutdown starts.
struct ShardOperationGate {
    // Shard IDs with an add/remove operation in progress.
    pending_shards: Mutex<HashSet<u64>>,
    // Blocks shutdown until in-flight add/remove operations finish.
    lifecycle: RwLock<()>,
    shutting_down: AtomicBool,
}

impl ShardOperationGate {
    fn new() -> Self {
        Self {
            pending_shards: Mutex::new(HashSet::new()),
            lifecycle: RwLock::new(()),
            shutting_down: AtomicBool::new(false),
        }
    }

    fn ensure_running(&self, op: &'static str) -> Result<()> {
        if self.shutting_down.load(Ordering::Acquire) {
            bail!("{op}: shard writer manager is shutting down");
        }
        Ok(())
    }

    fn enter_shard_op(&self, shard_id: u64, op: &'static str) -> Result<ShardOperationGuard<'_>> {
        // Check the shutdown flag both before and after taking the read guard.
        // Shutdown sets the flag before waiting for the write guard, so this
        // closes the gate immediately while allowing in-flight operations to drain.
        self.ensure_running(op)?;
        let Ok(lifecycle) = self.lifecycle.try_read() else {
            bail!("{op}: shard writer manager is shutting down");
        };
        self.ensure_running(op)?;

        if !self.pending_shards.lock().insert(shard_id) {
            bail!("{op}: shard operation already pending, shard_id={shard_id}");
        }

        Ok(ShardOperationGuard {
            _lifecycle: lifecycle,
            pending_shards: &self.pending_shards,
            shard_id,
        })
    }

    async fn begin_shutdown(&self) -> ShardOperationShutdownGuard<'_> {
        // Mark shutdown before waiting for the write guard. Otherwise new shard
        // operations could keep entering while shutdown is waiting for readers.
        let already_started = self.shutting_down.swap(true, Ordering::AcqRel);
        let lifecycle = self.lifecycle.write().await;
        ShardOperationShutdownGuard {
            _lifecycle: lifecycle,
            already_started,
        }
    }
}

struct ShardOperationGuard<'a> {
    _lifecycle: RwLockReadGuard<'a, ()>,
    pending_shards: &'a Mutex<HashSet<u64>>,
    shard_id: u64,
}

impl Drop for ShardOperationGuard<'_> {
    fn drop(&mut self) {
        self.pending_shards.lock().remove(&self.shard_id);
    }
}

struct ShardOperationShutdownGuard<'a> {
    _lifecycle: RwLockWriteGuard<'a, ()>,
    already_started: bool,
}

#[cfg(test)]
mod tests {
    use std::sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    };

    use anyhow::{Context, Result};
    use async_trait::async_trait;
    use futures::FutureExt;
    use log::info;
    use tempfile::TempDir;
    use tici_common::{
        codec::codec_v2::get_txn_keyspace_prefix, config::S3Config, index::IndexKey,
        keyspace::NULLSPACEID,
    };
    use tici_grpc::KeyRange;
    use tici_test_utils::test_utils::init_logger_for_ut;
    use tokio::sync::mpsc;
    use uuid::Uuid;

    use super::{ShardOperationGate, ShardWriterManager, ShardWriterManagerInner};
    use crate::{
        ShardWriterMeta,
        config::WorkerNodeConfig,
        data_types::codec::{Datum, encode_index_key},
        metadata::metadata_store::MetadataStore,
        reader::{frag_reader::FragReaderMeta, query_info::QueryInfo, shard_mirror::ShardMirror},
        test_utils::{
            TEST_S3_BUCKET, TEST_S3_BUCKET_PARTITION, TEST_S3_PREFIX,
            helpers::wait_index_ready_for_partition_table,
            index_meta::mock_custom_hdfs_logs_index_meta,
            is_sorted_frag, mock_custom_employees_partition_index_meta,
            shard_meta::{mock_custom_shard_meta, test_get_shard_writer_meta},
            test_create_backend_storage, test_get_inf_range, test_get_shard_mirror_meta,
            test_new_frag_path_generator, test_new_metadata_store, write_test_data,
        },
        utils::{
            frag_path_generator::{FragIDGenerator, FragPathGenerator},
            key::ShardingKeyRange,
            meta_client_pool::{MetaClientPool, MockMetaClientPool},
            s3_frag,
            schema::IndexMeta,
        },
        writer::{
            cdc_file_reader::CDCFileReader, cdc_file_reader_state::ShardInfo, cdc_log::CDCLogEntry,
            frag_batch_uploader::FragBatchUploader, shard_writer_manager::KeySpaceTableID,
        },
    };

    fn assert_gate_err<T>(result: Result<T>, expected: &str) {
        match result {
            Ok(_) => panic!("shard operation should fail with: {expected}"),
            Err(err) => {
                let msg = err.to_string();
                assert!(
                    msg.contains(expected),
                    "unexpected shard operation error: {msg}"
                );
            }
        }
    }

    struct FailingFragIDGenerator;

    #[async_trait]
    impl FragIDGenerator for FailingFragIDGenerator {
        async fn gen_id(&self) -> anyhow::Result<u64> {
            anyhow::bail!("mock frag id generator failure")
        }
    }

    struct FailOnNthFragIDGenerator {
        fail_on: usize,
        next_id: AtomicUsize,
    }

    impl FailOnNthFragIDGenerator {
        fn new(fail_on: usize) -> Self {
            Self {
                fail_on,
                next_id: AtomicUsize::new(0),
            }
        }
    }

    #[async_trait]
    impl FragIDGenerator for FailOnNthFragIDGenerator {
        async fn gen_id(&self) -> anyhow::Result<u64> {
            let current = self.next_id.fetch_add(1, Ordering::Relaxed) + 1;
            if current == self.fail_on {
                anyhow::bail!("mock frag id generator failure on call {current}")
            }
            Ok(current as u64)
        }
    }

    fn mock_reader_shard_info(
        shard_meta: &ShardWriterMeta,
        sender: mpsc::Sender<CDCLogEntry>,
        index_meta: IndexMeta,
    ) -> Result<ShardInfo> {
        let shard_range = ShardingKeyRange::half_open(
            shard_meta.shard_header.start_key.clone(),
            shard_meta.shard_header.end_key.clone(),
        )?;
        ShardInfo::new(
            shard_meta.shard_header.shard_id,
            sender,
            shard_range,
            shard_meta.shard_tracing_id(),
            shard_meta.cdc_checkpoint.last_file.clone(),
            index_meta,
        )
    }

    async fn test_sharding_by_tenant_id(keyspace_id: u32) -> Result<()> {
        let table_id = 123;
        let index_id = 3;
        let index_key = IndexKey {
            keyspace_id,
            table_id,
            index_id,
        };
        let frag_prefix = Uuid::new_v4().to_string();
        let index_meta = mock_custom_hdfs_logs_index_meta(index_key.clone(), &[4])?;

        let mut shard_metas = vec![];
        let shard_ranges = [(-100, 1), (1, 20), (20, 70), (70, 100), (100, 200)];
        let expected_counts = [0, 28571, 75076, 44821, 1532];
        for (i, (start, end)) in shard_ranges.iter().enumerate() {
            let mut start_key = vec![];
            if keyspace_id != NULLSPACEID {
                start_key.extend_from_slice(get_txn_keyspace_prefix(keyspace_id).as_slice());
            }
            encode_index_key(&mut start_key, table_id, index_id, &[Datum::Int64(*start)])?;

            let mut end_key = vec![];
            if keyspace_id != NULLSPACEID {
                end_key.extend_from_slice(get_txn_keyspace_prefix(keyspace_id).as_slice());
            }
            encode_index_key(&mut end_key, table_id, index_id, &[Datum::Int64(*end)])?;
            let shard_id = (i + 1) as u64;
            let shard_meta = mock_custom_shard_meta(
                index_key.clone(),
                shard_id,
                start_key,
                end_key,
                frag_prefix.clone(),
            );
            shard_metas.push(shard_meta);
        }

        let temp_dir = TempDir::new()?;
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());
        let backend_storage = test_create_backend_storage().await?;
        let meta_client_pool = Arc::new(MockMetaClientPool::new());

        write_test_data(
            &config,
            &shard_metas,
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await?;

        let mut mock_client = meta_client_pool.get_meta_client().await?;
        for (shard_meta, expected_count) in shard_metas.iter().zip(expected_counts) {
            let (frag_metas, _seq) = mock_client
                .get_frags_by_keyrange(
                    shard_meta.shard_header.shard_id,
                    &shard_meta.index_key,
                    KeyRange::default(),
                )
                .await?;
            let mut frag_reader_metas = Vec::with_capacity(frag_metas.len());
            for frag_meta in frag_metas {
                let property = frag_meta
                    .property
                    .as_ref()
                    .context("frag_meta.property is None")?;
                let local_frag_path = frag_meta.frag_path.replace(
                    &shard_meta.frag_location.prefix,
                    &config.storage.frag_base_path(),
                );
                assert!(!std::fs::exists(local_frag_path).unwrap());

                let local_frag_path = s3_frag::download_frag(
                    backend_storage.as_ref(),
                    TEST_S3_BUCKET,
                    &frag_meta,
                    &config.storage.frag_base_path(),
                )
                .await?;
                assert!(is_sorted_frag(&local_frag_path, &index_meta));
                frag_reader_metas.push(FragReaderMeta::new(
                    &local_frag_path,
                    property.min_handle.clone(),
                    property.max_handle.clone(),
                )?);
            }

            let shard_mirror_meta = test_get_shard_mirror_meta(shard_meta);
            let dedup_output = ShardMirror::dedup_frags(
                shard_mirror_meta.shard_tracing_id(),
                shard_mirror_meta.shard_range().unwrap(),
                frag_reader_metas.clone(),
                config.storage.temporary_path(),
            )
            .await??;
            let shard_mirror = ShardMirror::new(
                shard_mirror_meta,
                index_meta.clone(),
                frag_reader_metas,
                dedup_output.obsolete_doc_bitsets,
            )
            .await?;
            let reader = shard_mirror.get_reader();
            let query_info = QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )?;
            let (all_docs, _search_stat) = reader.search_with_query_info(query_info).unwrap();
            assert_eq!(all_docs.len(), expected_count);
        }
        Ok(())
    }

    #[tokio::test]
    async fn test_sharding_by_tenant_id_with_keyspace() -> Result<()> {
        init_logger_for_ut();
        test_sharding_by_tenant_id(1).await?;
        Ok(())
    }

    #[tokio::test]
    async fn test_sharding_by_tenant_id_without_keyspace() -> Result<()> {
        init_logger_for_ut();
        test_sharding_by_tenant_id(NULLSPACEID).await?;
        Ok(())
    }

    #[tokio::test]
    async fn test_shard_operation_gate_rejects_operations_after_shutdown() {
        let gate = ShardOperationGate::new();
        let shutdown_guard = gate.begin_shutdown().await;

        assert!(!shutdown_guard.already_started);
        assert_gate_err(
            gate.enter_shard_op(1, "add_shard"),
            "add_shard: shard writer manager is shutting down",
        );
        assert_gate_err(
            gate.enter_shard_op(2, "remove_shard"),
            "remove_shard: shard writer manager is shutting down",
        );
    }

    #[tokio::test]
    async fn test_shard_operation_gate_shutdown_waits_for_in_flight_operations() -> Result<()> {
        let gate = ShardOperationGate::new();
        let add_guard = gate.enter_shard_op(1, "add_shard")?;
        let remove_guard = gate.enter_shard_op(2, "remove_shard")?;
        let mut shutdown = Box::pin(gate.begin_shutdown());

        assert!(shutdown.as_mut().now_or_never().is_none());
        assert_gate_err(
            gate.enter_shard_op(3, "add_shard"),
            "add_shard: shard writer manager is shutting down",
        );

        drop(add_guard);
        assert!(shutdown.as_mut().now_or_never().is_none());

        drop(remove_guard);
        let shutdown_guard = shutdown.await;
        assert!(!shutdown_guard.already_started);

        Ok(())
    }

    #[test]
    fn test_shard_operation_gate_rejects_concurrent_operations_for_same_shard() -> Result<()> {
        let gate = ShardOperationGate::new();
        let add_guard = gate.enter_shard_op(1, "add_shard")?;

        assert_gate_err(
            gate.enter_shard_op(1, "remove_shard"),
            "remove_shard: shard operation already pending, shard_id=1",
        );
        let other_shard_guard = gate.enter_shard_op(2, "remove_shard")?;

        drop(add_guard);
        let same_shard_guard = gate.enter_shard_op(1, "remove_shard")?;

        drop(other_shard_guard);
        drop(same_shard_guard);
        Ok(())
    }

    #[tokio::test]
    async fn test_add_shard_does_not_modify_state_on_shard_writer_creation_failure() -> Result<()> {
        init_logger_for_ut();

        let temp_dir = TempDir::new()?;
        let mut config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());
        config.s3 = S3Config::default_for_test();
        let metadata_store = test_new_metadata_store(temp_dir.path().to_str().unwrap())?;
        let backend_storage = test_create_backend_storage().await?;
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let frag_path_gen = Arc::new(FragPathGenerator::new(
            &config.storage.frag_base_path(),
            Box::new(FailingFragIDGenerator),
        ));
        let manager = ShardWriterManager::new(
            config.clone(),
            backend_storage.clone(),
            meta_client_pool.clone(),
            metadata_store.clone(),
            frag_path_gen,
            Arc::new(FragBatchUploader::new(
                backend_storage,
                config.frag_upload.clone(),
            )),
        );

        let shard_meta = test_get_shard_writer_meta();
        let index_meta = mock_custom_hdfs_logs_index_meta(shard_meta.index_key.clone(), &[4])?;
        let shard_id = shard_meta.shard_header.shard_id;

        let err = manager
            .add_shard(&shard_meta, &index_meta)
            .await
            .expect_err("add_shard should fail when frag id generation fails");
        assert!(err.to_string().contains("Failed to create shard writer"));
        {
            let inner = manager.inner.lock();
            assert!(!inner.shard_writers.contains_key(&shard_id));
            assert!(!inner.file_readers.contains_key(&KeySpaceTableID {
                keyspace_id: index_meta.index_key().keyspace_id,
                physical_table_id: index_meta.index_key().table_id,
            }));
        }

        manager.shutdown().await;
        Ok(())
    }

    #[tokio::test]
    async fn test_add_shard_does_not_rewind_existing_reader_on_shard_writer_creation_failure()
    -> Result<()> {
        init_logger_for_ut();

        let temp_dir = TempDir::new()?;
        let mut config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());
        config.s3 = S3Config::default_for_test();
        let metadata_store = test_new_metadata_store(temp_dir.path().to_str().unwrap())?;
        let backend_storage = test_create_backend_storage().await?;
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let frag_path_gen = Arc::new(FragPathGenerator::new(
            &config.storage.frag_base_path(),
            Box::new(FailOnNthFragIDGenerator::new(2)),
        ));
        let manager = ShardWriterManager::new(
            config.clone(),
            backend_storage.clone(),
            meta_client_pool,
            metadata_store,
            frag_path_gen,
            Arc::new(FragBatchUploader::new(
                backend_storage,
                config.frag_upload.clone(),
            )),
        );

        let base_shard_meta = test_get_shard_writer_meta();
        let index_meta = mock_custom_hdfs_logs_index_meta(base_shard_meta.index_key.clone(), &[4])?;
        let first_last_file = format!(
            "{}/461869219056451591/CDC00000000000000000001.json",
            base_shard_meta.cdc_checkpoint.table_prefix
        );
        let mut first_shard_meta = mock_custom_shard_meta(
            base_shard_meta.index_key.clone(),
            11,
            vec![0],
            vec![1],
            Uuid::new_v4().to_string(),
        );
        first_shard_meta.cdc_checkpoint.last_file =
            tici_common::cdc_file_name::CDCFileName::parse(&first_last_file).unwrap();

        let mut second_shard_meta = mock_custom_shard_meta(
            base_shard_meta.index_key,
            12,
            vec![1],
            vec![2],
            Uuid::new_v4().to_string(),
        );
        second_shard_meta.cdc_checkpoint.last_file =
            tici_common::cdc_file_name::CDCFileName::empty();

        manager.add_shard(&first_shard_meta, &index_meta).await?;

        let keyspace_table = KeySpaceTableID {
            keyspace_id: index_meta.index_key().keyspace_id,
            physical_table_id: index_meta.index_key().table_id,
        };
        let expected_last_file = {
            let inner = manager.inner.lock();
            let reader = inner
                .file_readers
                .get(&keyspace_table)
                .context("file reader missing after first shard add")?;
            let expected_last_file = reader.get_state().get_last_file();
            assert_eq!(expected_last_file.to_string(), first_last_file);
            assert_eq!(reader.get_state().get_shard_infos().len(), 1);
            expected_last_file
        };

        let err = manager
            .add_shard(&second_shard_meta, &index_meta)
            .await
            .expect_err("second add_shard should fail when frag id generation fails");
        assert!(err.to_string().contains("Failed to create shard writer"));
        {
            let inner = manager.inner.lock();
            assert!(
                !inner
                    .shard_writers
                    .contains_key(&second_shard_meta.shard_header.shard_id)
            );

            let reader = inner
                .file_readers
                .get(&keyspace_table)
                .context("file reader missing after failed second shard add")?;
            let shard_infos = reader.get_state().get_shard_infos();
            assert_eq!(reader.get_state().get_last_file(), expected_last_file);
            assert_eq!(shard_infos.len(), 1);
            assert!(shard_infos.contains_key(&first_shard_meta.shard_header.shard_id));
            assert!(!shard_infos.contains_key(&second_shard_meta.shard_header.shard_id));
        }

        manager.shutdown().await;
        Ok(())
    }

    #[tokio::test]
    async fn test_remove_shard_from_reader_detaches_empty_reader() -> Result<()> {
        init_logger_for_ut();

        let shard_meta = test_get_shard_writer_meta();
        let shard_id = shard_meta.shard_header.shard_id;
        let index_meta = mock_custom_hdfs_logs_index_meta(shard_meta.index_key.clone(), &[4])?;
        let keyspace_table = KeySpaceTableID {
            keyspace_id: index_meta.index_key().keyspace_id,
            physical_table_id: index_meta.index_key().table_id,
        };
        let reader = CDCFileReader::new(
            test_create_backend_storage().await?,
            &shard_meta.cdc_checkpoint,
            std::time::Duration::from_secs(1),
            std::time::Duration::from_secs(1),
            1,
            true,
        );
        let (sender, _) = mpsc::channel(1);
        reader.add_shard(mock_reader_shard_info(&shard_meta, sender, index_meta)?);

        let mut inner = ShardWriterManagerInner::default();
        inner.file_readers.insert(keyspace_table.clone(), reader);

        let detached_reader =
            ShardWriterManager::remove_shard_from_reader(&mut inner, &keyspace_table, shard_id);

        assert!(detached_reader.is_some());
        assert!(!inner.file_readers.contains_key(&keyspace_table));

        if let Some(reader) = detached_reader {
            reader.stop().await;
        }

        Ok(())
    }

    #[tokio::test]
    async fn test_shard_writer_manager_for_partition_table() -> Result<()> {
        init_logger_for_ut();

        let keyspace_id = NULLSPACEID;
        let table_id = [437, 438, 439, 440];
        let index_id = 3;
        let frag_prefix = Uuid::new_v4().to_string();

        // Shard ranges must match id distribution in cdc_test_partition_data.zip:
        // - Partition 437: ids 1-4
        // - Partition 438: ids 5-9
        // - Partition 439: ids 10-14
        // - Partition 440: ids 15-18
        let mut shard_metas = vec![];
        let shard_ranges = [(1, 5), (5, 10), (10, 15), (15, 20)];
        let expected_counts = [4, 5, 5, 4];
        for (i, (start, end)) in shard_ranges.iter().enumerate() {
            let mut start_key = vec![];
            if keyspace_id != NULLSPACEID {
                start_key.extend_from_slice(get_txn_keyspace_prefix(keyspace_id).as_slice());
            }
            encode_index_key(
                &mut start_key,
                table_id[i],
                index_id,
                &[Datum::Int64(*start)],
            )?;
            let mut end_key = vec![];
            if keyspace_id != NULLSPACEID {
                end_key.extend_from_slice(get_txn_keyspace_prefix(keyspace_id).as_slice());
            }
            encode_index_key(&mut end_key, table_id[i], index_id, &[Datum::Int64(*end)])?;

            let shard_id = (i + 1) as u64;
            let mut shard_meta = mock_custom_shard_meta(
                IndexKey {
                    keyspace_id,
                    table_id: table_id[i],
                    index_id,
                },
                shard_id,
                start_key,
                end_key,
                frag_prefix.clone(),
            );
            shard_meta.cdc_checkpoint.bucket = TEST_S3_BUCKET_PARTITION.to_string();
            shard_meta.cdc_checkpoint.table_prefix =
                format!("{}/{}", TEST_S3_PREFIX.as_str(), table_id[i]);
            shard_meta.cdc_checkpoint.last_file = tici_common::cdc_file_name::CDCFileName::empty();
            shard_metas.push(shard_meta);
        }

        let temp_dir = TempDir::new()?;
        let mut config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());
        config.s3 = S3Config::default_for_test();
        config.storage.init_worker_data_dir()?;
        let backend_storage = test_create_backend_storage().await?;
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let metadata_store = Arc::new(MetadataStore::new(&config.storage.metadata_path())?);
        let frag_path_gen = test_new_frag_path_generator(&config.storage.frag_base_path());
        let shard_writer_manager = ShardWriterManager::new(
            config.clone(),
            backend_storage.clone(),
            meta_client_pool.clone(),
            metadata_store.clone(),
            frag_path_gen,
            Arc::new(FragBatchUploader::new(
                backend_storage.clone(),
                config.frag_upload.clone(),
            )),
        );
        let cdc_bucket = shard_metas[0].cdc_checkpoint.bucket.clone();

        for shard_meta in &shard_metas {
            assert_eq!(cdc_bucket, shard_meta.cdc_checkpoint.bucket);

            let index_meta =
                mock_custom_employees_partition_index_meta(shard_meta.index_key.clone())?;

            shard_writer_manager
                .add_shard(shard_meta, &index_meta)
                .await?;
        }

        let mock_client = meta_client_pool.get_meta_client().await?;
        wait_index_ready_for_partition_table(
            backend_storage.clone(),
            &cdc_bucket,
            &shard_metas,
            mock_client,
        )
        .await?;

        {
            let inner = shard_writer_manager.inner.lock();
            for reader in inner.file_readers.values() {
                info!(
                    "Reader, state of last file: {}",
                    reader.get_state().get_last_file()
                );
            }
        }

        // Verify expected_counts per shard (see test_sharding_by_tenant_id)
        let mut mock_client = meta_client_pool.get_meta_client().await?;
        for (shard_meta, expected_count) in shard_metas.iter().zip(expected_counts) {
            let index_meta =
                mock_custom_employees_partition_index_meta(shard_meta.index_key.clone())?;
            let (frag_metas, _seq) = mock_client
                .get_frags_by_keyrange(
                    shard_meta.shard_header.shard_id,
                    &shard_meta.index_key,
                    KeyRange::default(),
                )
                .await?;

            let mut frag_reader_metas = Vec::with_capacity(frag_metas.len());
            for frag_meta in frag_metas {
                let property = frag_meta
                    .property
                    .as_ref()
                    .context("frag_meta.property is None")?;
                let local_frag_path = frag_meta.frag_path.replace(
                    &shard_meta.frag_location.prefix,
                    &config.storage.frag_base_path(),
                );
                assert!(!std::fs::exists(local_frag_path).unwrap());

                let local_frag_path = s3_frag::download_frag(
                    backend_storage.as_ref(),
                    &shard_meta.frag_location.bucket,
                    &frag_meta,
                    &config.storage.frag_base_path(),
                )
                .await?;
                assert!(is_sorted_frag(&local_frag_path, &index_meta));
                frag_reader_metas.push(FragReaderMeta::new(
                    &local_frag_path,
                    property.min_handle.clone(),
                    property.max_handle.clone(),
                )?);
            }

            let shard_mirror_meta = test_get_shard_mirror_meta(shard_meta);
            let dedup_output = ShardMirror::dedup_frags(
                shard_mirror_meta.shard_tracing_id(),
                shard_mirror_meta.shard_range().unwrap(),
                frag_reader_metas.clone(),
                config.storage.temporary_path(),
            )
            .await??;
            let shard_mirror = ShardMirror::new(
                shard_mirror_meta,
                index_meta.clone(),
                frag_reader_metas,
                dedup_output.obsolete_doc_bitsets,
            )
            .await?;
            let reader = shard_mirror.get_reader();
            let query_info = QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )?;
            let (all_docs, _search_stat) = reader.search_with_query_info(query_info).unwrap();
            assert_eq!(all_docs.len(), expected_count);
        }

        shard_writer_manager.shutdown().await;

        Ok(())
    }
}