use std::{
    collections::{
        BTreeMap, HashMap,
        btree_map::Entry::{Occupied as BOccupied, Vacant as BVacant},
    },
    fmt::Debug,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU64},
    },
    time::Duration,
};

use anyhow::{Result, anyhow, bail};
use log::info;
use tici_common::index::IndexKey;
use tici_grpc::{IndexInfo, S3Location, TableInfo};
use tokio::sync::{
    OnceCell, RwLock,
    mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel},
    oneshot,
};

use crate::{
    HandleRange, RangeBound, ShardInfo, ShardManifest,
    meta::{
        IndexMetaProperties, MetaStore,
        util::BgRunner,
        writer::{RateCounter, ShardBuildProgress, ShardStats},
    },
};

type ShardMap = BTreeMap<RangeBound, Arc<RwLock<ShardInfo>>>;
type SplitTopologyResult = Result<(), anyhow::Error>;
type SplitTopologyRx = oneshot::Receiver<SplitTopologyResult>;
type MergeTopologyResult = Result<Arc<RwLock<ShardInfo>>, anyhow::Error>;
type MergeTopologyRx = oneshot::Receiver<MergeTopologyResult>;
type RemoveAllShardsResult = Result<Vec<(String, u64)>, anyhow::Error>;
type RemoveAllShardsTx = oneshot::Sender<RemoveAllShardsResult>;

struct TopologyActorHandle {
    tx: UnboundedSender<TopologyCmd>,
    _runner: BgRunner,
}

pub struct TopologyMutationPermit {
    tx: Option<UnboundedSender<TopologyCmd>>,
}

impl Drop for TopologyMutationPermit {
    fn drop(&mut self) {
        if let Some(tx) = self.tx.take() {
            let _ = tx.send(TopologyCmd::ReleaseMutation);
        }
    }
}

enum TopologyCmd {
    AcquireMutation {
        tx: oneshot::Sender<Result<(), anyhow::Error>>,
    },
    ReleaseMutation,
    Split {
        original_start: RangeBound,
        new_shards: Vec<ShardInfo>,
        tx: oneshot::Sender<SplitTopologyResult>,
    },
    Merge {
        input_starts: Vec<RangeBound>,
        new_shard: Box<ShardInfo>,
        tx: oneshot::Sender<MergeTopologyResult>,
    },
    RemoveAllShards {
        tx: RemoveAllShardsTx,
    },
}

#[derive(Debug, Clone, Default)]
pub struct FTSIndexStat {
    compaction_finish: Arc<AtomicU64>,
    compaction_abort: Arc<AtomicU64>,
    compaction_yield: Arc<AtomicU64>,
    split_finish: Arc<AtomicU64>,
    split_no_need: Arc<AtomicU64>,
    split_yield: Arc<AtomicU64>,
    split_unable: Arc<AtomicU64>,
    pub merge_finish: Arc<AtomicU64>,
    pub big_shard_after_write: Arc<AtomicU64>,
}

impl FTSIndexStat {
    pub fn to_short_string(&self) -> String {
        format!(
            "compaction: F{}+A{}+Y{}, split: F{}+N{}+Y{}+U{}, merge={}, big={}",
            self.get_compaction_finish(),
            self.get_compaction_abort(),
            self.get_compaction_yield(),
            self.get_split_finish(),
            self.get_split_no_need(),
            self.get_split_yield(),
            self.get_split_unable(),
            self.merge_finish.load(std::sync::atomic::Ordering::Relaxed),
            self.big_shard_after_write
                .load(std::sync::atomic::Ordering::Relaxed)
        )
    }

