mod append;
mod compaction;
mod reader;

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

use anyhow::{Result, anyhow};
use heapsize::HeapSizeOf as _;
use log::{error, info};
use rand::prelude::*;
use storage_backend::storage::create_backend_storage_for_test;
use tici::{
    RangeBound, ShardInfo,
    common::HandleRange,
    config::MetaServiceConfig,
    meta::{
        IndexMetaProperties, MetaServiceCore,
        metrics::*,
        writer::{FTSIndex, MergeShardReason, ShardMergeResult, WorkerNodeMeta},
    },
    worker::{
        mocks::{MockShardWriterManager, midpoint_bytes},
        server::DispatchShardWriter,
    },
};
use tici_common::{codec::table_raw_key, config::ServerConfig, index::IndexKey};
use tici_grpc::{
    CreateIndexRequest, GrpcMetaServiceClient, GrpcMetaServiceInternalClient, IndexInfo, TableInfo,
    WorkerNodeShardStatus,
};
use tici_shard::config::WorkerNodeConfig;
use tici_system::readable::ReadableDuration;
use tokio::sync::RwLock;
use tonic::transport::Channel;

pub use crate::test_utils::must_finally;
use crate::test_utils::{
    meta_handle::MetaServiceHandle,
    partition_table::make_create_partition_index_req,
    reader_handle::{ReaderNodeHandle, make_default_reader_config},
    runner,
    shard_cache_scaffold::MockReader,
    worker_handle::WorkerServiceHandle,
    workload_ci::make_create_index_req,
};
async fn clear_s3(meta_cfg: &MetaServiceConfig) -> Result<()> {
    let mut cfg = meta_cfg.s3.clone();
    cfg.use_path_style = true;
    let storage = create_backend_storage_for_test(&cfg).await?;
    let all_keys = storage
        .list_all(&meta_cfg.s3.bucket, &format!("{}s", meta_cfg.s3.prefix))
        .await
        .unwrap_or_default();
    if !all_keys.is_empty() {
        // Delete in batches of 1000 (S3 limit)
        for chunk in all_keys.chunks(1000) {
            let _ = storage.delete_objects(&meta_cfg.s3.bucket, chunk).await;
        }
    }
    Ok(())
}

pub struct E2ETestSuite {
    // We may restart the meta service, so we keep it as an Option
    meta: Option<MetaServiceHandle>,
    // multiple workers
    pub workers: HashMap<String, WorkerServiceHandle>,
    pub readers: HashMap<String, ReaderNodeHandle>,
    pub mocked_readers: HashMap<String, MockReader>,
    pub writer_config_template: WorkerNodeConfig,
    pub start: SystemTime,
}

impl Drop for E2ETestSuite {
    fn drop(&mut self) {
        runner::clear_failpoints();
    }
}

fn setup_panic_hook() {
    std::panic::set_hook(Box::new(|info| {
        eprintln!("panic_hook: task panicked: {}", info);
        println!("panic_hook: task panicked: {}", info);
    }));
}

impl E2ETestSuite {
    pub async fn new_with_cfg(
        meta_cfg: MetaServiceConfig,
        writer_cfg: WorkerNodeConfig,
        reset_meta: bool,
    ) -> Result<Self> {
        setup_panic_hook();
        tici::meta::metrics::clear_prom_metrics();
        clear_s3(&meta_cfg).await?;
        info!("creating meta service with config={:?}", meta_cfg);
        let meta = MetaServiceHandle::new(meta_cfg, reset_meta).await?;
        let mut suite = Self {
            meta: Some(meta),
            workers: HashMap::new(),
            writer_config_template: writer_cfg,
            start: SystemTime::now(),
            readers: HashMap::new(),
            mocked_readers: HashMap::new(),
        };
        suite.add_real_worker_node(None).await;
        Ok(suite)
    }

