use std::{
    collections::BTreeMap,
    sync::{Arc, Mutex, OnceLock},
};

use crate::{DURATION_BUCKETS_SECONDS, max_exemplar_slots};

#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct ExemplarTraceId([u8; 32]);

impl ExemplarTraceId {
    pub fn parse(value: &str) -> Option<Self> {
        if value.len() != 32
            || value
                .bytes()
                .any(|byte| !byte.is_ascii_digit() && !(b'a'..=b'f').contains(&byte))
            || value.bytes().all(|byte| byte == b'0')
        {
            return None;
        }
        let mut bytes = [0; 32];
        bytes.copy_from_slice(value.as_bytes());
        Some(Self(bytes))
    }

    pub fn as_str(&self) -> &str {
        std::str::from_utf8(&self.0).expect("validated ASCII trace id")
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct ExemplarObservation {
    pub metric: &'static str,
    pub labels: Vec<(&'static str, &'static str)>,
    pub bucket_upper_bound: Option<f64>,
    pub value: f64,
    pub trace_id: ExemplarTraceId,
}

#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
struct ExemplarKey {
    metric: &'static str,
    labels: Vec<(&'static str, &'static str)>,
    bucket_index: usize,
}

fn store() -> &'static Mutex<BTreeMap<ExemplarKey, Arc<ExemplarObservation>>> {
    static STORE: OnceLock<Mutex<BTreeMap<ExemplarKey, Arc<ExemplarObservation>>>> =
        OnceLock::new();
    STORE.get_or_init(|| Mutex::new(BTreeMap::new()))
}

pub(crate) fn record_exemplar(
    metric: &'static str,
    labels: Vec<(&'static str, &'static str)>,
    value: f64,
    trace_id: Option<ExemplarTraceId>,
) {
    let Some(trace_id) = trace_id else { return };
    let bucket_index = DURATION_BUCKETS_SECONDS
        .iter()
        .position(|bound| value <= *bound)
        .unwrap_or(DURATION_BUCKETS_SECONDS.len());
    let observation = ExemplarObservation {
        metric,
        labels: labels.clone(),
        bucket_upper_bound: DURATION_BUCKETS_SECONDS.get(bucket_index).copied(),
        value,
        trace_id,
    };
    let Ok(mut observations) = store().lock() else {
        return;
    };
    let key = ExemplarKey {
        metric,
        labels,
        bucket_index,
    };
    if observations.contains_key(&key) || observations.len() < max_exemplar_slots() {
        observations.insert(key, Arc::new(observation));
    }
}

pub fn exemplar_snapshot() -> Vec<Arc<ExemplarObservation>> {
    store()
        .lock()
        .map(|observations| observations.values().cloned().collect())
        .unwrap_or_default()
}