use std::{
    collections::{HashMap, HashSet},
    fmt::Debug,
    net::SocketAddr,
    sync::Arc,
    time::{Duration, Instant},
};

use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use flexi_logger::LoggerHandle;
use futures::stream::{self, StreamExt};
use lazy_static::lazy_static;
use log::{error, info, warn};
use scopeguard::defer;
use storage_backend::storage::{STORAGE_SERVICE_READER, create_backend_storage};
use tici_common::{config::GrpcConfig, index::IndexKey, logger::init_logger};
use tici_grpc::{
    FragMeta, GetDebugCacheInfoRequest, GetDebugCacheInfoResponse, GrpcReaderService,
    GrpcReaderServiceServer, IndexInfo, KeyRange, ShardManifestHeader, TableInfo,
    WarmupShardRequest, WarmupShardResponse, metrics::begin_grpc_handle_metrics,
};
use tici_shard::{
    FragMetaVersion, FragMetaWrapper, ShardMirrorMeta, ShardMirrorMetaVersion,
    metadata::metadata_store::{self, MetadataStore},
    reader::{
        dedup_persist::PersistedShardDedupBitsets, frag_data_cache::FragDataCache,
        frag_reader::FragReaderMeta, shard_mirror::ShardMirror,
        shard_mirror_manager::ShardMirrorManager,
    },
    utils::{
        dedup_bitset::FragAliveBitSetBytes,
        meta_client_pool::{GrpcMetaClientPool, MetaClientPool},
        schema::IndexMeta,
        shard_tracing_id,
    },
};
use tokio::{
    signal::unix::{SignalKind, signal},
    task::JoinHandle,
};
use tokio_util::sync::CancellationToken;
use tonic::{Request, Response, Status, transport::Server};

use crate::{
    config::Config, heartbeat::HeartbeatManager, metrics::observe_reader_warmup_duration,
    set_global_shard_mirror_manager,
};

#[derive(Clone, Copy, Debug, Default)]
struct AddShardAndPersistStats {
    frag_count: usize,
    build_frag_reader_metas: Duration,
    dedup_frags: Option<Duration>,
    add_shard: Duration,
}

#[derive(Debug)]
struct WarmupLogContext {
    shard_tracing_id: String,
    requested_seq: u64,
    current_warming_up_shards: usize,
    frag_upload_bucket: String,
    cleanup_metadata: Option<Duration>,
    fetch_prepare_meta: Option<Duration>,
    add_shard_and_persist: Option<AddShardAndPersistStats>,
    start: Instant,
}

impl WarmupLogContext {
    fn new(
        shard_tracing_id: String,
        requested_seq: u64,
        current_warming_up_shards: usize,
        frag_upload_bucket: String,
    ) -> Self {
        Self {
            shard_tracing_id,
            requested_seq,
            current_warming_up_shards,
            frag_upload_bucket,
            cleanup_metadata: None,
            fetch_prepare_meta: None,
            add_shard_and_persist: None,
            start: Instant::now(),
        }
    }

    fn log_success(&self, resolved_seq: u64) {
        let total = self.start.elapsed();
        let add_shard_and_persist = self.add_shard_and_persist.unwrap_or_default();
        info!(
            "warmup finished, shard={}, requested_seq={}, resolved_seq={}, frag_count={}, \
             bucket={}, current_warming_up_shards={}, total={:.3}s, cleanup_metadata={:.3}s, \
             fetch_prepare_meta={:.3}s, build_frag_reader_metas={:.3}s, dedup_frags={:.3}s, \
             add_shard={:.3}s",
            self.shard_tracing_id,
            self.requested_seq,
            resolved_seq,
            add_shard_and_persist.frag_count,
            self.frag_upload_bucket,
            self.current_warming_up_shards,
            total.as_secs_f64(),
            duration_secs(self.cleanup_metadata),
            duration_secs(self.fetch_prepare_meta),
            add_shard_and_persist.build_frag_reader_metas.as_secs_f64(),
            duration_secs(add_shard_and_persist.dedup_frags),
            add_shard_and_persist.add_shard.as_secs_f64(),
        );
    }

