use std::{collections::HashMap, sync::Arc};

use anyhow::{Result, anyhow, bail};
use log::info;
use tici_grpc::{ReaderCachedShard, ShardManifestHeader};
use tici_system::readable::ReadableSize;

use crate::{
    ShardMirrorMeta,
    config::FragReaderConfig,
    metadata::metadata_store::MetadataStore,
    reader::{
        doc_cache::DocCache, frag_data_cache::FragDataCache, frag_reader::FragReaderMeta,
        shard_mirror::ShardMirror,
    },
    utils::{
        dedup_bitset::FragAliveBitSetBytes, meta_client_pool::MetaClientPool,
        metrics::READER_SHARD_GAUGE, schema::IndexMeta,
    },
};

pub struct ShardStatus {
    pub shard_id: u64,
    pub curr_seq: u64,
    pub new_seq: u64,
    pub is_refresh_failed: bool,
}

#[derive(Default)]
pub struct ShardRefreshSummary {
    pub shard_statuses: Vec<ShardStatus>,
    pub missing_shards: Vec<u64>,
}

// ShardMirrorManager handles the addition, removal, and lookup operations for
// ShardMirrors.
pub struct ShardMirrorManager {
    config: FragReaderConfig,
    shard_mirrors: parking_lot::Mutex<HashMap<u64, Arc<ShardMirror>>>,
    doc_cache: Option<Arc<DocCache>>,
}

#[derive(Default)]
pub struct ShardsSnapshot {
    shards: HashMap<u64, Arc<ShardMirror>>,
}

impl ShardsSnapshot {
    pub fn get_shard(&self, shard_id: u64) -> Option<Arc<ShardMirror>> {
        self.shards.get(&shard_id).cloned()
    }
}

fn build_doc_cache(capacity: ReadableSize) -> Option<Arc<DocCache>> {
    if capacity.0 == 0 {
        info!("global doc cache disabled");
        None
    } else {
        info!("global doc cache enabled, capacity={}", capacity);
        Some(Arc::new(DocCache::new(capacity)))
    }
}

impl ShardMirrorManager {
    pub fn new(config: FragReaderConfig) -> Self {
        let doc_cache = build_doc_cache(config.global_doc_cache_size);
        ShardMirrorManager {
            config,
            shard_mirrors: parking_lot::Mutex::new(HashMap::new()),
            doc_cache,
        }
    }

    // Assume frags are already downloaded to local disk.
    pub async fn add_shard(
        &self,
        shard_meta: ShardMirrorMeta,
        index_meta: IndexMeta,
        frag_reader_metas: Vec<FragReaderMeta>,
        obsolete_doc_bitsets: Vec<FragAliveBitSetBytes>,
    ) -> Result<()> {
        let shard_id = shard_meta.shard_header.shard_id;
        let shard_tracing_id = shard_meta.shard_tracing_id();
        // Fast check
        if self.shard_mirrors.lock().contains_key(&shard_id) {
            info!("Shard {} already exists", shard_tracing_id);
            return Ok(());
        }

        let seq = shard_meta.shard_header.seq;
        let frag_paths = frag_reader_metas
            .iter()
            .map(|meta| meta.local_frag_path.clone())
            .collect::<Vec<_>>();
        let doc_cache = self.doc_cache.clone();
        // Load shard mirror
        let shard_mirror = Arc::new(
            ShardMirror::new_with_sample_counts(
                shard_meta,
                index_meta.clone(),
                frag_reader_metas,
                self.config.estimate_sample_counts(),
                obsolete_doc_bitsets,
                doc_cache,
            )
            .await?,
        );

        let mut shard_mirrors = self.shard_mirrors.lock();
        if shard_mirrors.contains_key(&shard_id) {
            info!("Shard {} already exists", shard_tracing_id);
            return Ok(());
        }
        shard_mirrors.insert(shard_id, shard_mirror);
        READER_SHARD_GAUGE.set(shard_mirrors.len() as i64);
        info!(
            "Shard {} is added, seq={}, frag_paths: {}=>{:?}",
            shard_tracing_id,
            seq,
            frag_paths.len(),
            frag_paths,
        );
        Ok(())
    }