    pub async fn new_with_mock_cfg(
        meta_cfg: MetaServiceConfig,
        writer_cfg: WorkerNodeConfig,
    ) -> Result<Self> {
        let mut suite = Self::new_with_mock_cfg_basic(meta_cfg, writer_cfg, true).await?;
        suite.add_mocked_worker_node(None).await;
        Ok(suite)
    }

    pub async fn new_with_mock_cfg_basic(
        meta_cfg: MetaServiceConfig,
        writer_cfg: WorkerNodeConfig,
        reset_meta: bool,
    ) -> Result<Self> {
        setup_panic_hook();
        tici::meta::metrics::clear_prom_metrics();
        clear_s3(&meta_cfg).await?;
        info!("creating meta service with config={:?}", meta_cfg);
        let meta = MetaServiceHandle::new(meta_cfg, reset_meta).await?;
        Ok(Self {
            meta: Some(meta),
            workers: HashMap::new(),
            writer_config_template: writer_cfg,
            start: SystemTime::now(),
            readers: HashMap::new(),
            mocked_readers: HashMap::new(),
        })
    }

    pub fn get_meta_handle(&self) -> &MetaServiceHandle {
        self.meta.as_ref().unwrap()
    }

    pub async fn add_mocked_worker_node(&mut self, addr: Option<String>) {
        let (a, mut config) = if let Some(addr) = addr {
            let mut config = self.writer_config_template.clone();
            config.server.addr = addr.clone();
            config.server.set_advertise_addr(addr.clone());
            (addr, config)
        } else {
            (
                self.writer_config_template.server.addr.clone(),
                self.writer_config_template.clone(),
            )
        };
        config.server.status_addr = String::new(); // Don't start status server in mock mode
        config.storage.data_dir = format!("{}/{}", config.storage.data_dir, uuid::Uuid::new_v4()); // Unique path, avoid conflict or reuse

        let mock_writer = DispatchShardWriter::Mock(Arc::new(MockShardWriterManager::new(
            self.get_meta_service(),
            config.server.clone(),
        )));
        info!(
            "adding mock worker node at {}, server_config={:?}",
            a, config.server
        );
        let worker = WorkerServiceHandle::new_with_worker(config, mock_writer, None)
            .await
            .unwrap();
        self.workers.insert(a, worker);
        // FIXME: wait for heartbeat is a workaround for shard not created if worker not
        // registered
        self.wait_heartbeat().await;
    }

    pub async fn add_real_worker_node(&mut self, addr: Option<String>) {
        let (a, config) = if let Some(addr) = addr {
            let mut config = self.writer_config_template.clone();
            config.server.addr = addr.clone();
            (addr, config)
        } else {
            (
                self.writer_config_template.server.addr.clone(),
                self.writer_config_template.clone(),
            )
        };
        info!("adding real worker node at {}", a);
        let worker = WorkerServiceHandle::new(config.clone(), None)
            .await
            .unwrap();
        self.workers.insert(a, worker);
        // FIXME: wait for heartbeat is a workaround for shard not created if worker not
        // registered
        self.wait_heartbeat().await;
    }

    pub async fn add_mocked_reader_node(
        &mut self,
        addr: String,
        tiflash_addr: String,
        start_hb: Option<Duration>,
    ) -> MockReader {
        let mut reader_service =
            MockReader::new(addr.clone(), tiflash_addr, self.get_meta_internal_client());
        reader_service.start_server();
        let _ = self.mocked_readers.insert(addr, reader_service.clone());
        if let Some(intv) = start_hb {
            reader_service.start_heartbeat(intv).await;
        }
        reader_service
    }

    pub async fn add_real_reader_node(&mut self, addr: String, tiflash_addr: String) {
        let searchlib_config = make_default_reader_config(
            &self.get_meta_service().config.load_full(),
            addr.clone(),
            tiflash_addr,
        );
        let reader_handle = ReaderNodeHandle::new(searchlib_config).await;
        let _ = self.readers.insert(addr, reader_handle);
    }