    fn log_failure<E>(&self, stage: &'static str, err: &E)
    where
        E: Debug + ?Sized,
    {
        error!(
            "warmup failed, shard={}, requested_seq={}, bucket={}, current_warming_up_shards={}, \
             stage={}, err={:?}",
            self.shard_tracing_id,
            self.requested_seq,
            self.frag_upload_bucket,
            self.current_warming_up_shards,
            stage,
            err,
        );
    }
}

fn duration_secs(duration: Option<Duration>) -> f64 {
    duration.map_or(0.0, |value| value.as_secs_f64())
}

#[derive(Clone)]
pub struct ReaderServiceImpl {
    warming_up_shards: Arc<parking_lot::Mutex<HashSet<u64>>>,
    frag_data_cache: Arc<FragDataCache>,
    pub heartbeat_manager: Arc<HeartbeatManager>,
    config: Config,
    metadata_store: Arc<MetadataStore>,
    meta_client_pool: Arc<dyn MetaClientPool>,
    pub shard_mirror_manager: Arc<ShardMirrorManager>,
}

impl ReaderServiceImpl {
    pub async fn new(
        config: &Config,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) -> Result<Self> {
        config.tici.storage.init_reader_data_dir()?;

        let backend_storage =
            create_backend_storage(&config.tici.s3, STORAGE_SERVICE_READER).await?;
        let frag_data_cache = Arc::new(FragDataCache::new(
            backend_storage,
            config.tici.storage.frag_base_path(),
            config.tici.storage.temporary_path(),
            config.tici.frag_reader.remove_zero_ref_frag_delay.0,
            config.tici.frag_reader.disk_reserve_size,
        )?);

        let metadata_store = Arc::new(MetadataStore::new(&config.tici.storage.metadata_path())?);

        // Ignore restore errors, recover from warmup shards later.
        let _ = Self::restore(
            shard_mirror_manager.clone(),
            frag_data_cache.clone(),
            metadata_store.clone(),
            config.tici.storage.temporary_path(),
        )
        .await
        .inspect_err(|e| warn!("restore failed, err={:?}", e));

        let meta_client_pool = Arc::new(
            GrpcMetaClientPool::new(
                &config.pd.get_pd_addresses(),
                &config.tici.security,
                config.tici.reader_node.grpc_client.clone(),
                -1, // Wait forever for meta service
            )
            .await?,
        );
        info!("Starting heartbeat manager");
        let manager = Arc::new(
            HeartbeatManager::new(
                config.clone(),
                frag_data_cache.clone(),
                meta_client_pool.clone(),
                metadata_store.clone(),
                shard_mirror_manager.clone(),
            )
            .await?,
        );
        info!("Heartbeat manager started");
        Ok(Self {
            warming_up_shards: Arc::new(parking_lot::Mutex::new(HashSet::new())),
            frag_data_cache,
            heartbeat_manager: manager,
            config: config.clone(),
            metadata_store,
            meta_client_pool,
            shard_mirror_manager,
        })
    }

    async fn fetch_and_prepare_meta(
        &self,
        table_info: &TableInfo,
        index_info: &IndexInfo,
        shard_info: &ShardManifestHeader,
        frag_bucket: &str,
    ) -> Result<(IndexMeta, ShardMirrorMeta, Vec<FragMeta>)> {
        let mut shard_mirror_meta = ShardMirrorMeta {
            version: ShardMirrorMetaVersion::LATEST,
            index_key: IndexKey::from(index_info),
            shard_header: shard_info.clone(),
            frag_bucket: frag_bucket.to_string(),
        };

        let mut meta_client = self.meta_client_pool.get_meta_client().await?;

        let key_range = KeyRange {
            start_key: shard_info.start_key.clone(),
            end_key: shard_info.end_key.clone(),
        };
        let (frag_metas, seq) = meta_client
            .get_frags_by_keyrange(
                shard_info.shard_id,
                &IndexKey::new(
                    table_info.keyspace_id,
                    table_info.table_id,
                    index_info.index_id,
                ),
                key_range,
            )
            .await?;
        if seq < shard_info.seq {
            bail!(
                "{} seq of frag metas is less than warmup request: {} < {}",
                shard_mirror_meta.shard_tracing_id(),
                seq,
                shard_info.seq
            );
        }
        shard_mirror_meta.shard_header.seq = seq;

        let index_meta = IndexMeta::new(table_info, index_info)?;

        Ok((index_meta, shard_mirror_meta, frag_metas))
    }

    #[allow(clippy::too_many_arguments)]
    async fn add_shard_and_persist(
        shard_mirror_manager: Arc<ShardMirrorManager>,
        shard_mirror_meta: ShardMirrorMeta,
        index_meta: IndexMeta,
        frag_metas: Vec<FragMeta>,
        metadata_store: Arc<MetadataStore>,
        frag_data_cache: Arc<FragDataCache>,
        obsolete_keys_temp_path: String,
    ) -> Result<AddShardAndPersistStats> {
        let mut stats = AddShardAndPersistStats {
            frag_count: frag_metas.len(),
            ..Default::default()
        };
        let shard_range = shard_mirror_meta.shard_range()?;
        let shard_id = shard_mirror_meta.shard_header.shard_id;
        let shard_tracing_id = shard_mirror_meta.shard_tracing_id();

        let build_frag_reader_start = Instant::now();
        let mut frag_reader_metas = Vec::with_capacity(frag_metas.len());
        for frag_meta in frag_metas.iter() {
            let frag_reader_meta = FragReaderMeta::from_frag_data_cache(
                frag_meta,
                frag_data_cache.clone(),
                &shard_mirror_meta.frag_bucket,
            )
            .await?;
            frag_reader_metas.push(frag_reader_meta);
        }
        stats.build_frag_reader_metas = build_frag_reader_start.elapsed();

        let persisted_obsolete_doc_bitsets = Self::load_persisted_dedup_bitsets(
            &metadata_store,
            &shard_mirror_meta,
            &frag_metas,
            &frag_reader_metas,
        );

        let (obsolete_doc_bitsets, persisted_dedup_bitsets) =
            if let Some(obsolete_doc_bitsets) = persisted_obsolete_doc_bitsets {
                (obsolete_doc_bitsets, None)
            } else {
                let dedup_start = Instant::now();
                let dedup_output = ShardMirror::dedup_frags(
                    shard_tracing_id.clone(),
                    shard_range.clone(),
                    frag_reader_metas.clone(),
                    obsolete_keys_temp_path,
                )
                .await??;
                stats.dedup_frags = Some(dedup_start.elapsed());
                let persisted_dedup_bitsets = match PersistedShardDedupBitsets::new(
                    &shard_mirror_meta.index_key,
                    shard_id,
                    &shard_range,
                    &frag_metas,
                    &dedup_output.obsolete_doc_bitsets,
                )
                .and_then(|persisted| persisted.encode_compressed())
                {
                    Ok(persisted) => Some(persisted),
                    Err(err) => {
                        warn!(
                            "build persisted shard dedup bitsets failed, shard={}, err={:?}",
                            shard_tracing_id, err
                        );
                        None
                    }
                };
                (dedup_output.obsolete_doc_bitsets, persisted_dedup_bitsets)
            };

        let add_shard_start = Instant::now();
        shard_mirror_manager
            .add_shard(
                shard_mirror_meta.clone(),
                index_meta.clone(),
                frag_reader_metas,
                obsolete_doc_bitsets,
            )
            .await?;
        stats.add_shard = add_shard_start.elapsed();

        if let Some(persisted_dedup_bitsets) = persisted_dedup_bitsets.as_ref() {
            if let Err(err) = async {
                metadata_store
                    .set_shard_dedup_bitset(shard_id, persisted_dedup_bitsets)
                    .await?
                    .await??;
                Ok::<(), anyhow::Error>(())
            }
            .await
            {
                warn!(
                    "persist shard dedup bitsets failed, shard={}, err={:?}",
                    shard_tracing_id, err
                );
            }
        }

        Ok(stats)
    }