    pub fn get_shard(&self, shard_id: u64) -> Option<Arc<ShardMirror>> {
        self.shard_mirrors.lock().get(&shard_id).cloned()
    }

    pub fn remove_shards(&self, shard_ids: &[u64]) {
        let mut shard_mirrors = self.shard_mirrors.lock();
        for shard_id in shard_ids {
            if let Some(shard_mirror) = shard_mirrors.remove(shard_id) {
                info!("Shard {} is removed", shard_mirror.tracing_id());
            } else {
                info!("Shard {} not exist, no need to remove", shard_id);
            }
        }
        READER_SHARD_GAUGE.set(shard_mirrors.len() as i64);
    }

    pub fn get_cached_shard_ids(&self) -> Vec<u64> {
        self.shard_mirrors.lock().keys().copied().collect()
    }

    pub fn get_cached_shards(&self) -> Vec<ReaderCachedShard> {
        self.shard_mirrors
            .lock()
            .iter()
            .map(|(&shard_id, shard_mirror)| ReaderCachedShard {
                shard_id,
                seq: shard_mirror.get_curr_seq(),
            })
            .collect()
    }

    pub fn check_shards(&self, shard_ids: &[u64]) -> Vec<bool> {
        let shard_mirrors = self.shard_mirrors.lock();
        shard_ids
            .iter()
            .map(|shard_id| shard_mirrors.contains_key(shard_id))
            .collect()
    }

    pub fn check_shards_and_acquire_snapshot(
        &self,
        shard_ids: &[u64],
    ) -> (ShardsSnapshot, Vec<bool>) {
        let mirrors = self.shard_mirrors.lock();
        let mut shards = HashMap::new();
        let mut hits = Vec::with_capacity(shard_ids.len());

        for &id in shard_ids {
            if let Some(mirror) = mirrors.get(&id) {
                shards.insert(id, mirror.clone());
                hits.push(true);
            } else {
                hits.push(false);
            }
        }
        (ShardsSnapshot { shards }, hits)
    }

    pub async fn try_refresh_shard(
        &self,
        resp_shard: &ShardManifestHeader,
        frag_data_cache: Arc<FragDataCache>,
        meta_client_pool: Arc<dyn MetaClientPool>,
        metadata_store: Arc<MetadataStore>,
    ) -> Result<()> {
        let shard_mirror = self
            .get_shard(resp_shard.shard_id)
            .ok_or_else(|| anyhow!("Shard {} not exist", resp_shard.shard_id))?;
        let curr_seq = shard_mirror.get_curr_seq();
        if curr_seq == resp_shard.seq {
            return Ok(());
        }
        if curr_seq > resp_shard.seq {
            bail!(
                "Shard {} current seq {} is greater than repsonsed seq {}, expect smaller",
                shard_mirror.tracing_id(),
                curr_seq,
                resp_shard.seq
            );
        }

        let meta_client = meta_client_pool.get_meta_client().await?;
        shard_mirror
            .refresh(frag_data_cache, meta_client, metadata_store)
            .await?;
        Ok(())
    }