    pub async fn add_real_reader_node_with_cfg(
        &mut self,
        searchlib_config: tici_search_lib::config::Config,
    ) {
        let addr = searchlib_config.tici.reader_node.addr.clone();
        let reader_handle = ReaderNodeHandle::new(searchlib_config).await;
        let _ = self.readers.insert(addr, reader_handle);
    }

    pub async fn check_meta_subscriber(&self) {
        let shards = self.collect_shards_from_memory().await;
        assert_eq!(
            shards.len(),
            self.get_meta_handle().get_subscribed_shards().len()
        );
        for shard in shards {
            assert_eq!(
                &shard,
                &*self
                    .get_meta_handle()
                    .get_subscribed_shards()
                    .get(&shard.shard_id)
                    .unwrap()
            );
        }
    }

    pub async fn must_meta_subscriber_consistent(&self) {
        must_finally(async || {
            let shards = self.collect_shards_from_memory().await;
            let sub_ss = self.get_meta_handle().get_subscribed_shards();

            if shards.len() != sub_ss.len() {
                return false;
            }
            for shard in shards {
                let sk = sub_ss.get(&shard.shard_id).unwrap();
                if shard.manifest != sk.manifest {
                    info!(
                        "must_meta_subscriber_consistent manifest: {:?} <=> {:?}",
                        shard.manifest, sk.manifest
                    );
                    return false;
                }
                if shard.get_header() != sk.get_header() {
                    info!(
                        "must_meta_subscriber_consistent get_header: {:?} <=> {:?}",
                        shard.manifest, sk.manifest
                    );
                    return false;
                }
                if shard.progress.cdc_s3_last_file != sk.progress.cdc_s3_last_file {
                    info!(
                        "must_meta_subscriber_consistent cdc_s3_last_file: {:?} <=> {:?}",
                        shard.manifest, sk.manifest
                    );
                    return false;
                }
            }
            true
        })
        .await
    }

    pub async fn stop_readers(&mut self) -> Result<()> {
        info!("stop in total {} real readers", self.readers.len());
        // Take all the workers and stop them
        for (_, mut reader) in self.readers.drain() {
            // Properly stop the reader. The ReaderHandle has a `stop` async method.
            reader.stop().await;
        }
        info!("stop in total {} mock readers", self.mocked_readers.len());
        for (_, mut reader) in self.mocked_readers.drain() {
            // Properly stop the reader. The ReaderHandle has a `stop` async method.
            reader.shutdown().await;
        }
        Ok(())
    }

    pub async fn stop_worker(&mut self) -> Result<()> {
        info!("stop in total {} workers", self.workers.len());
        // Take all the workers and stop them
        for (_, worker_handle) in self.workers.drain() {
            worker_handle.stop().await?;
        }
        Ok(())
    }

    pub async fn stop_meta(&mut self) -> Result<()> {
        if let Some(meta) = self.meta.take() {
            meta.stop().await
        } else {
            Ok(())
        }
    }

    pub async fn restart_meta(&mut self, config: MetaServiceConfig) -> Result<()> {
        error!("restart_meta: stop meta");
        self.stop_meta().await?;
        error!("restart_meta: stop meta finished");
        self.meta = Some(MetaServiceHandle::new(config, false).await?);
        for (_, w) in self.workers.iter() {
            let shard_write_manager = w.shard_write_manager.clone();
            if let DispatchShardWriter::Mock(mock) = &shard_write_manager {
                mock.reset_meta(self.get_meta_handle().meta_service.clone())
                    .await;
            }
        }
        error!("restart_meta: restart meta finished");
        Ok(())
    }

    pub async fn stop(&mut self) -> Result<()> {
        self.stop_worker().await?;
        self.stop_meta().await?;
        Ok(())
    }

    pub fn get_meta_service(&self) -> Arc<MetaServiceCore> {
        self.get_meta_handle().meta_service.clone()
    }