    async fn submit_shard_metadata_persist(
        metadata_store: &Arc<MetadataStore>,
        shard_mirror_meta: &ShardMirrorMeta,
        index_meta: &IndexMeta,
        frag_metas: &[FragMeta],
    ) -> Result<tokio::sync::oneshot::Receiver<Result<()>>> {
        let frag_meta_wrapper = FragMetaWrapper {
            version: FragMetaVersion::LATEST,
            seq: shard_mirror_meta.shard_header.seq,
            frag_metas: frag_metas.to_vec(),
        };
        metadata_store
            .persist_reader_local_shard_metadata(shard_mirror_meta, index_meta, &frag_meta_wrapper)
            .await
    }

    #[cfg(test)]
    async fn persist_shard_metadata(
        metadata_store: &Arc<MetadataStore>,
        shard_mirror_meta: &ShardMirrorMeta,
        index_meta: &IndexMeta,
        frag_metas: &[FragMeta],
    ) -> Result<()> {
        let rx = Self::submit_shard_metadata_persist(
            metadata_store,
            shard_mirror_meta,
            index_meta,
            frag_metas,
        )
        .await?;
        rx.await??;

        Ok(())
    }

    fn load_persisted_dedup_bitsets(
        metadata_store: &Arc<MetadataStore>,
        shard_mirror_meta: &ShardMirrorMeta,
        frag_metas: &[FragMeta],
        frag_reader_metas: &[FragReaderMeta],
    ) -> Option<Vec<FragAliveBitSetBytes>> {
        let shard_id = shard_mirror_meta.shard_header.shard_id;
        let shard_tracing_id = shard_mirror_meta.shard_tracing_id();
        let persisted_bytes = match metadata_store.get_shard_dedup_bitset(shard_id) {
            Ok(bytes) => bytes,
            Err(err) => {
                warn!(
                    "read persisted shard dedup bitset failed, shard={}, err={:?}",
                    shard_tracing_id, err
                );
                return None;
            }
        }?;

        let persisted = match PersistedShardDedupBitsets::decode_compressed(&persisted_bytes) {
            Ok(persisted) => persisted,
            Err(err) => {
                warn!(
                    "decode persisted shard dedup bitsets failed, shard={}, err={:?}",
                    shard_tracing_id, err
                );
                return None;
            }
        };

        let shard_range = match shard_mirror_meta.shard_range() {
            Ok(r) => r,
            Err(err) => {
                warn!(
                    "load persisted dedup bitsets: shard_range failed, shard={}, err={:?}",
                    shard_tracing_id, err
                );
                return None;
            }
        };

        match persisted.validate_and_extract(
            &shard_mirror_meta.index_key,
            shard_id,
            &shard_range,
            frag_metas,
            frag_reader_metas,
        ) {
            Ok(obsolete_doc_bitsets) => {
                info!(
                    "reuse persisted shard dedup bitsets, shard={}, frag_count={}",
                    shard_tracing_id,
                    obsolete_doc_bitsets.len()
                );
                Some(obsolete_doc_bitsets)
            }
            Err(err) => {
                warn!(
                    "validate persisted shard dedup bitsets failed, shard={}, err={:?}",
                    shard_tracing_id, err
                );
                None
            }
        }
    }

