use std::{
    collections::{BTreeMap, HashSet},
    sync::Arc,
};

use anyhow::Result;
use either::Either;
use log::{debug, error, info};
use tici_common::{
    codec::{table_end_key, table_start_key},
    index::IndexKey,
};
use tici_grpc::{IndexInfo, TableInfo};

use crate::{
    HandleRange, RangeBound, ShardInfo, ShardInfoAppendUpdate,
    meta::{
        metrics::SHARD_CACHE_EVENT_COUNTER,
        reader::{
            AllocateReaderReason, CacheAllocator, ManualRemoveResult, ReaderPool, UnifiedShardInfo,
            allocation_map::{QueryResult, ShardAllocMap},
        },
    },
};

/// Index metadata stored in cache manager (unified with cache allocations)
#[derive(Clone)]
pub struct IndexMeta {
    pub table_info: TableInfo,
    pub index_info: IndexInfo,
    pub index_key: IndexKey,

    pub shard_map: ShardAllocMap,
    pub cache_allocator: CacheAllocator,
}

impl std::fmt::Debug for IndexMeta {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("IndexMeta")
            .field("table_info", &self.table_info)
            .field("index_info", &self.index_info)
            .field("shards_len", &self.shard_map.len())
            .field("index_key", &self.index_key)
            .finish()
    }
}

impl IndexMeta {
    pub fn new(
        table_info: TableInfo,
        index_info: IndexInfo,
        cache_allocator: CacheAllocator,
        optimize_for_fast_scan: bool,
    ) -> Self {
        let index_key = IndexKey::from(&index_info);
        let shard_map =
            ShardAllocMap::new_uninitialized(format!("{:?}", index_key), optimize_for_fast_scan);
        Self {
            table_info,
            index_info,
            index_key,
            shard_map,
            cache_allocator,
        }
    }

    pub fn get_shard_by_id(&self, shard_id: &u64) -> Option<UnifiedShardInfo> {
        self.shard_map.shards_by_id.get(shard_id).cloned()
    }

    pub fn shard_len(&self) -> usize {
        self.shard_map.len()
    }

    pub fn init(&mut self, v: Vec<UnifiedShardInfo>) {
        self.shard_map.init(v);
    }

    pub async fn find_overlapping_shards(
        &self,
        key_ranges: &[tici_grpc::KeyRange],
        reason: AllocateReaderReason,
        perform_warmup: Option<Arc<ReaderPool>>,
        auto_merge_range: bool,
    ) -> Result<(Vec<UnifiedShardInfo>, usize, bool)> {
        fn push_or_merge_adjacent_same_shard(
            overlapping_shards: &mut Vec<UnifiedShardInfo>,
            shard: UnifiedShardInfo,
            auto_merge_range: bool,
        ) {
            if auto_merge_range {
                if let Some(last) = overlapping_shards.last_mut() {
                    if last.shard_id == shard.shard_id
                        && last.end_key_ref() == shard.start_key_ref()
                    {
                        last.range =
                            HandleRange::new_from_sorted(last.start_key(), shard.end_key());
                        return;
                    }
                }
            }
            overlapping_shards.push(shard);
        }

        let mut overlapping_shards = Vec::new();
        if key_ranges.is_empty() {
            anyhow::bail!("must specify ranges, index_key={:?}", self.index_key);
        }
        let table_bound = HandleRange::new(
            table_start_key(&self.index_key),
            table_end_key(&self.index_key),
        );
        let mut not_ready_count: usize = 0;
        let mut any_is_trimed = false;
        let mut is_no_reader = false;
        for range in key_ranges {
            let handle_range = HandleRange::new(range.start_key.clone(), range.end_key.clone());
            let handle_range = handle_range.intersect(&table_bound);
            let (res, to_remedy_warmup) = self.shard_map.query(&handle_range).await?;
            for s in to_remedy_warmup {
                if let Some(reader_pool) = perform_warmup.clone() {
                    if reader_pool.get_readers_len().await != 0 {
                        debug!(
                            "start warmup on read remedy, shard={}, reason={:?}",
                            s.to_long_ident(),
                            &reason
                        );
                        SHARD_CACHE_EVENT_COUNTER
                            .with_label_values(&["warm_remedy_suc"])
                            .inc();
                        reader_pool.background_warmup(
                            &s,
                            reader_pool.clone(),
                            self.table_info.clone(),
                            self.index_info.clone(),
                            reason.clone(),
                        )?;
                    } else {
                        debug!(
                            "skip start warmup on read remedy, shard={}, reason={:?}",
                            s.to_long_ident(),
                            &reason
                        );
                        // If no reader, we don't warmup.
                        SHARD_CACHE_EVENT_COUNTER
                            .with_label_values(&["warm_remedy_failed_no_reader"])
                            .inc();
                        is_no_reader = true;
                    }
                }
            }
            for tuple in res.into_iter() {
                let QueryResult {
                    shard,
                    valid_range,
                    is_trimed,
                } = tuple;
                any_is_trimed |= is_trimed;
                match shard {
                    Either::Left(s) => {
                        let mut s = s;
                        s.range = valid_range;
                        push_or_merge_adjacent_same_shard(
                            &mut overlapping_shards,
                            s,
                            auto_merge_range,
                        );
                    }
                    Either::Right(s) => {
                        // Not warmed up
                        if let Some(reader_pool) = perform_warmup.clone() {
                            // We don't blocked wait here for this exact shard, because a newer
                            // shard maybe warmed.
                            if reader_pool.get_readers_len().await != 0 {
                                debug!(
                                    "start warmup on read request, shard={}, reason={:?}",
                                    s.to_long_ident(),
                                    &reason
                                );
                                reader_pool.background_warmup(
                                    &s,
                                    reader_pool.clone(),
                                    self.table_info.clone(),
                                    self.index_info.clone(),
                                    reason.clone(),
                                )?;
                            } else {
                                debug!(
                                    "skip start warmup on read request, shard={}, reason={:?}",
                                    s.to_long_ident(),
                                    &reason
                                );
                                // If no reader, we don't warmup.
                                SHARD_CACHE_EVENT_COUNTER
                                    .with_label_values(&["warm_failed_no_reader"])
                                    .inc();
                                is_no_reader = true;
                            }
                        }
                        let mut s = s;
                        s.range = valid_range;
                        push_or_merge_adjacent_same_shard(
                            &mut overlapping_shards,
                            s,
                            auto_merge_range,
                        );
                        not_ready_count += 1;
                    }
                }
            }
        }

        if not_ready_count > 0 && any_is_trimed {
            SHARD_CACHE_EVENT_COUNTER
                .with_label_values(&["read_partial_shard"])
                .inc();
        }
        Ok((overlapping_shards, not_ready_count, is_no_reader))
    }

