use std::{
    collections::HashMap,
    fmt::Display,
    sync::{
        Arc, LazyLock, Mutex,
        atomic::{AtomicUsize, Ordering},
    },
};

pub static COUNTERS: LazyLock<Mutex<HashMap<&'static str, Arc<AtomicUsize>>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

#[macro_export]
macro_rules! counter {
    ($key:expr) => {
        $crate::common::stats::CounterRef::from(
            &*$crate::common::stats::COUNTERS
                .lock()
                .unwrap()
                .entry($key)
                .or_default(),
        )
    };
}

#[derive(Debug, Clone)]
pub struct CounterRef(Arc<AtomicUsize>);

#[allow(unused)]
impl CounterRef {
    pub fn inc(&self, by: usize) {
        self.0.fetch_add(by, Ordering::Relaxed);
    }
    pub fn set(&self, n: usize) {
        self.0.store(n, Ordering::Relaxed);
    }
    pub fn get(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
}

impl From<&Arc<AtomicUsize>> for CounterRef {
    fn from(value: &Arc<AtomicUsize>) -> Self {
        Self(value.clone())
    }
}

impl Display for CounterRef {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.get())
    }
}

pub fn get_stats() -> Vec<(&'static str, usize)> {
    let mut entries: Vec<_> = COUNTERS
        .lock()
        .unwrap()
        .iter()
        .map(|(k, v)| (*k, CounterRef::from(v).get()))
        .collect();
    entries.sort_unstable_by_key(|(k, _)| *k);
    entries
}

#[cfg(test)]
mod tests {

    #[test]
    fn test_counter() {
        counter!("test_counter").inc(4);
        assert_eq!(counter!("test_counter").get(), 4);
    }
}