    async fn restore(
        shard_mirror_manager: Arc<ShardMirrorManager>,
        frag_data_cache: Arc<FragDataCache>,
        metadata_store: Arc<MetadataStore>,
        obsolete_keys_temp_path: String,
    ) -> Result<()> {
        let start = Instant::now();
        let index_metas = metadata_store
            .scan_index_metas_skip_invalid()?
            .into_iter()
            .collect::<HashMap<String, IndexMeta>>();
        let frag_metas = metadata_store
            .scan_frag_metas_skip_invalid()?
            .into_iter()
            .collect::<HashMap<String, FragMetaWrapper>>();
        let shard_mirror_metas = metadata_store.scan_shard_mirror_metas_skip_invalid()?;
        let shard_count = shard_mirror_metas.len();
        info!(
            "restore: scan meta data: {} index_metas, {} frag_metas, {} shard_mirror_metas, use {:?}",
            index_metas.len(),
            frag_metas.len(),
            shard_count,
            start.elapsed()
        );

        let parallelism = std::thread::available_parallelism()?;
        let results = stream::iter(shard_mirror_metas)
            .map(|(_, mut shard_meta)| async {
                let index_meta_key = metadata_store::index_meta_key(&shard_meta.index_key);
                let index_meta = index_metas.get(&index_meta_key).cloned().ok_or_else(|| {
                    anyhow!(
                        "index meta not found for {:?}",
                        shard_meta.shard_tracing_id()
                    )
                })?;
                let frag_meta_key = metadata_store::frag_meta_key(shard_meta.shard_header.shard_id);
                let frag_meta_wrapper =
                    frag_metas.get(&frag_meta_key).cloned().ok_or_else(|| {
                        anyhow!(
                            "frag meta not found for {:?}",
                            shard_meta.shard_tracing_id()
                        )
                    })?;
                shard_meta.shard_header.seq = frag_meta_wrapper.seq;
                let shard_mirror_manager = shard_mirror_manager.clone();
                let frag_data_cache = frag_data_cache.clone();
                let metadata_store = metadata_store.clone();
                let obsolete_keys_temp_path = obsolete_keys_temp_path.clone();
                tokio::spawn(async move {
                    Self::add_shard_and_persist(
                        shard_mirror_manager,
                        shard_meta,
                        index_meta,
                        frag_meta_wrapper.frag_metas,
                        metadata_store,
                        frag_data_cache,
                        obsolete_keys_temp_path,
                    )
                    .await
                    .map(|_| ())
                })
                .await?
            })
            .buffer_unordered(parallelism.get())
            .collect::<Vec<Result<()>>>()
            .await;

        let mut err_cnt = 0;
        for res in results {
            // Ignore the error
            let _ = res.inspect_err(|e| {
                err_cnt += 1;
                warn!("restore: add shard failed: {:?}", e)
            });
        }
        info!(
            "restore: done, {} parallelism, {} shards, {} errors, use {:?}",
            parallelism.get(),
            shard_count,
            err_cnt,
            start.elapsed()
        );
        Ok(())
    }
}

#[async_trait]
impl GrpcReaderService for ReaderServiceImpl {
    async fn warmup_shard(
        &self,
        request: Request<WarmupShardRequest>,
    ) -> Result<Response<WarmupShardResponse>, Status> {
        let rpc_metrics = begin_grpc_handle_metrics("reader", "warmup_shard");
        let req = request.into_inner();
        let shard_info = req.shard.as_ref().ok_or_else(|| {
            rpc_metrics.mark_error();
            Status::invalid_argument("Shard info is required")
        })?;
        let index_info = req.index_info.as_ref().ok_or_else(|| {
            rpc_metrics.mark_error();
            Status::invalid_argument("Index info is required")
        })?;
        let table_info = req.table_info.as_ref().ok_or_else(|| {
            rpc_metrics.mark_error();
            Status::invalid_argument("Table info is required")
        })?;
        let shard_tracing_id = shard_tracing_id(
            &IndexKey::new(
                table_info.keyspace_id,
                table_info.table_id,
                index_info.index_id,
            ),
            shard_info.shard_id,
        );

        // Avoid concurrent warmup requests of the same shard.
        let current_warming_up_shards = {
            let mut warming_up_shards = self.warming_up_shards.lock();
            if warming_up_shards.contains(&shard_info.shard_id) {
                rpc_metrics.mark_error();
                return Err(Status::invalid_argument(format!(
                    "{} is warming up",
                    shard_tracing_id
                )));
            }
            warming_up_shards.insert(shard_info.shard_id);
            warming_up_shards.len()
        };
        let mut warmup_log_ctx = WarmupLogContext::new(
            shard_tracing_id.clone(),
            shard_info.seq,
            current_warming_up_shards,
            req.frag_upload_bucket.clone(),
        );
        defer! {
            let mut warming_up_shards = self.warming_up_shards.lock();
            warming_up_shards.remove(&shard_info.shard_id);
        }

        // Check if the shard is already cached.
        if self
            .shard_mirror_manager
            .check_shards(&[shard_info.shard_id])[0]
        {
            rpc_metrics.mark_error();
            return Err(Status::invalid_argument(format!(
                "{} is already cached",
                shard_tracing_id
            )));
        }

        info!(
            "Received warmup request for {}, seq={}, bucket={}, current_warming_up_shards={}",
            shard_tracing_id, shard_info.seq, req.frag_upload_bucket, current_warming_up_shards
        );

        // Cleanup metadata.
        let cleanup_metadata_start = Instant::now();
        let cleanup_res: anyhow::Result<()> = async {
            let rx = self
                .metadata_store
                .remove_reader_local_shard_metadata(shard_info.shard_id)
                .await?;
            rx.await??;
            Ok(())
        }
        .await;
        cleanup_res.map_err(|e| {
            rpc_metrics.mark_error();
            warmup_log_ctx.log_failure("cleanup_metadata", &e);
            Status::internal(format!(
                "Failed to cleanup metadata for {}, err={}",
                shard_tracing_id, e
            ))
        })?;
        let cleanup_metadata_elapsed = cleanup_metadata_start.elapsed();
        warmup_log_ctx.cleanup_metadata = Some(cleanup_metadata_elapsed);
        observe_reader_warmup_duration("cleanup_metadata", cleanup_metadata_elapsed);

        // Fetch and prepare meta.
        let fetch_prepare_meta_start = Instant::now();
        let (index_meta, shard_mirror_meta, frag_metas) = self
            .fetch_and_prepare_meta(table_info, index_info, shard_info, &req.frag_upload_bucket)
            .await
            .map_err(|e| {
                rpc_metrics.mark_error();
                warmup_log_ctx.log_failure("fetch_prepare_meta", &e);
                Status::internal(format!(
                    "Failed to fetch_and_prepare_meta for {}, err={}",
                    shard_tracing_id, e
                ))
            })?;
        let fetch_prepare_meta_elapsed = fetch_prepare_meta_start.elapsed();
        warmup_log_ctx.fetch_prepare_meta = Some(fetch_prepare_meta_elapsed);
        observe_reader_warmup_duration("fetch_prepare_meta", fetch_prepare_meta_elapsed);

        let persist_shard_mirror_meta = shard_mirror_meta.clone();
        let persist_index_meta = index_meta.clone();
        let persist_frag_metas = frag_metas.clone();
        let resolved_seq = shard_mirror_meta.shard_header.seq;
        let add_shard_and_persist_stats = Self::add_shard_and_persist(
            self.shard_mirror_manager.clone(),
            shard_mirror_meta,
            index_meta,
            frag_metas,
            self.metadata_store.clone(),
            self.frag_data_cache.clone(),
            self.config.tici.storage.temporary_path(),
        )
        .await
        .map_err(|e| {
            rpc_metrics.mark_error();
            warmup_log_ctx.log_failure("add_shard_and_persist", &e);
            Status::internal(format!("Failed to add shard, err={}", e))
        })?;
        observe_reader_warmup_duration(
            "build_frag_reader_metas",
            add_shard_and_persist_stats.build_frag_reader_metas,
        );
        if let Some(dedup_frags) = add_shard_and_persist_stats.dedup_frags {
            observe_reader_warmup_duration("dedup_frags", dedup_frags);
        }
        observe_reader_warmup_duration("add_shard", add_shard_and_persist_stats.add_shard);

        if let Err(err) = async {
            let rx = Self::submit_shard_metadata_persist(
                &self.metadata_store,
                &persist_shard_mirror_meta,
                &persist_index_meta,
                &persist_frag_metas,
            )
            .await?;
            rx.await??;
            Ok::<(), anyhow::Error>(())
        }
        .await
        {
            warn!(
                "persist warmup shard metadata failed, shard={}, err={:?}",
                shard_tracing_id, err
            );
        }

        warmup_log_ctx.add_shard_and_persist = Some(add_shard_and_persist_stats);
        warmup_log_ctx.log_success(resolved_seq);

        Ok(Response::new(WarmupShardResponse {
            status: 0,
            ..Default::default()
        }))
    }