    pub fn get_worker_node_heartbeat_interval(&self) -> ReadableDuration {
        self.writer_config_template.heartbeat_interval
    }

    pub fn get_worker_node_template_addr(&self) -> String {
        self.writer_config_template.server.addr.clone()
    }

    pub async fn get_worker_node_config(&self, addr: &String) -> ServerConfig {
        self.workers.get(addr).unwrap().config.server.clone()
    }

    pub fn get_one_worker_node_addr(&self) -> String {
        self.workers.iter().next().unwrap().0.clone()
    }

    pub fn get_one_worker_node_handle(&self) -> &WorkerServiceHandle {
        self.workers.iter().next().unwrap().1
    }

    pub fn get_worker_node_handle(&self, addr: &str) -> &WorkerServiceHandle {
        self.workers.get(addr).unwrap()
    }

    pub fn get_meta_config(&self) -> MetaServiceConfig {
        self.get_meta_handle().config.clone()
    }

    pub async fn get_mock_worker(&self, addr: &str) -> Arc<MockShardWriterManager> {
        if let DispatchShardWriter::Mock(ref x) =
            self.get_worker_node_handle(addr).shard_write_manager
        {
            return x.clone();
        }
        panic!("not in mock env");
    }

    pub fn get_meta_internal_client(&self) -> GrpcMetaServiceInternalClient<Channel> {
        self.meta
            .as_ref()
            .unwrap()
            .get_meta_internal_client()
            .clone()
    }

    pub fn get_meta_client(&self) -> GrpcMetaServiceClient<Channel> {
        self.get_meta_handle().get_meta_client().clone()
    }
}

impl E2ETestSuite {
    pub async fn validate_size_histo(&self) {
        let mut total_size = 0;
        let mut total_size2 = 0;
        for s in self.collect_shards_from_restore().await {
            total_size += s.get_total_estimated_bytes();
            total_size2 += s.stats.get_size();
        }
        assert!(SHARD_SIZE_HISTOGRAM.get_sample_sum() > 0.0);
        assert_eq!(total_size, total_size2);
    }

    pub async fn get_index(&self, index_key: &IndexKey) -> Arc<FTSIndex> {
        self.get_meta_service().get_index(index_key).await.unwrap()
    }

    pub async fn compute_total_size(&self, index_key: &IndexKey) -> u64 {
        let sr = self.collect_shards_from_index(index_key).await;
        sr.into_iter().map(|e| e.get_total_estimated_bytes()).sum()
    }

    pub async fn compute_total_rows(&self, index_key: &IndexKey) -> u64 {
        let sr = self.collect_shards_from_index(index_key).await;
        sr.into_iter().map(|e| e.get_total_estimated_rows()).sum()
    }

    pub async fn compute_total_frags(&self, index_key: &IndexKey) -> u64 {
        let sr = self.collect_shards_from_index(index_key).await;
        let mut set: HashSet<String> = Default::default();
        for s in sr {
            for frag in s.manifest.get_fragments() {
                set.insert(frag.frag_path.clone());
            }
        }
        set.len() as u64
    }

    pub async fn print_all_frags(&self, index_key: &IndexKey) {
        let sr = self.collect_shards_from_index(index_key).await;
        let mut hm: HashMap<String, (u64, f64)> = Default::default();
        for s in sr {
            for frag in s.manifest.get_fragments() {
                hm.insert(
                    frag.frag_path.clone(),
                    (frag.property.as_ref().unwrap().count, frag.p),
                );
            }
        }
        info!(
            "{}",
            hm.iter()
                .map(|x| format!("{}/{:?}", x.0, x.1))
                .collect::<Vec<_>>()
                .join("\n")
        );
    }

    pub async fn collect_shards_from_memory(&self) -> Vec<ShardInfo> {
        let mut res = vec![];
        for (index, ..) in self.collect_indexes_from_restore().await {
            let shards = self
                .collect_shards_from_index(&IndexKey::from(&index))
                .await;
            res.extend(shards.into_iter());
        }
        res
    }