    pub async fn on_warmup_finished(
        &mut self,
        unified_shard: &UnifiedShardInfo,
        reason: AllocateReaderReason,
    ) -> Result<HashSet<u64>> {
        fail::fail_point!("on_warmup_finished_start", |_| Ok(HashSet::new()));
        let r = if self.shard_map.scan_optimized() {
            self.shard_map
                .do_whole_replace(&unified_shard.shard_id)
                .await?
        } else {
            self.shard_map.do_replace(&unified_shard.shard_id)?
        };
        if !r.is_empty() {
            let coalesced = if reason == AllocateReaderReason::Merge {
                self.shard_map.coalesce(&unified_shard.range).await
            } else {
                vec![]
            };
            info!(
                "on_warmup_finished, replaced shards, reason={:?}, by={}, removed={}, coalesced={:?}",
                reason,
                unified_shard.to_long_ident(),
                r.iter()
                    .map(|x| format!("shard_id={}", x))
                    .collect::<Vec<_>>()
                    .join(":"),
                coalesced,
            );
        }
        Ok(r)
    }

    /// Get cache allocations for overlapping shards
    pub async fn get_cache_allocations_for_shards(
        &self,
        overlapping_shards: &[&UnifiedShardInfo],
    ) -> BTreeMap<RangeBound, UnifiedShardInfo> {
        let mut result = BTreeMap::new();
        for unified_shard in overlapping_shards.iter() {
            if unified_shard.has_ready_readers().await {
                result.insert(unified_shard.start_key(), (*unified_shard).clone());
            }
        }
        result
    }

    /// Remove cache allocations for expired reader
    pub async fn on_reader_node_down(&mut self, expired_reader_addr: &str) -> usize {
        let mut removed_count = 0;

        // TODO(P2) Use the shard info in read node meta to avoid iterating all over the
        // `shards_by_id`.
        // Although we failover all read nodes, however, if the reader node recovered,
        // it will report with its remained nodes, and as a result, these nodes will be
        // remedied again.
        for unified_shard in self.shard_map.shards_by_id.values_mut() {
            info!(
                "remove_reader for node down, addr={}, unified_shard_info={}",
                expired_reader_addr,
                unified_shard.to_long_ident()
            );
            let callbacks = unified_shard
                .take_remove_callbacks(expired_reader_addr)
                .await;
            if unified_shard
                .finalize_remove_reader(expired_reader_addr)
                .await
            {
                for callback in callbacks {
                    let _ = callback.send(ManualRemoveResult::ReaderExpired);
                }
                removed_count += 1;
            }
        }

        removed_count
    }