    async fn get_debug_cache_info(
        &self,
        _request: Request<GetDebugCacheInfoRequest>,
    ) -> Result<Response<GetDebugCacheInfoResponse>, Status> {
        let _rpc_metrics = begin_grpc_handle_metrics("reader", "get_debug_cache_info");
        let cached_shard_ids = self.shard_mirror_manager.get_cached_shard_ids();

        info!(
            "Reader {}/{} cache info={} shards cached, cached_shard_ids={:?}",
            self.config.tici.reader_node.addr,
            self.config.tici.reader_node.advertise_addr,
            cached_shard_ids.len(),
            cached_shard_ids
        );

        Ok(Response::new(GetDebugCacheInfoResponse {
            status: 0,
            error_message: String::new(),
            cached_shard_count: cached_shard_ids.len() as i32,
            cached_shard_ids,
        }))
    }
}

async fn wait_for_shutdown_signal() -> Result<()> {
    let mut sigterm =
        signal(SignalKind::terminate()).context("Failed to register SIGTERM handler")?;
    let mut sigint =
        signal(SignalKind::interrupt()).context("Failed to register SIGINT handler")?;

    tokio::select! {
        _ = sigterm.recv() => info!("Received SIGTERM, shutting down..."),
        _ = sigint.recv() => info!("Received SIGINT (Ctrl+C), shutting down..."),
    }
    Ok(())
}

pub struct ReaderServer {
    service_impl: Arc<ReaderServiceImpl>,
    frag_data_gc_task: Option<(JoinHandle<()>, CancellationToken)>,
    grpc_handle: Option<tokio::task::JoinHandle<Result<(), tonic::transport::Error>>>,
}

impl ReaderServer {
    pub async fn new(
        config: &Config,
        shard_mirror_manager: Arc<ShardMirrorManager>,
    ) -> Result<Self> {
        let service_impl = Arc::new(ReaderServiceImpl::new(config, shard_mirror_manager).await?);
        let server = Self {
            service_impl,
            frag_data_gc_task: None,
            grpc_handle: None,
        };
        Ok(server)
    }

    pub fn get_service(&self) -> Arc<ReaderServiceImpl> {
        self.service_impl.clone()
    }