    pub fn shutdown(self) {
        let mut shard_mirrors = self.shard_mirrors.lock();
        info!(
            "Shutting down shard mirror manager with {} shards",
            shard_mirrors.len()
        );
        for (_shard_id, shard_mirror) in shard_mirrors.drain() {
            info!("Shard {} is removed", shard_mirror.tracing_id());
        }
        READER_SHARD_GAUGE.set(0);
        info!("Shard mirror manager is shutdown");
    }
}

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

    use tempfile::TempDir;
    use tici_grpc::KeyRange;
    use tici_system::readable::ReadableSize;
    use tici_test_utils::test_utils::init_logger_for_ut;

    use crate::{
        config::{FragReaderConfig, WorkerNodeConfig},
        metadata::metadata_store::MetadataStore,
        reader::{
            frag_data_cache::FragDataCache, frag_reader::FragReaderMeta, query_info::QueryInfo,
            shard_mirror::ShardMirror, shard_mirror_manager::ShardMirrorManager,
        },
        test_utils::{
            TEST_HDFSLOGS_INDEX_META, TEST_S3_BUCKET, is_sorted_frag, 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},
            s3_frag,
        },
    };

    #[tokio::test]
    async fn test_add_shard_and_read() {
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new().unwrap();
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());
        let shard_id = shard_meta.shard_header.shard_id;

        let backend_storage = test_create_backend_storage().await.unwrap();
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        let mut mock_client = meta_client_pool.get_meta_client().await.unwrap();

        mock_client.reset_shard_progress(shard_id).await.unwrap();
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

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

        let shard_mirror_manager = ShardMirrorManager::new(default_reader_cfg.clone());
        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().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,
            )
            .unwrap(),
        );
        let mut mock_client_for_read = meta_client_pool.get_meta_client().await.unwrap();
        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
            .unwrap();
        let mut frag_reader_metas = Vec::new();
        for frag_meta in &frag_metas {
            frag_reader_metas.push(
                FragReaderMeta::from_frag_data_cache(
                    frag_meta,
                    frag_data_cache.clone(),
                    TEST_S3_BUCKET,
                )
                .await
                .unwrap(),
            );
            assert!(is_sorted_frag(
                &frag_reader_metas.last().unwrap().local_frag_path,
                &index_meta
            ));
        }

        {
            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 dedup_output = ShardMirror::dedup_frags(
                shard_mirror_meta.shard_tracing_id(),
                shard_mirror_meta.shard_range().unwrap(),
                frag_reader_metas.clone(),
                config.storage.temporary_path(),
            )
            .await
            .unwrap()
            .unwrap();
            shard_mirror_manager
                .add_shard(
                    shard_mirror_meta,
                    index_meta.clone(),
                    frag_reader_metas,
                    dedup_output.obsolete_doc_bitsets,
                )
                .await
                .unwrap();
        }
        let shard_mirror = shard_mirror_manager
            .get_shard(shard_meta.shard_header.shard_id)
            .unwrap();

        let reader1 = shard_mirror.get_reader();
        let total_count1 = 150000;
        let query_info = Arc::new(
            QueryInfo::query_all_for_test(
                &index_meta,
                &reader1.get_schema().unwrap(),
                vec![test_get_inf_range()],
            )
            .unwrap(),
        );
        let (all_docs1, search_stat1) = reader1.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(search_stat1.docs.tantivy_matched, total_count1 as u64);
        assert_eq!(search_stat1.docs.read_range_filtered, 0);
        assert_eq!(search_stat1.frags.disjoint, 0);
        assert_eq!(search_stat1.frags.overlap, 0);
        let frags_len1 = frag_metas.len();
        assert_eq!(search_stat1.frags.subset, frags_len1);

        assert_eq!(all_docs1.len(), total_count1);

        mock_client.reset_shard_progress(shard_id).await.unwrap();
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();
        let frags_len2 = mock_client_for_read
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await
            .unwrap()
            .0
            .len();
        assert!(
            frags_len2 > frags_len1,
            "frags_len2={}, frags_len1={}",
            frags_len2,
            frags_len1
        );

        let temp_dir = TempDir::new().unwrap();
        let metadata_path = format!("{}/metadata", temp_dir.path().to_str().unwrap());
        let metadata_store = Arc::new(MetadataStore::new(&metadata_path).unwrap());

        shard_mirror
            .refresh(
                frag_data_cache.clone(),
                meta_client_pool.get_meta_client().await.unwrap(),
                metadata_store.clone(),
            )
            .await
            .unwrap();

        tokio::time::sleep(Duration::from_secs(8)).await; // wait bg refresh

        let reader2 = shard_mirror.get_reader();
        let (all_docs2, search_stat2) = reader2.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(search_stat2.docs.tantivy_matched, total_count1 as u64);
        assert_eq!(search_stat2.docs.read_range_filtered, 0);
        assert_eq!(search_stat2.frags.disjoint, 0);
        assert_eq!(search_stat2.frags.overlap, 0);
        assert_eq!(search_stat2.frags.subset, frags_len2);
        assert_eq!(all_docs2.len(), total_count1);

        mock_client.reset_shard_progress(shard_id).await.unwrap();
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();
        let frags_len3 = mock_client_for_read
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &shard_meta.index_key,
                KeyRange::default(),
            )
            .await
            .unwrap()
            .0
            .len();
        assert!(frags_len3 > frags_len2);

        shard_mirror
            .refresh(
                frag_data_cache.clone(),
                meta_client_pool.get_meta_client().await.unwrap(),
                metadata_store.clone(),
            )
            .await
            .unwrap();

        let reader3 = shard_mirror.get_reader();
        let (all_docs3, search_stat3) = reader3.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(search_stat3.docs.tantivy_matched, total_count1 as u64);
        assert_eq!(search_stat3.docs.read_range_filtered, 0);
        assert_eq!(search_stat3.frags.disjoint, 0);
        assert_eq!(search_stat3.frags.overlap, 0);
        assert_eq!(search_stat3.frags.subset, frags_len3);
        assert_eq!(all_docs3.len(), total_count1);

        let query_info = Arc::new(
            QueryInfo::query_word(
                &index_meta,
                &reader1.get_schema().unwrap(),
                &[test_get_inf_range()],
                "column_2".to_string(),
                "1e17ab02-baf4-4199-88eb-bd036faa8cd3".to_string(),
            )
            .unwrap(),
        );
        let (pks1, search_stat1) = reader1.search_with_query_info(query_info.clone()).unwrap();
        let (pks2, search_stat2) = reader2.search_with_query_info(query_info.clone()).unwrap();
        let (pks3, search_stat3) = reader3.search_with_query_info(query_info.clone()).unwrap();
        assert_eq!(pks1.len(), pks2.len());
        assert_eq!(pks2.len(), pks3.len());
        assert_eq!(search_stat1.frags.subset, frags_len1);
        assert_eq!(search_stat2.frags.subset, frags_len2);
        assert_eq!(search_stat3.frags.subset, frags_len3);
    }

    #[tokio::test]
    async fn test_remove_shard() {
        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().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let config = WorkerNodeConfig::new_for_test(local_data_path);

        let backend_storage = test_create_backend_storage().await.unwrap();

        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

        let shard_mirror_manager = ShardMirrorManager::new(FragReaderConfig::default());
        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path.clone(),
                config.storage.temporary_path(),
                Duration::from_secs(0),
                ReadableSize::mb(64),
            )
            .unwrap(),
        );
        let mut mock_client_for_read = meta_client_pool.get_meta_client().await.unwrap();
        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
            .unwrap();
        let mut frag_reader_metas = Vec::new();
        for frag_meta in &frag_metas {
            frag_reader_metas.push(
                FragReaderMeta::from_frag_data_cache(
                    frag_meta,
                    frag_data_cache.clone(),
                    TEST_S3_BUCKET,
                )
                .await
                .unwrap(),
            );
            assert!(is_sorted_frag(
                &frag_reader_metas.last().unwrap().local_frag_path,
                &index_meta
            ));
        }
        frag_data_cache.remove_zero_ref_frags();
        let mut local_files = s3_frag::list_all_files_in_absolute(&frag_data_cache_path).unwrap();
        local_files.sort();
        assert!(!local_files.is_empty());

        {
            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 dedup_output = ShardMirror::dedup_frags(
                shard_mirror_meta.shard_tracing_id(),
                shard_mirror_meta.shard_range().unwrap(),
                frag_reader_metas.clone(),
                config.storage.temporary_path(),
            )
            .await
            .unwrap()
            .unwrap();
            shard_mirror_manager
                .add_shard(
                    shard_mirror_meta,
                    index_meta.clone(),
                    frag_reader_metas,
                    dedup_output.obsolete_doc_bitsets,
                )
                .await
                .unwrap();
        }

        let shard_mirror = shard_mirror_manager
            .get_shard(shard_meta.shard_header.shard_id)
            .unwrap();
        shard_mirror_manager.remove_shards(&[shard_meta.shard_header.shard_id]);
        assert!(
            shard_mirror_manager
                .get_shard(shard_meta.shard_header.shard_id)
                .is_none()
        );
        frag_data_cache.remove_zero_ref_frags();
        let mut local_files2 = s3_frag::list_all_files_in_absolute(&frag_data_cache_path).unwrap();
        local_files2.sort();
        assert_eq!(local_files, local_files2);

        let shard_reader = shard_mirror.get_reader();
        drop(shard_mirror);

        // Nothing will be removed because shard_reader holds the reference.
        frag_data_cache.remove_zero_ref_frags();
        let mut local_files3 = s3_frag::list_all_files_in_absolute(&frag_data_cache_path).unwrap();
        local_files3.sort();
        assert_eq!(local_files, local_files3);

        // The ​shard_reader​ retains read access despite the removal of the ​shard_mirror.
        let query_info = QueryInfo::query_all_for_test(
            &index_meta,
            &shard_reader.get_schema().unwrap(),
            vec![test_get_inf_range()],
        )
        .unwrap();
        let (all_docs1, _search_stat1) = shard_reader.search_with_query_info(query_info).unwrap();
        assert_eq!(all_docs1.len(), 150000);

        drop(shard_reader);
        // After shard_reader is dropped, reference count became 0. All data will be
        // removed.
        frag_data_cache.remove_zero_ref_frags();
        let local_files4 = s3_frag::list_all_files_in_absolute(&frag_data_cache_path).unwrap();
        assert!(local_files4.is_empty());
    }

    #[tokio::test]
    async fn test_snapshot_query_after_remove() {
        // Step 1: Prepare one shard's test data and persist it to backend storage.
        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().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let config = WorkerNodeConfig::new_for_test(local_data_path);
        let shard_id = shard_meta.shard_header.shard_id;

        let backend_storage = test_create_backend_storage().await.unwrap();
        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

        // Step 2: Build a reader-side shard mirror and register it into the manager.
        let shard_mirror_manager = ShardMirrorManager::new(FragReaderConfig::default());
        let frag_data_cache_path = TempDir::new().unwrap().path().to_str().unwrap().to_string();
        let frag_data_cache = Arc::new(
            FragDataCache::new(
                backend_storage.clone(),
                frag_data_cache_path,
                config.storage.temporary_path(),
                Duration::from_secs(0),
                ReadableSize::mb(64),
            )
            .unwrap(),
        );
        let mut mock_client_for_read = meta_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, seq) = mock_client_for_read
            .get_frags_by_keyrange(shard_id, &shard_meta.index_key, KeyRange::default())
            .await
            .unwrap();
        let mut frag_reader_metas = Vec::new();
        for frag_meta in &frag_metas {
            frag_reader_metas.push(
                FragReaderMeta::from_frag_data_cache(
                    frag_meta,
                    frag_data_cache.clone(),
                    TEST_S3_BUCKET,
                )
                .await
                .unwrap(),
            );
        }

        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 dedup_output = ShardMirror::dedup_frags(
            shard_mirror_meta.shard_tracing_id(),
            shard_mirror_meta.shard_range().unwrap(),
            frag_reader_metas.clone(),
            config.storage.temporary_path(),
        )
        .await
        .unwrap()
        .unwrap();
        shard_mirror_manager
            .add_shard(
                shard_mirror_meta,
                index_meta.clone(),
                frag_reader_metas,
                dedup_output.obsolete_doc_bitsets,
            )
            .await
            .unwrap();

        // Step 3: Acquire a snapshot and verify this shard is visible in the snapshot.
        let (snapshot, hits) = shard_mirror_manager.check_shards_and_acquire_snapshot(&[shard_id]);
        assert_eq!(hits, vec![true]);
        assert!(snapshot.get_shard(shard_id).is_some());

        // Step 4: Remove shard from the latest manager state and verify it disappears
        // there.
        shard_mirror_manager.remove_shards(&[shard_id]);
        assert!(shard_mirror_manager.get_shard(shard_id).is_none());

        // Step 5: Query through the previously acquired snapshot and ensure it still
        // works.
        let snapshot_shard = snapshot.get_shard(shard_id).unwrap();
        let shard_reader = snapshot_shard.get_reader();
        let query_info = QueryInfo::query_all_for_test(
            &index_meta,
            &shard_reader.get_schema().unwrap(),
            vec![test_get_inf_range()],
        )
        .unwrap();
        let (all_docs, _search_stat) = shard_reader.search_with_query_info(query_info).unwrap();
        assert_eq!(all_docs.len(), 150000);

        // Step 6: Re-acquire latest snapshot and confirm the removed shard is no longer
        // present.
        let (latest_snapshot, latest_hits) =
            shard_mirror_manager.check_shards_and_acquire_snapshot(&[shard_id]);
        assert_eq!(latest_hits, vec![false]);
        assert!(latest_snapshot.get_shard(shard_id).is_none());
    }
}