    pub async fn collect_shards_from_index(&self, index_key: &IndexKey) -> Vec<ShardInfo> {
        let sr = self
            .get_meta_service()
            .get_index(index_key)
            .await
            .unwrap()
            .clone_shards()
            .await;
        futures::future::join_all(sr.into_iter().map(async |e| e.read().await.clone())).await
    }

    pub async fn compute_index_memory_usage(&self, index_key: &IndexKey) -> usize {
        self.collect_shards_from_index(index_key)
            .await
            .iter()
            .map(|x| x.heap_size_of_children() + std::mem::size_of::<ShardInfo>())
            .sum::<usize>()
    }

    pub async fn collect_shards_from_restore(&self) -> Vec<ShardInfo> {
        self.get_meta_service()
            .meta_store
            .restore_shards()
            .await
            .unwrap()
    }

    pub async fn collect_indexes_from_restore(
        &self,
    ) -> Vec<(IndexInfo, TableInfo, String, String, IndexMetaProperties)> {
        self.get_meta_service()
            .meta_store
            .restore_indexes()
            .await
            .unwrap()
    }

    pub async fn get_shard_from_restore(&self, shard_id: u64) -> Result<ShardInfo> {
        let m_s = self.collect_shards_from_restore().await;
        m_s.into_iter()
            .find(|s| s.shard_id == shard_id)
            .ok_or_else(|| anyhow::anyhow!("restore: shard_id={} not found", shard_id))
    }

    pub async fn get_shard_from_index(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
    ) -> Result<Arc<RwLock<ShardInfo>>> {
        Ok(self.get_shard_and_index(index_key, shard_id).await?.1)
    }

    pub async fn get_shard_and_index(
        &self,
        index_key: &IndexKey,
        shard_id: u64,
    ) -> Result<(Arc<FTSIndex>, Arc<RwLock<ShardInfo>>)> {
        let index = self.get_meta_service().get_index(index_key).await?;
        let shards = index.clone_shards().await;
        for shard in shards {
            if shard.read().await.shard_id == shard_id {
                return Ok((index, shard.clone()));
            }
        }
        Err(anyhow::anyhow!(
            "get_shard_and_index: shard_id={} not found in index {:?}",
            shard_id,
            index_key,
        ))
    }

    pub async fn collect_worker_nodes(&self) -> Vec<(String, Arc<WorkerNodeMeta>)> {
        self.get_meta_service()
            .worker_node_mgr
            .collect_worker_nodes()
            .await
    }

    pub async fn get_first_shard_id(&self) -> u64 {
        self.collect_shards_from_restore()
            .await
            .first()
            .unwrap()
            .shard_id
    }

    pub async fn get_first_shard(&self) -> ShardInfo {
        self.collect_shards_from_restore()
            .await
            .first()
            .unwrap()
            .clone()
    }

    pub async fn get_random_shard(&self) -> ShardInfo {
        let shards = self.collect_shards_from_restore().await;
        let mut rng = rand::rng();
        shards
            .as_slice()
            .choose(&mut rng)
            .cloned()
            .expect("No shards found")
    }

    pub async fn must_create_index_with_req(
        &self,
        create_index_req: CreateIndexRequest,
    ) -> Result<()> {
        let mut att = 0;
        loop {
            let resp = self
                .get_meta_client()
                .create_index(create_index_req.clone())
                .await;
            let res = match resp {
                Err(e) => {
                    info!("error, e={:?}", e);
                    att += 1;
                    tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
                    assert!(att <= 10);
                    continue;
                }
                Ok(o) => o,
            };
            att += 1;
            if res.get_ref().status != 0 {
                info!("create index error {:?}", res.get_ref().error_message);
                att += 1;
                tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
                assert!(att <= 10);
            } else {
                break;
            }
        }

        Ok(())
    }