    pub async fn shutdown(
        &mut self,
        grpc_server_close_tx: tokio::sync::oneshot::Sender<()>,
    ) -> Result<()> {
        grpc_server_close_tx
            .send(())
            .map_err(|e| anyhow!("grpc_server_close_tx send failed: {:?}", e))?;
        if let Some((task_handle, cancel_token)) = self.frag_data_gc_task.take() {
            info!("Shutting down frag data gc task");
            cancel_token.cancel();
            task_handle.await?;
            info!("Frag data gc task shut down");
        }
        Ok(())
    }

    pub fn build_server(
        addr: SocketAddr,
        service_impl: ReaderServiceImpl,
        rx: tokio::sync::oneshot::Receiver<()>,
        grpc_config: GrpcConfig,
    ) -> impl Future<Output = Result<(), tonic::transport::Error>> {
        Server::builder()
            .add_service(
                GrpcReaderServiceServer::new(service_impl)
                    .max_decoding_message_size(grpc_config.max_decoding_message_size.0 as usize)
                    .max_encoding_message_size(grpc_config.max_encoding_message_size.0 as usize),
            )
            .serve_with_shutdown(addr, async {
                let _ = rx.await.inspect_err(|e| {
                    error!("grpc_server_close_tx receive failed: {:?}", e);
                });
            })
    }

    pub fn start(&mut self, rx: tokio::sync::oneshot::Receiver<()>) -> Result<()> {
        let s_addr = &self.service_impl.config.tici.reader_node.addr;
        let addr = s_addr
            .parse()
            .with_context(|| format!("Address={}", s_addr))?;

        info!("Starting frag data gc task");
        let cancel_token = CancellationToken::new();
        let task_handle = FragDataCache::spawn_gc_task(
            self.service_impl.frag_data_cache.clone(),
            cancel_token.child_token(),
        );
        self.frag_data_gc_task = Some((task_handle, cancel_token));

        info!("Starting process memory status monitor");
        tici_system::metrics::monitor_process_memory_status();

        info!("Starting reader gRPC server at {}", addr);
        self.grpc_handle = Some(tokio::spawn(Self::build_server(
            addr,
            self.service_impl.as_ref().clone(),
            rx,
            self.service_impl.config.tici.reader_node.grpc.clone(),
        )));

        Ok(())
    }
}

lazy_static! {
    static ref READER_SERVER_HANDLE: parking_lot::Mutex<Option<std::thread::JoinHandle<Result<()>>>> =
        parking_lot::Mutex::new(None);
    static ref READER_LOGGER_HANDLE: parking_lot::Mutex<Option<LoggerHandle>> =
        parking_lot::Mutex::new(None);
}

pub fn start_server(
    conf_path: &str,
    read_node_addr: &str,
    service_addr: &str,
    pd_addr: &str,
) -> Result<()> {
    tici_system::panic_hook::setup_panic_hook();

    let mut config = Config::load(conf_path)
        .with_context(|| format!("Failed to load config from {}", conf_path))?;

    if config.tici.security.is_valid() {
        info!("Use aws_lc_rs for rustls::crypto");
        use rustls::crypto::CryptoProvider;
        CryptoProvider::install_default(rustls::crypto::aws_lc_rs::default_provider())
            .map_err(|e| anyhow!("Failed to install crypto provider: {:?}", e))?;
    }

    let logger_handle = init_logger(&config.tici.logger).with_context(|| {
        format!(
            "Failed to initialize reader server logger: {:?}",
            config.tici.logger
        )
    })?;
    {
        let mut reader_logger_handle = READER_LOGGER_HANDLE.lock();
        *reader_logger_handle = Some(logger_handle);
    }

    info!("Build Info={}", tici_shard::utils::build_info());

    config.tici.reader_node.addr = read_node_addr.to_string();
    config.flash.service_addr = service_addr.to_string();
    config.pd.pd_addr = pd_addr.to_string();
    config.resolve_reader_node_addr()?;
    info!("Config: {:?}", config);

    let rt = tokio::runtime::Runtime::new().context("Failed to create tokio runtime")?;

    let shard_mirror_manager = Arc::new(ShardMirrorManager::new(config.tici.frag_reader.clone()));
    set_global_shard_mirror_manager(shard_mirror_manager.clone())?;

    // TODO(P1) remove this block_on
    let mut reader_server = rt.block_on(ReaderServer::new(&config, shard_mirror_manager))?;

    let handle = std::thread::spawn(move || {
        rt.block_on(async move {
            let (tx, rx) = tokio::sync::oneshot::channel();
            if let Err(e) = reader_server.start(rx) {
                error!("Failed to start reader server, err={:?}", e);
            }
            if let Err(e) = wait_for_shutdown_signal().await {
                error!("Failed to wait reader server shutdown signal, err={:?}", e);
            }
            reader_server.shutdown(tx).await
        })
    });

    let mut reader_handle = READER_SERVER_HANDLE.lock();
    *reader_handle = Some(handle);
    Ok(())
}

pub fn stop_server() {
    let mut reader_handle = READER_SERVER_HANDLE.lock();
    if let Some(handle) = reader_handle.take() {
        let e = handle.join();
        info!("stop_server: {:?}", e);
    }
    let mut reader_logger_handle = READER_LOGGER_HANDLE.lock();
    if let Some(handle) = reader_logger_handle.take() {
        handle.shutdown();
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use anyhow::{Context, Result};
    use tempfile::TempDir;
    use tici_grpc::KeyRange;
    use tici_shard::{
        config::{FragReaderConfig, WorkerNodeConfig},
        metadata::metadata_store::{self, MetadataStore},
        reader::{
            dedup_persist::PersistedShardDedupBitsets, frag_data_cache::FragDataCache,
            query_info::QueryInfo, shard_mirror_manager::ShardMirrorManager,
        },
        test_utils::{
            TEST_HDFSLOGS_INDEX_META, test_create_backend_storage, test_get_inf_range,
            test_get_shard_mirror_meta, test_get_shard_writer_meta, write_test_data,
        },
        utils::meta_client_pool::{MetaClientPool, MockMetaClientPool},
    };
    use tici_system::readable::ReadableSize;

    use crate::reader_server::ReaderServiceImpl;

    async fn corrupt_persisted_dedup_bitset(
        metadata_store: &Arc<MetadataStore>,
        shard_id: u64,
    ) -> Result<()> {
        let persisted_bytes = metadata_store
            .get_shard_dedup_bitset(shard_id)?
            .context("persisted dedup bitsets should exist")?;
        let mut persisted = PersistedShardDedupBitsets::decode_compressed(&persisted_bytes)?;
        let frag = persisted
            .frags
            .first_mut()
            .context("persisted dedup bitsets should contain a frag")?;
        let segment = frag
            .bitset
            .as_mut()
            .context("persisted dedup bitsets should contain a bitset payload")?
            .segments
            .first_mut()
            .context("persisted dedup bitsets should contain a segment")?;
        segment.alive_bitset_bytes = vec![1, 2, 3];
        let corrupted = persisted.encode_compressed()?;
        metadata_store
            .set_shard_dedup_bitset(shard_id, &corrupted)
            .await?
            .await??;
        Ok(())
    }

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

        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new()?;
        let config = WorkerNodeConfig::new_for_test(
            temp_dir
                .path()
                .to_str()
                .context("Failed to get temp dir path")?,
        );
        let backend_storage = test_create_backend_storage()
            .await
            .context("Failed to create backend storage")?;
        let mock_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            mock_client_pool.clone(),
        )
        .await?;

        let mut mock_client_for_read = mock_client_pool.get_meta_client().await?;
        let (frag_metas, seq) = mock_client_for_read
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await?;

        let mut shard_meta = shard_meta.clone();
        shard_meta.shard_header.seq = seq;
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);
        let valid_shard_id = shard_mirror_meta.shard_header.shard_id;
        let invalid_shard_id = valid_shard_id + 1;

