use {
    crate::blockstore_meta::SlotMeta,
    bitflags::bitflags,
    lru::LruCache,
    solana_clock::Slot,
    std::{collections::HashMap, sync::Mutex},
};

const SLOTS_STATS_CACHE_CAPACITY: usize = 300;

#[derive(Copy, Clone, Debug)]
pub(crate) enum ShredSource {
    Turbine,
    Repaired,
    Recovered,
}

bitflags! {
    #[derive(Copy, Clone, Default)]
    struct SlotFlags: u8 {
        const DEAD   = 0b00000001;
        const FULL   = 0b00000010;
        const ROOTED = 0b00000100;
    }
}

#[derive(Clone, Default)]
pub struct SlotStats {
    turbine_fec_set_index_counts: HashMap</*fec_set_index*/ u32, /*count*/ usize>,
    num_repaired: usize,
    num_recovered: usize,
    last_index: u64,
    flags: SlotFlags,
}

impl SlotStats {
    pub fn get_min_index_count(&self) -> usize {
        self.turbine_fec_set_index_counts
            .values()
            .min()
            .copied()
            .unwrap_or_default()
    }

    fn report(&self, slot: Slot) {
        let min_fec_set_count = self.get_min_index_count();
        datapoint_info!(
            "slot_stats_tracking_complete",
            ("slot", slot, i64),
            ("last_index", self.last_index, i64),
            ("num_repaired", self.num_repaired, i64),
            ("num_recovered", self.num_recovered, i64),
            ("min_turbine_fec_set_count", min_fec_set_count, i64),
            ("is_full", self.flags.contains(SlotFlags::FULL), bool),
            ("is_rooted", self.flags.contains(SlotFlags::ROOTED), bool),
            ("is_dead", self.flags.contains(SlotFlags::DEAD), bool),
        );
    }
}

pub struct SlotsStats {
    pub stats: Mutex<LruCache<Slot, SlotStats>>,
}

impl Default for SlotsStats {
    fn default() -> Self {
        Self {
            stats: Mutex::new(LruCache::new(SLOTS_STATS_CACHE_CAPACITY)),
        }
    }
}

impl SlotsStats {
    /// Returns a mutable reference to [`SlotStats`] associated with the slot in the stats LruCache
    /// and a possibly evicted cache entry.
    ///
    /// A new SlotStats entry will be inserted if there is not one present for `slot`; insertion
    /// may cause an existing entry to be evicted.
    fn get_or_default_with_eviction_check(
        stats: &mut LruCache<Slot, SlotStats>,
        slot: Slot,
    ) -> (&mut SlotStats, Option<(Slot, SlotStats)>) {
        let evicted = if stats.contains(&slot) {
            None
        } else {
            // insert slot in cache which might potentially evict an entry
            let evicted = stats.push(slot, SlotStats::default());
            if let Some((evicted_slot, _)) = evicted {
                assert_ne!(evicted_slot, slot);
            }
            evicted
        };
        (stats.get_mut(&slot).unwrap(), evicted)
    }

    pub(crate) fn record_shred(
        &self,
        slot: Slot,
        fec_set_index: u32,
        source: ShredSource,
        slot_meta: Option<&SlotMeta>,
    ) {
        let (slot_full_reporting_info, evicted) = {
            let mut stats = self.stats.lock().unwrap();
            let (slot_stats, evicted) = Self::get_or_default_with_eviction_check(&mut stats, slot);
            match source {
                ShredSource::Recovered => slot_stats.num_recovered += 1,
                ShredSource::Repaired => slot_stats.num_repaired += 1,
                ShredSource::Turbine => {
                    *slot_stats
                        .turbine_fec_set_index_counts
                        .entry(fec_set_index)
                        .or_default() += 1
                }
            }
            let mut slot_full_reporting_info = None;
            if let Some(meta) = slot_meta {
                if meta.is_full() {
                    slot_stats.last_index = meta.last_index.unwrap();
                    if !slot_stats.flags.contains(SlotFlags::FULL) {
                        slot_stats.flags |= SlotFlags::FULL;
                        slot_full_reporting_info =
                            Some((slot_stats.num_repaired, slot_stats.num_recovered));
                    }
                }
            }
            (slot_full_reporting_info, evicted)
        };
        if let Some((num_repaired, num_recovered)) = slot_full_reporting_info {
            let slot_meta = slot_meta.unwrap();
            let total_time_ms =
                solana_time_utils::timestamp().saturating_sub(slot_meta.first_shred_timestamp);
            let last_index = slot_meta
                .last_index
                .and_then(|ix| i64::try_from(ix).ok())
                .unwrap_or(-1);
            datapoint_info!(
                "shred_insert_is_full",
                ("slot", slot, i64),
                ("total_time_ms", total_time_ms, i64),
                ("last_index", last_index, i64),
                ("num_repaired", num_repaired, i64),
                ("num_recovered", num_recovered, i64),
            );
        }
        if let Some((evicted_slot, evicted_stats)) = evicted {
            evicted_stats.report(evicted_slot);
        }
    }

    fn add_flag(&self, slot: Slot, flag: SlotFlags) {
        let evicted = {
            let mut stats = self.stats.lock().unwrap();
            let (slot_stats, evicted) = Self::get_or_default_with_eviction_check(&mut stats, slot);
            slot_stats.flags |= flag;
            evicted
        };
        if let Some((evicted_slot, evicted_stats)) = evicted {
            evicted_stats.report(evicted_slot);
        }
    }

    pub fn mark_dead(&self, slot: Slot) {
        self.add_flag(slot, SlotFlags::DEAD);
    }

    pub fn mark_rooted(&self, slot: Slot) {
        self.add_flag(slot, SlotFlags::ROOTED);
    }
}