    pub async fn must_create_index(&self, index_key: &IndexKey) -> Result<()> {
        let create_index_req = make_create_index_req(index_key);
        self.must_create_index_with_req(create_index_req).await
    }

    pub async fn must_create_partition_index(
        &self,
        logical_index_key: &IndexKey,
        physical_table_ids: &[i64],
    ) -> Result<()> {
        let create_index_req =
            make_create_partition_index_req(logical_index_key, physical_table_ids);
        self.must_create_index_with_req(create_index_req).await
    }

    pub async fn must_drop_index(&self, index_key: &IndexKey) -> Result<()> {
        let drop_index_req = tici_grpc::DropIndexRequest {
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            keyspace_id: index_key.keyspace_id,
        };

        let mut att = 0;
        loop {
            let res = self
                .get_meta_client()
                .drop_index(drop_index_req.clone())
                .await?;
            att += 1;
            if res.get_ref().status != 0 {
                info!("drop index error {:?}", res.get_ref().error_message);
                att += 1;
                tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
                assert!(att <= 10);
            } else {
                break;
            }
        }

        assert!(self.get_meta_service().get_index(index_key).await.is_err());
        for (_, wn) in self.collect_worker_nodes().await.iter() {
            assert!(!wn.contains_shard(index_key).await);
        }
        Ok(())
    }

    pub async fn wait_heartbeat(&self) {
        let hb_interval = self.get_worker_node_heartbeat_interval().as_millis();
        info!("Waiting for 2 heartbeat interval: 2 * {} ms", hb_interval);
        std::thread::sleep(std::time::Duration::from_millis(hb_interval * 2));
    }

    pub async fn validate_add_index_result(&self, index_key: &IndexKey) {
        let mut in_index = self.collect_shards_from_index(index_key).await;
        assert!(!in_index.is_empty(), "Expected at least one shard in index");
        let in_restore = self.collect_shards_from_restore().await;
        let mut in_restore = in_restore
            .into_iter()
            .filter(|x| x.index_id() == index_key.index_id && x.table_id() == index_key.table_id)
            .collect::<Vec<_>>();
        in_index.sort_by_key(|x| x.shard_id);
        in_restore.sort_by_key(|x| x.shard_id);
        assert_eq!(
            in_restore, in_index,
            "Shards in index and restore should match"
        );
        assert_eq!(in_index[0].shard_id, in_restore[0].shard_id);
        assert_eq!(
            in_index[0].range.sortable_start(),
            in_restore[0].range.sortable_start()
        );
    }

    pub async fn must_split_shard_simple(
        &self,
        index_key: &IndexKey,
        shard_range: HandleRange,
        split_point: Option<&[u8]>,
    ) -> Result<Vec<ShardInfo>> {
        let start = shard_range.sortable_start();
        let end = shard_range.sortable_end();
        let start1 = shard_range.encoded_start();
        let end1 = shard_range.encoded_end();
        let ranges = if let Some(x) = split_point {
            let key1 = table_raw_key(index_key, x);
            let key = RangeBound::Key(key1.clone());
            assert!(key > start);
            assert!(key < end);
            vec![
                HandleRange::new(start1, key1.clone()),
                HandleRange::new(key1.clone(), end1),
            ]
        } else {
            let mid_handle = midpoint_bytes(start1.as_slice(), end1.as_slice());
            vec![
                HandleRange::new_from_sorted(
                    RangeBound::Key(start1.clone()),
                    RangeBound::Key(mid_handle.clone()),
                ),
                HandleRange::new_from_sorted(
                    RangeBound::Key(mid_handle.clone()),
                    RangeBound::Key(end1.clone()),
                ),
            ]
        };
        let res = self
            .must_manual_split_shard(index_key, start, ranges, 0)
            .await?;
        Ok(res)
    }

