use bevy_ecs::prelude::*;
use std::{
    collections::HashMap,
    sync::{atomic, Mutex, RwLock},
};

use crate::{entity_registry::EntityRegistry, Stages};

#[derive(Default)]
pub struct Histogram<T>
where
    T: Copy + Send + Sync,
{
    inner: Mutex<Vec<T>>,
}

impl<T> Histogram<T>
where
    T: Copy + Send + Sync,
{
    #[inline(always)]
    pub fn record(&self, val: T) {
        self.inner.lock().unwrap().push(val);
    }
}

#[derive(Default)]
pub struct Counter {
    inner: atomic::AtomicU64,
}

impl Counter {
    #[inline(always)]
    pub fn increment(&self, value: u64) {
        self.inner.fetch_add(value, atomic::Ordering::Relaxed);
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MetricKey {
    entity: Option<Entity>,
    label: &'static str,
}

#[derive(Resource, Default)]
pub struct Metrics {
    pub(crate) active: bool,
    counters: RwLock<HashMap<MetricKey, Counter>>,
    histograms: RwLock<HashMap<MetricKey, Histogram<u32>>>,
}

impl Metrics {
    pub fn record_histogram(&self, entity: Option<Entity>, label: &'static str, value: u32) {
        if self.active {
            let key = MetricKey { entity, label };
            let r_hist = self.histograms.read().unwrap();
            if let Some(hist) = r_hist.get(&key) {
                hist.record(value);
            } else {
                std::mem::drop(r_hist);

                let mut histograms = self.histograms.write().unwrap();
                histograms.insert(
                    key,
                    Histogram {
                        inner: vec![value].into(),
                    },
                );
            }
        }
    }

    pub fn increment_counter(&self, entity: Option<Entity>, label: &'static str, value: u64) {
        if self.active {
            let key = MetricKey { entity, label };
            let r_counters = self.counters.read().unwrap();
            if let Some(counter) = r_counters.get(&key) {
                counter.increment(value);
            } else {
                std::mem::drop(r_counters);

                let mut counters = self.counters.write().unwrap();
                counters.insert(
                    key,
                    Counter {
                        inner: value.into(),
                    },
                );
            }
        }
    }
}

pub(crate) fn consume_metrics(
    world: &World,
) -> (
    Vec<proxisim_models::dto::metrics::Counter>,
    Vec<proxisim_models::dto::metrics::Histogram>,
) {
    let metrics = world.resource::<Metrics>();
    let entities = world.resource::<EntityRegistry>();

    let counters = metrics
        .counters
        .try_write()
        .unwrap()
        .drain()
        .map(|(key, value)| proxisim_models::dto::metrics::Counter {
            entity: key
                .entity
                .as_ref()
                .map(|e| entities.0.get(e).unwrap().clone())
                .unwrap_or(proxisim_models::dto::metrics::EntityInfo::Global),
            value: value.inner.load(atomic::Ordering::Relaxed),
            label: key.label,
        })
        .collect();

    let histograms = metrics
        .histograms
        .try_write()
        .unwrap()
        .drain()
        .map(|(key, value)| proxisim_models::dto::metrics::Histogram {
            entity: key
                .entity
                .as_ref()
                .map(|e| entities.0.get(e).unwrap().clone())
                .unwrap_or(proxisim_models::dto::metrics::EntityInfo::Global),
            values: value.inner.into_inner().unwrap(),
            label: key.label,
        })
        .collect();

    (counters, histograms)
}

pub(crate) fn configure(stages: &mut Stages) {
    stages.world.init_resource::<Metrics>();
}

#[cfg(test)]
mod test {}