        let restore_temp_dir = TempDir::new()?;
        let metadata_path = format!(
            "{}/metadata",
            restore_temp_dir
                .path()
                .to_str()
                .context("Failed to get temp dir path")?
        );
        let metadata_store = Arc::new(MetadataStore::new(&metadata_path)?);
        let invalid_frag_key = metadata_store::frag_meta_key(invalid_shard_id);
        for (key, value) in [
            ("i_invalid", "{bad json"),
            ("f_invalid", "{bad json"),
            ("sm_invalid", "{bad json"),
            (invalid_frag_key.as_str(), "{bad json"),
        ] {
            metadata_store.insert_raw_for_test(key, value)?;
        }
        ReaderServiceImpl::persist_shard_metadata(
            &metadata_store,
            &shard_mirror_meta,
            &index_meta,
            &frag_metas,
        )
        .await?;
        let mut invalid_shard_mirror_meta = shard_mirror_meta.clone();
        invalid_shard_mirror_meta.shard_header.shard_id = invalid_shard_id;
        metadata_store
            .set_shard_mirror_meta(invalid_shard_id, &invalid_shard_mirror_meta)
            .await?
            .await??;

        let default_reader_cfg = FragReaderConfig {
            disk_reserve_size: ReadableSize::mb(64),
            ..FragReaderConfig::default()
        };
        let frag_data_cache_dir = TempDir::new()?;
        let restore_temporary_dir = TempDir::new()?;
        let temporary_path = restore_temporary_dir
            .path()
            .to_str()
            .context("Failed to get temporary dir path")?
            .to_string();
        let frag_data_cache = Arc::new(FragDataCache::new(
            backend_storage,
            frag_data_cache_dir
                .path()
                .to_str()
                .context("Failed to get frag data cache dir path")?
                .to_string(),
            temporary_path.clone(),
            default_reader_cfg.remove_zero_ref_frag_delay.0,
            default_reader_cfg.disk_reserve_size,
        )?);

        let manager = Arc::new(ShardMirrorManager::new(default_reader_cfg));
        ReaderServiceImpl::restore(
            manager.clone(),
            frag_data_cache,
            metadata_store,
            temporary_path,
        )
        .await?;

        let mut cached_shard_ids = manager.get_cached_shard_ids();
        cached_shard_ids.sort_unstable();
        assert_eq!(cached_shard_ids, vec![valid_shard_id]);
        assert!(!cached_shard_ids.contains(&invalid_shard_id));