    pub async fn must_manual_split_shard(
        &self,
        index_key: &IndexKey,
        start_key: tici::RangeBound,
        new_ranges: Vec<HandleRange>,
        batch: usize,
    ) -> Result<Vec<ShardInfo>> {
        let new_wn = self.collect_worker_nodes().await.first().unwrap().1.clone();
        let index = self.get_meta_service().get_index(index_key).await?;
        let origin_len = index.shards_len().await;
        let expected_len = origin_len + new_ranges.len() - 1;
        let shard = index
            .clone_shard(&start_key)
            .await
            .ok_or(anyhow!("not found"))?;
        let origin_wn = self
            .get_meta_service()
            .worker_node_mgr
            .get_worker_node(&shard.read().await.writer)
            .await
            .unwrap();
        assert!(origin_wn.shards_len().await > 0);
        let mut expected_wn_len = origin_wn.shards_len().await - 1;
        let new_shards = self
            .get_meta_service()
            .shard_task_manager
            .debug_split_shards(shard.clone(), new_ranges.clone(), batch)
            .await?;
        for s in new_shards.iter() {
            if s.writer == origin_wn.addr {
                expected_wn_len += 1;
            }
        }
        assert_eq!(new_shards.len(), new_ranges.len());
        assert_eq!(
            index.shards_len().await,
            expected_len,
            "expected_len={}, new_ranges={:?}",
            expected_len,
            new_ranges
        );
        assert_eq!(new_wn.shards_len().await, expected_wn_len);
        Ok(new_shards)
    }

    pub async fn must_merge_shard(
        &self,
        shards: &[Arc<RwLock<ShardInfo>>],
        _wn: String,
    ) -> Arc<RwLock<ShardInfo>> {
        let reason = MergeShardReason::Manual;
        match self
            .get_meta_service()
            .shard_task_manager
            .debug_maybe_merge_shards(shards.to_vec(), &reason)
            .await
            .unwrap()
        {
            ShardMergeResult::Merged(merged) => merged,
            ShardMergeResult::NoNeed => panic!("must_merge_shard: unexpectedly got NoNeed"),
            ShardMergeResult::NeedMergeOverlapCompaction(_) => {
                panic!("must_merge_shard: unexpectedly needs overlap compaction")
            }
        }
    }

    pub async fn must_merge_shard_by_id(
        &self,
        index_key: &IndexKey,
        shards: &[RangeBound],
        wn: String,
    ) -> Arc<RwLock<ShardInfo>> {
        let mut new_shards_arc = vec![];
        for shard in shards.iter() {
            let x = self
                .get_meta_service()
                .shard_task_manager
                .get_shard(index_key, shard)
                .await
                .unwrap();
            new_shards_arc.push(x);
        }
        self.must_merge_shard(new_shards_arc.as_slice(), wn).await
    }

    pub async fn must_not_merge_shard_for_overlap(
        &self,
        shards: &[Arc<RwLock<ShardInfo>>],
        _wn: Arc<WorkerNodeMeta>,
    ) {
        let reason = MergeShardReason::Manual;
        match self
            .get_meta_service()
            .shard_task_manager
            .debug_maybe_merge_shards(shards.to_vec(), &reason)
            .await
            .unwrap()
        {
            ShardMergeResult::NeedMergeOverlapCompaction(_) => {}
            ShardMergeResult::NoNeed => {
                panic!("must_not_merge_shard_for_overlap: unexpectedly got NoNeed")
            }
            ShardMergeResult::Merged(_) => {
                panic!("must_not_merge_shard_for_overlap: unexpectedly merged")
            }
        }
    }

    pub async fn collect_shards_by_worker_node(
        &self,
        addr: &String,
    ) -> Vec<(u64, WorkerNodeShardStatus)> {
        self.workers
            .get(addr)
            .unwrap()
            .worker_node_service
            .shard_statuses()
            .into_iter()
            .map(|status| (status.shard_id, status))
            .collect::<Vec<_>>()
    }
}