    pub fn get_compaction_yield(&self) -> u64 {
        self.compaction_yield
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_compaction_finish(&self) -> u64 {
        self.compaction_finish
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_compaction_abort(&self) -> u64 {
        self.compaction_abort
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn observe_compaction_finish(&self) {
        self.compaction_finish
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_compaction_abort(&self) {
        self.compaction_abort
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_compaction_yield(&self) {
        self.compaction_yield
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn get_split_finish(&self) -> u64 {
        self.split_finish.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_split_no_need(&self) -> u64 {
        self.split_no_need
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_split_yield(&self) -> u64 {
        self.split_yield.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn get_split_unable(&self) -> u64 {
        self.split_unable.load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn observe_split_finish(&self) {
        self.split_finish
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_split_no_need(&self) {
        self.split_no_need
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_split_yield(&self) {
        self.split_yield
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn observe_split_unable(&self) {
        self.split_unable
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }
}

/// An `FTSIndex` is for one FTS index of a table.
/// It manages several shards.
pub struct FTSIndex {
    meta_store: Arc<dyn MetaStore>,
    // There are two way to identify a shard:
    // - (index_id, start_key)
    // - shard_id
    shards: Arc<RwLock<ShardMap>>,
    topology_actor: OnceCell<TopologyActorHandle>,
    pub table_info: TableInfo,
    pub index_info: IndexInfo,
    cdc_s3_bucket: String,
    cdc_s3_prefix: String,
    pub stats: Arc<FTSIndexStat>,
    pub append_qps: Arc<RateCounter>,
    topology_frozen: AtomicBool,
    finished_append_base: AtomicBool,
}

impl Debug for FTSIndex {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FTSIndex")
            .field("table_info", &self.table_info)
            .field("index_info", &self.index_info)
            .field(
                "cdc_path",
                &format!("{}/{}", &self.cdc_s3_bucket, &self.cdc_s3_prefix),
            )
            .field("topology_frozen", &self.is_topology_frozen())
            .field("finished_append_base", &self.is_finished_append_base())
            .field("stats", &self.stats)
            .finish()
    }
}

unsafe impl Send for FTSIndex {}
unsafe impl Sync for FTSIndex {}

fn spawn_topology_actor(
    meta_store: Arc<dyn MetaStore>,
    shards: Arc<RwLock<ShardMap>>,
    index_key: IndexKey,
) -> TopologyActorHandle {
    let (tx, rx) = unbounded_channel();
    let runner_name = format!("FTSIndexTopology({:?})", index_key);
    let runner = BgRunner::spawn_name(runner_name, move |rf| async move {
        run_topology_actor(meta_store, shards, index_key, rx, rf).await;
    });
    TopologyActorHandle {
        tx,
        _runner: runner,
    }
}

async fn handle_split_cmd(
    shards: &Arc<RwLock<ShardMap>>,
    index_key: &IndexKey,
    original_start: RangeBound,
    new_shards: Vec<ShardInfo>,
) -> Result<(), anyhow::Error> {
    let mut guard = shards.write().await;
    match guard.remove(&original_start) {
        Some(_) => {}
        None => bail!(
            "split_shards: shard missing while applying topology, index={:?}, start_key={:?}",
            index_key,
            original_start
        ),
    }
    for shard_info in new_shards.iter() {
        let sortable_start = shard_info.range.sortable_start();
        if guard.contains_key(&sortable_start) {
            bail!(
                "split_shards: new shard already exists, index={:?}, start_key={:?}",
                index_key,
                sortable_start
            );
        }
        guard.insert(sortable_start, Arc::new(RwLock::new(shard_info.clone())));
    }
    Ok(())
}

async fn handle_merge_cmd(
    shards: &Arc<RwLock<ShardMap>>,
    index_key: &IndexKey,
    input_starts: Vec<RangeBound>,
    new_shard: Box<ShardInfo>,
) -> Result<Arc<RwLock<ShardInfo>>, anyhow::Error> {
    if input_starts.len() < 2 {
        bail!(
            "merge_shards: can't merge, less than 2 shards, index={:?}",
            index_key
        );
    }
    let new_start = new_shard.range.sortable_start();
    let new_shard = Arc::new(RwLock::new(*new_shard));

    let mut guard = shards.write().await;
    for start_key in input_starts {
        if guard.remove(&start_key).is_none() {
            bail!(
                "merge_shards: shard missing while applying topology, index={:?}, start_key={:?}",
                index_key,
                start_key
            );
        }
    }
    if guard.contains_key(&new_start) {
        bail!(
            "merge_shards: merged shard already exists, index={:?}, start_key={:?}",
            index_key,
            new_start
        );
    }
    guard.insert(new_start, new_shard.clone());
    Ok(new_shard)
}

async fn handle_remove_all_shards_cmd(
    meta_store: &Arc<dyn MetaStore>,
    shards: &Arc<RwLock<ShardMap>>,
    index_key: &IndexKey,
) -> Result<Vec<(String, u64)>, anyhow::Error> {
    let shard_arcs = {
        info!("drop_index: acquire lock {:?}", index_key,);
        let mut guard = shards.write().await;
        info!(
            "drop_index: start {:?}, current has {} shards",
            index_key,
            guard.len()
        );
        std::mem::take(&mut *guard)
            .into_values()
            .collect::<Vec<_>>()
    };

    let mut removed = vec![];
    for shard in shard_arcs {
        let shard_info = shard.read().await;
        meta_store.remove_shard(&shard_info).await?;
        removed.push((shard_info.writer.clone(), shard_info.shard_id));
    }
    Ok(removed)
}

async fn try_finish_remove_all_shards(
    meta_store: &Arc<dyn MetaStore>,
    shards: &Arc<RwLock<ShardMap>>,
    index_key: &IndexKey,
    inflight_mutations: usize,
    pending_remove: &mut Option<RemoveAllShardsTx>,
) {
    if inflight_mutations != 0 {
        info!(
            "try_finish_remove_all_shards: pending, index_key={:?}, inflight_mutations={}",
            index_key, inflight_mutations
        );
        return;
    }
    if let Some(tx) = pending_remove.take() {
        let _ = tx.send(handle_remove_all_shards_cmd(meta_store, shards, index_key).await);
    }
}

async fn run_topology_actor(
    meta_store: Arc<dyn MetaStore>,
    shards: Arc<RwLock<ShardMap>>,
    index_key: IndexKey,
    mut rx: UnboundedReceiver<TopologyCmd>,
    mut stop_rx: oneshot::Receiver<()>,
) {
    info!("FTSIndexTopology({:?}): start", index_key);
    let mut inflight_mutations = 0usize;
    let mut closing = false;
    let mut pending_remove = None;
    loop {
        tokio::select! {
            _ = &mut stop_rx => {
                info!("FTSIndexTopology({:?}): recv stop signal", index_key);
                break;
            }
            cmd = rx.recv() => {
                let Some(cmd) = cmd else {
                    info!("FTSIndexTopology({:?}): sender dropped", index_key);
                    break;
                };
                match cmd {
                    TopologyCmd::AcquireMutation { tx } => {
                        if closing {
                            let _ = tx.send(Err(anyhow!(
                                "topology actor is closing, index={:?}",
                                index_key
                            )));
                        } else {
                            inflight_mutations += 1;
                            let _ = tx.send(Ok(()));
                        }
                    }
                    TopologyCmd::ReleaseMutation => {
                        inflight_mutations = inflight_mutations.saturating_sub(1);
                        if closing {
                            try_finish_remove_all_shards(
                                &meta_store,
                                &shards,
                                &index_key,
                                inflight_mutations,
                                &mut pending_remove,
                            )
                            .await;
                        }
                    }
                    TopologyCmd::Split {
                        original_start,
                        new_shards,
                        tx,
                    } => {
                        let _ = tx.send(
                            handle_split_cmd(&shards, &index_key, original_start, new_shards).await,
                        );
                    }
                    TopologyCmd::Merge {
                        input_starts,
                        new_shard,
                        tx,
                    } => {
                        let _ = tx.send(
                            handle_merge_cmd(&shards, &index_key, input_starts, new_shard).await,
                        );
                    }
                    TopologyCmd::RemoveAllShards { tx } => {
                        if closing {
                            let _ = tx.send(Err(anyhow!(
                                "remove_all_shards: already closing, index={:?}",
                                index_key
                            )));
                            continue;
                        }
                        closing = true;
                        pending_remove = Some(tx);
                        try_finish_remove_all_shards(
                            &meta_store,
                            &shards,
                            &index_key,
                            inflight_mutations,
                            &mut pending_remove,
                        )
                        .await;
                    }
                }
            }
        }
    }
    info!("FTSIndexTopology({:?}): quit", index_key);
}

impl FTSIndex {
    pub fn new(
        meta_store: Arc<dyn MetaStore>,
        table_info: TableInfo,
        index_info: IndexInfo,
        cdc_s3_bucket: String,
        cdc_s3_prefix: String,
    ) -> Self {
        Self::new_from_restore(
            meta_store,
            table_info,
            index_info,
            cdc_s3_bucket,
            cdc_s3_prefix,
            IndexMetaProperties::default(),
        )
    }

    pub fn new_from_restore(
        meta_store: Arc<dyn MetaStore>,
        table_info: TableInfo,
        index_info: IndexInfo,
        cdc_s3_bucket: String,
        cdc_s3_prefix: String,
        properties: IndexMetaProperties,
    ) -> Self {
        Self {
            meta_store,
            shards: Arc::new(RwLock::new(BTreeMap::new())),
            topology_actor: OnceCell::new(),
            table_info,
            index_info,
            cdc_s3_bucket,
            cdc_s3_prefix,
            stats: Arc::new(FTSIndexStat::default()),
            append_qps: Arc::new(RateCounter::new(Duration::from_secs(30))),
            topology_frozen: AtomicBool::new(false),
            finished_append_base: AtomicBool::new(properties.finished_append_base),
        }
    }

    async fn topology_actor(&self) -> &TopologyActorHandle {
        let meta_store = self.meta_store.clone();
        let shards = self.shards.clone();
        let index_key = self.index_key();
        self.topology_actor
            .get_or_init(|| async move { spawn_topology_actor(meta_store, shards, index_key) })
            .await
    }

    async fn topology_tx(&self) -> &UnboundedSender<TopologyCmd> {
        &self.topology_actor().await.tx
    }

    pub async fn acquire_topology_mutation_permit(&self) -> Result<TopologyMutationPermit> {
        let tx = self.topology_tx().await.clone();
        let (reply_tx, reply_rx) = oneshot::channel();
        tx.send(TopologyCmd::AcquireMutation { tx: reply_tx })
            .map_err(|e| anyhow!("acquire_topology_mutation_permit: send failed, err={:?}", e))?;
        reply_rx.await.map_err(|e| {
            anyhow!(
                "acquire_topology_mutation_permit: actor canceled, err={:?}",
                e
            )
        })??;
        Ok(TopologyMutationPermit { tx: Some(tx) })
    }

    #[cfg(any(test, feature = "testexport"))]
    #[cfg_attr(feature = "testexport", allow(dead_code))]
    pub(super) async fn topology_split_shards(
        &self,
        original_start: RangeBound,
        new_shards: Vec<ShardInfo>,
    ) -> Result<(), anyhow::Error> {
        let rx = self
            .submit_topology_split_shards(original_start, new_shards)
            .await?;
        rx.await
            .map_err(|e| anyhow!("split_shards: topology actor canceled, err={:?}", e))?
    }

    pub(super) async fn submit_topology_split_shards(
        &self,
        original_start: RangeBound,
        new_shards: Vec<ShardInfo>,
    ) -> Result<SplitTopologyRx, anyhow::Error> {
        let (tx, rx) = oneshot::channel();
        self.topology_tx()
            .await
            .send(TopologyCmd::Split {
                original_start,
                new_shards,
                tx,
            })
            .map_err(|e| anyhow!("split_shards: send topology actor failed, err={:?}", e))?;
        Ok(rx)
    }

    #[cfg(any(test, feature = "testexport"))]
    #[cfg_attr(feature = "testexport", allow(dead_code))]
    pub(super) async fn topology_merge_shards(
        &self,
        input_starts: Vec<RangeBound>,
        new_shard: ShardInfo,
    ) -> Result<Arc<RwLock<ShardInfo>>, anyhow::Error> {
        let rx = self
            .submit_topology_merge_shards(input_starts, new_shard)
            .await?;
        rx.await
            .map_err(|e| anyhow!("merge_shards: topology actor canceled, err={:?}", e))?
    }

    pub(super) async fn submit_topology_merge_shards(
        &self,
        input_starts: Vec<RangeBound>,
        new_shard: ShardInfo,
    ) -> Result<MergeTopologyRx, anyhow::Error> {
        let (tx, rx) = oneshot::channel();
        self.topology_tx()
            .await
            .send(TopologyCmd::Merge {
                input_starts,
                new_shard: Box::new(new_shard),
                tx,
            })
            .map_err(|e| anyhow!("merge_shards: send topology actor failed, err={:?}", e))?;
        Ok(rx)
    }

    pub fn keyspace_id(&self) -> u32 {
        self.table_info.keyspace_id
    }

    pub fn table_id(&self) -> i64 {
        self.index_info.table_id
    }

    pub fn index_id(&self) -> i64 {
        self.index_info.index_id
    }

    pub fn get_cdc_location(&self) -> S3Location {
        S3Location {
            bucket: self.cdc_s3_bucket.clone(),
            prefix: self.cdc_s3_prefix.clone(),
        }
    }

    pub fn get_cdc_s3_bucket(&self) -> &str {
        &self.cdc_s3_bucket
    }

    pub fn get_cdc_s3_prefix(&self) -> &str {
        &self.cdc_s3_prefix
    }

    pub fn index_key(&self) -> IndexKey {
        IndexKey::new(self.keyspace_id(), self.table_id(), self.index_id())
    }

    pub fn is_topology_frozen(&self) -> bool {
        self.topology_frozen
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub fn set_topology_frozen(&self, frozen: bool) {
        self.topology_frozen
            .store(frozen, std::sync::atomic::Ordering::Relaxed);
    }

    pub fn is_finished_append_base(&self) -> bool {
        self.finished_append_base
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    pub async fn finish_append_base(&self) -> Result<bool> {
        info!(
            "finish_append_base: snapshotting shards, index={:?}, finished_append_base={}",
            self.index_key(),
            self.is_finished_append_base()
        );
        let shards = self.clone_shards().await;
        let mut changed = false;
        if !self.is_finished_append_base() {
            info!(
                "finish_append_base: persisting index flag, index={:?}",
                self.index_key()
            );
            self.meta_store
                .finish_index_append_base(&self.index_key())
                .await?;
            self.finished_append_base
                .store(true, std::sync::atomic::Ordering::Relaxed);
            info!(
                "finish_append_base: persisted index flag, index={:?}",
                self.index_key()
            );
            changed = true;
        }

        for shard in shards {
            let mut shard = shard.write().await;
            info!(
                "finish_append_base: process, index={:?}, shard_id={}, seq={}, base_len={}, delta_len={}, compactions={}",
                self.index_key(),
                shard.shard_id,
                shard.seq,
                shard.manifest.base_slice().len(),
                shard.manifest.delta_slice().len(),
                shard.progress.compaction_state.len()
            );
            if !shard.switch_base_to_delta(&self.meta_store).await? {
                continue;
            }
            info!(
                "finish_append_base: switched shard manifest, index={:?}, shard_id={}, seq={}, base_len={}, delta_len={}",
                self.index_key(),
                shard.shard_id,
                shard.seq,
                shard.manifest.base_slice().len(),
                shard.manifest.delta_slice().len()
            );
            changed = true;
        }

        info!(
            "finish_append_base: finished switch, index={:?}, changed={}",
            self.index_key(),
            changed
        );
        Ok(changed)
    }

    pub async fn get_all_shard_ids(&self) -> Vec<u64> {
        let mut shard_ids = vec![];
        for (_, shard) in self.shards.read().await.iter() {
            shard_ids.push(shard.read().await.shard_id);
        }
        shard_ids
    }
}

impl FTSIndex {
    // Create a shard which is the initially-first shard of that index.
    // Don't use this function if the shard is created by split or migration.
    pub async fn bootstrap_shard(
        &self,
        range: &HandleRange,
        epoch: u64,
        manifest: ShardManifest,
        writer: String,
        bucket: String,
        prefix: String,
        progress: ShardBuildProgress,
    ) -> Result<Arc<RwLock<ShardInfo>>, anyhow::Error> {
        let mut guard = self.shards.write().await;
        match guard.entry(range.sortable_start()) {
            BOccupied(_) => Err(anyhow::format_err!(
                "The shard is already bootstrapped {:?} range={:?}",
                self.index_key(),
                range
            )),
            BVacant(v) => {
                let mut shard_info = ShardInfo {
                    index_key: self.index_key(),
                    shard_id: 0, // Assigned by MetaStore
                    range: range.clone(),
                    epoch,
                    seq: 0,
                    manifest,
                    writer,
                    progress,
                    frag_upload_location: S3Location { bucket, prefix },
                    stats: ShardStats::default(),
                    append_empty_persist_count: Arc::new(AtomicU64::new(0)),
                };
                match self.meta_store.init_shard(&shard_info).await {
                    Err(e) => Err(e),
                    Ok(shard_id) => {
                        shard_info.shard_id = shard_id;
                        shard_info.progress.mark_cdc_last_file_persisted();
                        let shard_info = Arc::new(RwLock::new(shard_info));
                        v.insert(shard_info.clone());
                        Ok(shard_info)
                    }
                }
            }
        }
    }
}

impl FTSIndex {
    pub async fn clone_shards_map(&self) -> BTreeMap<RangeBound, Arc<RwLock<ShardInfo>>> {
        self.shards.read().await.clone()
    }

    pub async fn clone_shards(&self) -> Vec<Arc<RwLock<ShardInfo>>> {
        self.shards
            .read()
            .await
            .iter()
            .map(|e| e.1.clone())
            .collect::<Vec<_>>()
    }

    pub async fn select_first_shard(&self) -> Option<Arc<RwLock<ShardInfo>>> {
        self.shards
            .read()
            .await
            .first_key_value()
            .map(|x| x.1.clone())
    }

    pub async fn clone_shard(&self, start_key: &RangeBound) -> Option<Arc<RwLock<ShardInfo>>> {
        self.shards.read().await.get(start_key).cloned()
    }

    pub async fn shards_len(&self) -> usize {
        self.shards.read().await.len()
    }

    pub(super) async fn insert_restored_shard(
        &self,
        start_key: RangeBound,
        shard_info: Arc<RwLock<ShardInfo>>,
    ) {
        self.shards.write().await.insert(start_key, shard_info);
    }

    #[cfg(any(test, feature = "testexport"))]
    pub async fn add_shard_for_test(
        &self,
        start_key: RangeBound,
        shard_info: Arc<RwLock<ShardInfo>>,
    ) {
        self.insert_restored_shard(start_key, shard_info).await;
    }

    pub async fn get_shard(&self, start_key: &RangeBound) -> Option<Arc<RwLock<ShardInfo>>> {
        self.shards.read().await.get(start_key).cloned()
    }

    pub async fn find_by_lower_bound(
        &self,
        start_key: &RangeBound,
    ) -> Option<Arc<RwLock<ShardInfo>>> {
        self.shards
            .read()
            .await
            .lower_bound(std::collections::Bound::Included(start_key))
            .next()
            .map(|x| x.1.clone())
    }

    pub async fn find_shard_by_key(&self, key: &RangeBound) -> Option<Arc<RwLock<ShardInfo>>> {
        let shards = self.shards.read().await;
        let candidate = shards
            .range(..=key)
            .next_back()
            .map(|(_, shard)| shard.clone());
        drop(shards);
        match candidate {
            Some(shard) => {
                if shard.read().await.range.contains_key(key) {
                    Some(shard)
                } else {
                    None
                }
            }
            None => None,
        }
    }

    pub async fn remove_all_shards(&self) -> Result<Vec<(String, u64)>> {
        let (tx, rx) = oneshot::channel();
        self.topology_tx()
            .await
            .send(TopologyCmd::RemoveAllShards { tx })
            .map_err(|e| anyhow!("remove_all_shards: send topology actor failed, err={:?}", e))?;
        rx.await
            .map_err(|e| anyhow!("remove_all_shards: topology actor canceled, err={:?}", e))?
    }

    /// Find shards that overlap with the requested key ranges
    /// - `key_ranges` each key_range could be [start_key, end_key) or
    ///   [start_key, end_key]. If no key ranges are provided, return all shards
    /// - `include_end_key` indicates whether the end_key of each key_range is
    ///   inclusive
    pub async fn find_overlapping_shards(
        &self,
        key_ranges: &[tici_grpc::KeyRange],
        include_end_key: bool,
    ) -> Vec<Arc<RwLock<ShardInfo>>> {
        let shard_arcs = {
            let shards = self.shards.read().await;

            // Return all shards if no key ranges are provided.
            if key_ranges.is_empty() {
                return shards.values().cloned().collect();
            }

            // Snapshot shard handles first so we don't hold the index lock while
            // awaiting per-shard reads below.
            shards.values().cloned().collect::<Vec<_>>()
        };

        // TODO: To optimize the performance when there are more than thousands of
        // shards inside one index, should optimize this loop by binary searching the
        // start_key
        let mut overlapping_shards = Vec::new();
        for shard_arc in shard_arcs {
            let shard = shard_arc.read().await;
            let overlaps = key_ranges.iter().any(|key_range| {
                if include_end_key {
                    shard
                        .range
                        .range_overlaps_minmax(&key_range.start_key, &key_range.end_key)
                } else {
                    shard
                        .range
                        .range_overlaps_range(&key_range.start_key, &key_range.end_key)
                }
            });
            drop(shard);
            if overlaps {
                overlapping_shards.push(shard_arc);
            }
        }
        overlapping_shards
    }

    pub async fn get_all_referenced_size(&self) -> usize {
        let shards = self.clone_shards().await;
        let mut frags = HashMap::new();
        for shard in shards {
            for f in shard.read().await.manifest.cloned_fragments().into_iter() {
                let b = f.bytes() as usize;
                let _ = frags.try_insert(f.frag_path, b);
            }
        }
        let s = frags.values().sum();
        s
    }

    #[cfg(feature = "testexport")]
    pub fn new_for_test() -> Self {
        use crate::meta::store::test_utils::DummyMetaStore;
        Self {
            meta_store: Arc::new(DummyMetaStore::default()),
            shards: Arc::new(RwLock::new(BTreeMap::new())),
            topology_actor: OnceCell::new(),
            table_info: TableInfo::default(),
            index_info: IndexInfo::default(),
            cdc_s3_bucket: "test_bucket".to_string(),
            cdc_s3_prefix: "test_prefix".to_string(),
            stats: Arc::new(FTSIndexStat::default()),
            append_qps: Arc::new(RateCounter::new(Duration::from_secs(30))),
            topology_frozen: AtomicBool::new(false),
            finished_append_base: AtomicBool::new(false),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        HandleRange, ShardInfo,
        meta::store::test_utils::{DummyMetaStore, RecordingMetaStore},
    };

    #[test]
    fn test_topology_freeze_flag() {
        let index = FTSIndex::new(
            Arc::new(DummyMetaStore::default()),
            TableInfo::default(),
            IndexInfo::default(),
            String::new(),
            String::new(),
        );

        assert!(!index.is_topology_frozen());

        index.set_topology_frozen(true);
        assert!(index.is_topology_frozen());

        index.set_topology_frozen(false);
        assert!(!index.is_topology_frozen());
    }

    #[test]
    fn test_finished_append_base_flag_defaults_to_false() {
        let index = FTSIndex::new(
            Arc::new(DummyMetaStore::default()),
            TableInfo::default(),
            IndexInfo::default(),
            String::new(),
            String::new(),
        );

        assert!(!index.is_finished_append_base());
    }

    #[tokio::test]
    async fn test_find_overlapping_shards_returns_matching_shards() {
        let index = FTSIndex::new(
            Arc::new(DummyMetaStore::default()),
            TableInfo::default(),
            IndexInfo::default(),
            String::new(),
            String::new(),
        );

        let left_range = HandleRange::new(vec![b'a'], vec![b'c']);
        let middle_range = HandleRange::new(vec![b'c'], vec![b'f']);
        let right_range = HandleRange::new(vec![b'f'], vec![b'z']);

        index
            .insert_restored_shard(
                left_range.sortable_start(),
                Arc::new(RwLock::new(ShardInfo::new_for_test(left_range.clone()))),
            )
            .await;
        index
            .insert_restored_shard(
                middle_range.sortable_start(),
                Arc::new(RwLock::new(ShardInfo::new_for_test(middle_range.clone()))),
            )
            .await;
        index
            .insert_restored_shard(
                right_range.sortable_start(),
                Arc::new(RwLock::new(ShardInfo::new_for_test(right_range.clone()))),
            )
            .await;

        let overlapping = index
            .find_overlapping_shards(
                &[tici_grpc::KeyRange {
                    start_key: vec![b'b'],
                    end_key: vec![b'f'],
                }],
                true,
            )
            .await;

        assert_eq!(overlapping.len(), 3);
        assert_eq!(overlapping[0].read().await.range, left_range);
        assert_eq!(overlapping[1].read().await.range, middle_range);
        assert_eq!(overlapping[2].read().await.range, right_range);
    }

    #[tokio::test]
    async fn test_topology_split_shards_updates_index() {
        let store = Arc::new(RecordingMetaStore::default());
        let index = FTSIndex::new(
            store,
            TableInfo::default(),
            IndexInfo::default(),
            String::new(),
            String::new(),
        );
        let original_range = HandleRange::new(vec![b'a'], vec![b'z']);
        let original_shard = Arc::new(RwLock::new(ShardInfo::new_for_test(original_range.clone())));
        index
            .insert_restored_shard(original_range.sortable_start(), original_shard.clone())
            .await;

        let new_ranges = [
            HandleRange::new(vec![b'a'], vec![b'm']),
            HandleRange::new(vec![b'm'], vec![b'z']),
        ];
        let new_shards = vec![
            ShardInfo::new_for_test(new_ranges[0].clone()),
            ShardInfo::new_for_test(new_ranges[1].clone()),
        ];
        index
            .topology_split_shards(original_range.sortable_start(), new_shards)
            .await
            .unwrap();

        assert_eq!(index.shards_len().await, 2);
        assert!(
            index
                .get_shard(&original_range.sortable_start())
                .await
                .is_some()
        );
        assert!(
            index
                .get_shard(&new_ranges[1].sortable_start())
                .await
                .is_some()
        );
    }

    #[tokio::test]
    async fn test_remove_all_shards_clears_index_via_actor() {
        let store = Arc::new(RecordingMetaStore::default());
        let index = FTSIndex::new(
            store.clone(),
            TableInfo::default(),
            IndexInfo::default(),
            String::new(),
            String::new(),
        );
        for (shard_id, range) in [
            (1, HandleRange::new(vec![b'a'], vec![b'c'])),
            (2, HandleRange::new(vec![b'c'], vec![b'f'])),
        ] {
            let mut shard = ShardInfo::new_for_test(range.clone());
            shard.shard_id = shard_id;
            index
                .insert_restored_shard(range.sortable_start(), Arc::new(RwLock::new(shard)))
                .await;
        }

        let removed = index.remove_all_shards().await.unwrap();

        assert_eq!(removed.len(), 2);
        assert_eq!(index.shards_len().await, 0);
        assert_eq!(store.removed_shard_ids.lock().unwrap().len(), 2);
    }

    #[tokio::test]
    async fn test_remove_all_shards_waits_for_inflight_topology_mutation() {
        let store = Arc::new(RecordingMetaStore::default());
        let index = Arc::new(FTSIndex::new(
            store.clone(),
            TableInfo::default(),
            IndexInfo::default(),
            String::new(),
            String::new(),
        ));
        let original_range = HandleRange::new(vec![b'a'], vec![b'z']);
        let original_shard = Arc::new(RwLock::new(ShardInfo::new_for_test(original_range.clone())));
        index
            .insert_restored_shard(original_range.sortable_start(), original_shard)
            .await;

        let permit = index.acquire_topology_mutation_permit().await.unwrap();
        let remove_index = index.clone();
        let remove_task =
            tokio::spawn(async move { remove_index.remove_all_shards().await.unwrap() });
        tokio::task::yield_now().await;

        let new_ranges = [
            HandleRange::new(vec![b'a'], vec![b'm']),
            HandleRange::new(vec![b'm'], vec![b'z']),
        ];
        index
            .topology_split_shards(
                original_range.sortable_start(),
                vec![
                    ShardInfo::new_for_test(new_ranges[0].clone()),
                    ShardInfo::new_for_test(new_ranges[1].clone()),
                ],
            )
            .await
            .unwrap();
        drop(permit);

        let removed = remove_task.await.unwrap();
        assert_eq!(removed.len(), 2);
        assert_eq!(index.shards_len().await, 0);
        assert_eq!(store.removed_shard_ids.lock().unwrap().len(), 2);
    }
}