    pub async fn handle_split_shard(
        &mut self,
        original: ShardInfo,
        new_shards: &[ShardInfo],
        reader_pool: Arc<ReaderPool>,
    ) {
        let warm_derived_when_parent_warming = reader_pool
            .warmup_mgr()
            .should_warm_derived_when_parent_warming()
            .await;
        let original_has_allocation = if let Some(x) = self.get_shard_by_id(&original.shard_id) {
            x.has_ready_readers().await || (warm_derived_when_parent_warming && x.is_warming())
        } else {
            false
        };

        // Add new shards with their cache allocations (if any)
        let mut uni_shards = vec![];
        for new_shard in new_shards {
            let unishard = UnifiedShardInfo::from_shard_info(new_shard);
            // The new shard must be warmed up, if the original is warmed.
            if original_has_allocation {
                debug!("warmup for splitted, shard={}", unishard.to_long_ident());
                if let Err(e) = reader_pool.background_warmup(
                    &unishard,
                    reader_pool.clone(),
                    self.table_info.clone(),
                    self.index_info.clone(),
                    AllocateReaderReason::Split,
                ) {
                    error!(
                        "handle_split_shard warmup error, origin={:?}, new={:?}, e={:?}",
                        original.to_long_ident(),
                        new_shards,
                        e
                    )
                }
            }
            uni_shards.push(unishard);
        }

        if let Err(e) = self.shard_map.split(
            original.range.clone(),
            original.shard_id,
            uni_shards.clone(),
        ) {
            error!(
                "handle_split_shard error, origin={:?}, new={:?}, e={:?}",
                original.to_long_ident(),
                new_shards,
                e
            )
        }

        info!(
            "handle_split_shard, shard_id={} ===> [{}]{}, after_node_len={}, after_len={}, original_has_allocation={}",
            original.shard_id,
            new_shards.len(),
            new_shards
                .iter()
                .map(|x| format!("shard_id={}", x.shard_id))
                .collect::<Vec<_>>()
                .join(", "),
            self.shard_map.segments_len(),
            self.shard_map.shards_by_id.len(),
            original_has_allocation,
        );
    }

    pub async fn handle_merge_shard(
        &mut self,
        old_shards: Vec<ShardInfo>,
        new_shard: ShardInfo,
        reader_pool: Arc<ReaderPool>,
    ) {
        let mut old_shards_allocation_count = 0;
        for shard in old_shards.iter() {
            if let Some(x) = self.get_shard_by_id(&shard.shard_id) {
                if x.has_ready_readers().await || x.is_warming() {
                    old_shards_allocation_count += 1;
                }
            }
        }

        let any_old_shard_has_allocation = old_shards_allocation_count > 0;

        let unishard = UnifiedShardInfo::from_shard_info(&new_shard);
        // The new shard must be warmed up, if the original is warmed.
        if any_old_shard_has_allocation {
            debug!("warmup for merge, shard={}", unishard.to_long_ident());
            if let Err(e) = reader_pool.background_warmup(
                &unishard,
                reader_pool.clone(),
                self.table_info.clone(),
                self.index_info.clone(),
                AllocateReaderReason::Merge,
            ) {
                error!(
                    "handle_merge_shard warmup error, origin={:?}, new={:?}, e={:?}",
                    old_shards, unishard, e
                )
            }
        }

        if let Err(e) = self.shard_map.merge(unishard.clone()) {
            error!(
                "handle_merge_shard error, origin={:?}, new={:?}, e={:?}",
                old_shards, unishard, e
            )
        }

        info!(
            "handle_merge_shard, [{}]{} ===> shard_id={}, after_node_len={}, after_len={}, original_has_allocation={}",
            old_shards.len(),
            old_shards
                .iter()
                .map(|x| format!("shard_id={}", x.shard_id))
                .collect::<Vec<_>>()
                .join(", "),
            new_shard.shard_id,
            self.shard_map.segments_len(),
            self.shard_map.shards_by_id.len(),
            any_old_shard_has_allocation,
        );
    }

    pub fn handle_add_shard(&mut self, unified_shard: UnifiedShardInfo) -> Result<()> {
        self.shard_map.add_shard(unified_shard)
    }

    pub fn handle_remove_shard(&mut self, shard_ids: &[u64]) -> Result<()> {
        let _ = self.shard_map.remove_shards(shard_ids)?;
        Ok(())
    }

    pub async fn handle_update_shard(&self, new_shard: &ShardInfo) -> Result<()> {
        if let Some(existing_shard) = self.get_shard_by_id(&new_shard.shard_id) {
            existing_shard.apply_update(new_shard).await;
            Ok(())
        } else {
            Err(anyhow::anyhow!(
                "handle_update_shard: shard not found, shard_id={}, index_key={:?}",
                new_shard.shard_id,
                new_shard.index_key
            ))
        }
    }

    pub async fn handle_append_shard(&self, append: &ShardInfoAppendUpdate) -> Result<()> {
        if let Some(existing_shard) = self.shard_map.shards_by_id.get(&append.shard_id) {
            existing_shard.apply_append(append).await;
            Ok(())
        } else {
            Err(anyhow::anyhow!(
                "handle_append_shard: shard not found, shard_id={}, index_key={:?}",
                append.shard_id,
                self.index_key
            ))
        }
    }
}