use std::{
    collections::hash_map::Entry::{Occupied as HOccupied, Vacant as HVacant},
    sync::Arc,
};

use anyhow::Result;
use log::info;
use tici_common::index::IndexKey;
use tici_grpc::{IndexInfo, TableInfo};
use tokio::sync::RwLock;

use crate::{
    RangeBound, ShardInfo,
    meta::{
        metrics::*,
        writer::{FTSIndex, ShardTaskManager},
    },
};

impl ShardTaskManager {
    pub async fn bootstrap_index(
        &self,
        table_info: TableInfo,
        index_info: IndexInfo,
    ) -> Result<bool> {
        // TODO(P1): Create CDC changefeed here.
        let index_key = IndexKey::from(&index_info);
        info!("bootstrap_index, index_key={:?}", index_key);
        let prefix = tici_common::s3_keys::cdc_path_prefix(
            &self.config.load().s3.prefix,
            table_info.table_id,
        );
        let mut guard = self.indices.write().await;
        match guard.entry(IndexKey::from(&index_info)) {
            HOccupied(_) => Ok(false),
            HVacant(v) => {
                self.meta_store
                    .init_index(
                        &table_info,
                        &index_info,
                        self.config.load().s3.bucket.clone(),
                        prefix.clone(),
                    )
                    .await?;
                v.insert(Arc::new(FTSIndex::new(
                    self.meta_store.clone(),
                    table_info,
                    index_info,
                    self.config.load().s3.bucket.clone(),
                    prefix.clone(),
                )));

                INDEX_COUNT_GAUGE_VEC
                    .with_label_values(&["total"])
                    .set(guard.len() as i64);
                Ok(true)
            }
        }
    }

    pub async fn get_index(&self, index_key: &IndexKey) -> Result<Arc<FTSIndex>, anyhow::Error> {
        let guard = self.indices.read().await;
        if let Some(o) = guard.get(index_key) {
            Ok(o.clone())
        } else {
            Err(anyhow::format_err!(
                "The index is not bootstrapped {:?}, total index={}",
                index_key,
                guard.len(),
            ))
        }
    }

    pub async fn finish_append_base(&self, index_key: &IndexKey) -> Result<bool> {
        self.get_index(index_key).await?.finish_append_base().await
    }

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

    pub async fn get_shard(
        &self,
        index_key: &IndexKey,
        start_key: &RangeBound,
    ) -> Result<Arc<RwLock<ShardInfo>>, anyhow::Error> {
        Ok(self.get_index_and_shard(index_key, start_key).await?.1)
    }

    pub async fn find_shard_by_id(
        &self,
        shard_id: u64,
    ) -> Option<(Arc<FTSIndex>, Arc<RwLock<ShardInfo>>)> {
        let all_indexes: Vec<Arc<FTSIndex>> = self.indices.read().await.values().cloned().collect();

        for index in all_indexes {
            for shard in index.clone_shards().await {
                if shard.read().await.shard_id == shard_id {
                    return Some((index.clone(), shard));
                }
            }
        }

        None
    }

    pub async fn get_index_and_shard(
        &self,
        index_key: &IndexKey,
        start_key: &RangeBound,
    ) -> Result<(Arc<FTSIndex>, Arc<RwLock<ShardInfo>>), anyhow::Error> {
        let index = self.get_index(index_key).await?;

        if let Some(x) = index.get_shard(start_key).await {
            return Ok((index, x.clone()));
        }
        Err(anyhow::format_err!(
            "The shard is not bootstrapped {:?}, start_key={:?}",
            index_key,
            start_key,
        ))
    }

    pub async fn remove_index(&self, index_key: &IndexKey) -> Result<()> {
        if let Ok(index) = self.get_index(index_key).await {
            let v = index.remove_all_shards().await?;
            let shard_len = v.len();
            self.helper_after_remove_shard(v, index.clone()).await?;
            info!(
                "remove_index, index_key={:?}, shard_count={}",
                index_key, shard_len
            );
        } else {
            info!("drop_index: can't find index {:?}", index_key);
        }
        // Drop index
        self.meta_store.drop_index(index_key).await?;
        self.indices.write().await.remove(index_key);
        Ok(())
    }

    pub async fn collect_partition_index_keys(
        &self,
        keyspace_id: u32,
        logical_table_id: i64,
        index_id: i64,
    ) -> Vec<IndexKey> {
        let guard = self.indices.read().await;
        guard
            .iter()
            .filter_map(|(index_key, index)| {
                if index_key.keyspace_id != keyspace_id || index_key.index_id != index_id {
                    return None;
                }
                if index.table_info.logical_table_id != logical_table_id {
                    return None;
                }
                Some(index_key.clone())
            })
            .collect()
    }

    pub fn on_index_change(&self) {
        SHARD_COUNT_GAUGE_VEC.with_label_values(&["total"]).set(
            self.stats
                .total_shards_count
                .load(std::sync::atomic::Ordering::Relaxed) as i64,
        );
    }
}