        Ok(())
    }

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

        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new()?;
        let config = WorkerNodeConfig::new_for_test(
            temp_dir
                .path()
                .to_str()
                .context("Failed to get temp dir path")?,
        );
        let backend_storage = test_create_backend_storage()
            .await
            .context("Failed to create backend storage")?;
        let mock_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            mock_client_pool.clone(),
        )
        .await?;

        let default_reader_cfg = FragReaderConfig {
            disk_reserve_size: ReadableSize::mb(64),
            ..FragReaderConfig::default()
        };

        let frag_data_cache_path = TempDir::new()?
            .path()
            .to_str()
            .context("Failed to get temp dir path")?
            .to_string();
        let frag_data_cache = Arc::new(FragDataCache::new(
            backend_storage.clone(),
            frag_data_cache_path,
            config.storage.temporary_path(),
            default_reader_cfg.remove_zero_ref_frag_delay.0,
            default_reader_cfg.disk_reserve_size,
        )?);
        let mut mock_client_for_read = mock_client_pool.get_meta_client().await?;
        let (frag_metas, seq) = mock_client_for_read
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await?;

        log::info!("frag_metas={:?}", frag_metas);

        let mut shard_meta = shard_meta.clone();
        shard_meta.shard_header.seq = seq;
        let shard_mirror_meta = test_get_shard_mirror_meta(&shard_meta);

        let temp_dir = TempDir::new()?;
        let metadata_path = format!(
            "{}/metadata",
            temp_dir
                .path()
                .to_str()
                .context("Failed to get temp dir path")?
        );
        let metadata_store = Arc::new(MetadataStore::new(&metadata_path)?);

        let manager1 = Arc::new(ShardMirrorManager::new(default_reader_cfg.clone()));
        let stats = ReaderServiceImpl::add_shard_and_persist(
            manager1.clone(),
            shard_mirror_meta.clone(),
            index_meta.clone(),
            frag_metas.clone(),
            metadata_store.clone(),
            frag_data_cache.clone(),
            format!(
                "{}/temp",
                temp_dir
                    .path()
                    .to_str()
                    .context("Failed to get temp dir path")?
            ),
        )
        .await?;
        assert_eq!(stats.frag_count, frag_metas.len());
        assert!(stats.dedup_frags.is_some());
        ReaderServiceImpl::persist_shard_metadata(
            &metadata_store,
            &shard_mirror_meta,
            &index_meta,
            &frag_metas,
        )
        .await?;

        let shard_id = shard_mirror_meta.shard_header.shard_id;

        {
            let shard_mirror = manager1
                .get_shard(shard_id)
                .context("Failed to get shard mirror")?;
            let reader = shard_mirror.get_reader();
            let total_count = 150000;
            let query_info = QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema()?,
                vec![test_get_inf_range()],
            )?;
            let (all_docs, search_stat) = reader.search_with_query_info(query_info)?;
            assert_eq!(search_stat.docs.tantivy_matched, total_count as u64);
            assert_eq!(search_stat.docs.read_range_filtered, 0);
            assert_eq!(search_stat.frags.disjoint, 0);
            assert_eq!(search_stat.frags.overlap, 0);
            let frags_len = frag_metas.len();
            assert_eq!(search_stat.frags.subset, frags_len);

            assert_eq!(all_docs.len(), total_count);
        }

        let manager3 = Arc::new(ShardMirrorManager::new(default_reader_cfg.clone()));
        let cache_hit_stats = ReaderServiceImpl::add_shard_and_persist(
            manager3,
            shard_mirror_meta.clone(),
            index_meta.clone(),
            frag_metas.clone(),
            metadata_store.clone(),
            frag_data_cache.clone(),
            format!(
                "{}/temp",
                temp_dir
                    .path()
                    .to_str()
                    .context("Failed to get temp dir path")?
            ),
        )
        .await?;
        assert!(cache_hit_stats.dedup_frags.is_none());

        corrupt_persisted_dedup_bitset(&metadata_store, shard_id).await?;
        let manager4 = Arc::new(ShardMirrorManager::new(default_reader_cfg.clone()));
        let fallback_stats = ReaderServiceImpl::add_shard_and_persist(
            manager4.clone(),
            shard_mirror_meta.clone(),
            index_meta.clone(),
            frag_metas.clone(),
            metadata_store.clone(),
            frag_data_cache.clone(),
            format!(
                "{}/temp",
                temp_dir
                    .path()
                    .to_str()
                    .context("Failed to get temp dir path")?
            ),
        )
        .await?;
        assert!(fallback_stats.dedup_frags.is_some());
        manager4
            .get_shard(shard_id)
            .context("Failed to get shard mirror after falling back from invalid cache")?;

        let manager2 = Arc::new(ShardMirrorManager::new(default_reader_cfg));
        ReaderServiceImpl::restore(
            manager2.clone(),
            frag_data_cache.clone(),
            metadata_store.clone(),
            format!(
                "{}/temp",
                temp_dir
                    .path()
                    .to_str()
                    .context("Failed to get temp dir path")?
            ),
        )
        .await?;

        {
            let shard_mirror = manager2
                .get_shard(shard_id)
                .context("Failed to get shard mirror")?;
            let reader = shard_mirror.get_reader();
            let total_count = 150000;
            let query_info = QueryInfo::query_all_for_test(
                &index_meta,
                &reader.get_schema()?,
                vec![test_get_inf_range()],
            )?;
            let (all_docs, search_stat) = reader.search_with_query_info(query_info)?;
            assert_eq!(search_stat.docs.tantivy_matched, total_count as u64);
            assert_eq!(search_stat.docs.read_range_filtered, 0);
            assert_eq!(search_stat.frags.disjoint, 0);
            assert_eq!(search_stat.frags.overlap, 0);
            let frags_len = frag_metas.len();
            assert_eq!(search_stat.frags.subset, frags_len);

            assert_eq!(all_docs.len(), total_count);
        }

        Ok(())
    }
}