lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! Metrics registry and collection

use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec::Vec;

use super::types::*;

/// Metrics registry
pub struct MetricsRegistry {
    /// Next metric ID
    next_id: u64,
    /// Registered metrics
    metrics: BTreeMap<MetricId, MetricDefinition>,
    /// Name to ID mapping
    name_to_id: BTreeMap<String, MetricId>,
    /// Metric values
    values: BTreeMap<MetricId, MetricValues>,
}

/// Metric definition
#[derive(Debug, Clone)]
struct MetricDefinition {
    /// Metric ID
    id: MetricId,
    /// Metric name
    name: String,
    /// Help text
    help: String,
    /// Metric type
    metric_type: MetricType,
    /// Label names
    labels: Vec<String>,
    /// Histogram buckets (if histogram)
    buckets: Vec<f64>,
}

/// Metric values storage
#[derive(Debug)]
enum MetricValues {
    /// Counter values per label set
    Counter(BTreeMap<String, f64>),
    /// Gauge values per label set
    Gauge(BTreeMap<String, f64>),
    /// Histogram values per label set
    Histogram(BTreeMap<String, HistogramData>),
    /// Summary values per label set
    Summary(BTreeMap<String, SummaryData>),
}

impl MetricsRegistry {
    /// Create a new metrics registry
    pub fn new() -> Self {
        Self {
            next_id: 1,
            metrics: BTreeMap::new(),
            name_to_id: BTreeMap::new(),
            values: BTreeMap::new(),
        }
    }

    /// Register a counter metric
    pub fn register_counter(
        &mut self,
        name: &str,
        help: &str,
        labels: &[&str],
    ) -> Result<MetricId, TelemetryError> {
        validate_metric_name(name)?;

        if self.name_to_id.contains_key(name) {
            return Err(TelemetryError::MetricExists(name.to_string()));
        }

        let id = MetricId::new(self.next_id);
        self.next_id += 1;

        let def = MetricDefinition {
            id,
            name: name.to_string(),
            help: help.to_string(),
            metric_type: MetricType::Counter,
            labels: labels.iter().map(|s| s.to_string()).collect(),
            buckets: Vec::new(),
        };

        self.metrics.insert(id, def);
        self.name_to_id.insert(name.to_string(), id);
        self.values
            .insert(id, MetricValues::Counter(BTreeMap::new()));

        Ok(id)
    }

    /// Register a gauge metric
    pub fn register_gauge(
        &mut self,
        name: &str,
        help: &str,
        labels: &[&str],
    ) -> Result<MetricId, TelemetryError> {
        validate_metric_name(name)?;

        if self.name_to_id.contains_key(name) {
            return Err(TelemetryError::MetricExists(name.to_string()));
        }

        let id = MetricId::new(self.next_id);
        self.next_id += 1;

        let def = MetricDefinition {
            id,
            name: name.to_string(),
            help: help.to_string(),
            metric_type: MetricType::Gauge,
            labels: labels.iter().map(|s| s.to_string()).collect(),
            buckets: Vec::new(),
        };

        self.metrics.insert(id, def);
        self.name_to_id.insert(name.to_string(), id);
        self.values.insert(id, MetricValues::Gauge(BTreeMap::new()));

        Ok(id)
    }

    /// Register a histogram metric
    pub fn register_histogram(
        &mut self,
        name: &str,
        help: &str,
        labels: &[&str],
        buckets: &[f64],
    ) -> Result<MetricId, TelemetryError> {
        validate_metric_name(name)?;

        if self.name_to_id.contains_key(name) {
            return Err(TelemetryError::MetricExists(name.to_string()));
        }

        if buckets.is_empty() {
            return Err(TelemetryError::InvalidBuckets);
        }

        let id = MetricId::new(self.next_id);
        self.next_id += 1;

        let def = MetricDefinition {
            id,
            name: name.to_string(),
            help: help.to_string(),
            metric_type: MetricType::Histogram,
            labels: labels.iter().map(|s| s.to_string()).collect(),
            buckets: buckets.to_vec(),
        };

        self.metrics.insert(id, def.clone());
        self.name_to_id.insert(name.to_string(), id);
        self.values
            .insert(id, MetricValues::Histogram(BTreeMap::new()));

        Ok(id)
    }

    /// Increment a counter
    pub fn counter_inc(&mut self, id: MetricId, labels: &[&str], value: f64) {
        if let Some(MetricValues::Counter(values)) = self.values.get_mut(&id) {
            let key = make_label_key(labels);
            *values.entry(key).or_insert(0.0) += value;
        }
    }

    /// Set a gauge value
    pub fn gauge_set(&mut self, id: MetricId, labels: &[&str], value: f64) {
        if let Some(MetricValues::Gauge(values)) = self.values.get_mut(&id) {
            let key = make_label_key(labels);
            values.insert(key, value);
        }
    }

    /// Add to a gauge value
    pub fn gauge_add(&mut self, id: MetricId, labels: &[&str], value: f64) {
        if let Some(MetricValues::Gauge(values)) = self.values.get_mut(&id) {
            let key = make_label_key(labels);
            *values.entry(key).or_insert(0.0) += value;
        }
    }

    /// Observe a histogram value
    pub fn histogram_observe(&mut self, id: MetricId, labels: &[&str], value: f64) {
        let buckets = self.metrics.get(&id).map(|d| d.buckets.clone());

        if let (Some(MetricValues::Histogram(values)), Some(bucket_bounds)) =
            (self.values.get_mut(&id), buckets)
        {
            let key = make_label_key(labels);
            let hist = values
                .entry(key)
                .or_insert_with(|| HistogramData::new(&bucket_bounds));
            hist.observe(value);
        }
    }

    /// Get metric information
    pub fn get_metric(&self, id: MetricId) -> Option<MetricInfo> {
        self.metrics.get(&id).map(|def| MetricInfo {
            id: def.id,
            name: def.name.clone(),
            help: def.help.clone(),
            metric_type: def.metric_type,
            labels: def.labels.clone(),
        })
    }

    /// Get metric by name
    pub fn get_metric_by_name(&self, name: &str) -> Option<MetricInfo> {
        self.name_to_id
            .get(name)
            .and_then(|id| self.get_metric(*id))
    }

    /// List all metrics
    pub fn list_metrics(&self) -> Vec<MetricInfo> {
        self.metrics
            .values()
            .map(|def| MetricInfo {
                id: def.id,
                name: def.name.clone(),
                help: def.help.clone(),
                metric_type: def.metric_type,
                labels: def.labels.clone(),
            })
            .collect()
    }

    /// Get all metric definitions and values for export
    pub fn get_all_for_export(&self) -> Vec<ExportMetric> {
        let mut result = Vec::new();

        for (id, def) in &self.metrics {
            if let Some(values) = self.values.get(id) {
                result.push(ExportMetric {
                    name: def.name.clone(),
                    help: def.help.clone(),
                    metric_type: def.metric_type,
                    label_names: def.labels.clone(),
                    buckets: def.buckets.clone(),
                    values: values.clone_for_export(),
                });
            }
        }

        result
    }

    /// Reset all metrics
    pub fn reset(&mut self) {
        for values in self.values.values_mut() {
            match values {
                MetricValues::Counter(v) => v.clear(),
                MetricValues::Gauge(v) => v.clear(),
                MetricValues::Histogram(v) => {
                    for hist in v.values_mut() {
                        hist.reset();
                    }
                }
                MetricValues::Summary(v) => {
                    for sum in v.values_mut() {
                        sum.reset();
                    }
                }
            }
        }
    }
}

impl Default for MetricsRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl MetricValues {
    fn clone_for_export(&self) -> ExportValues {
        match self {
            MetricValues::Counter(v) => ExportValues::Counter(v.clone()),
            MetricValues::Gauge(v) => ExportValues::Gauge(v.clone()),
            MetricValues::Histogram(v) => ExportValues::Histogram(v.clone()),
            MetricValues::Summary(v) => ExportValues::Summary(v.clone()),
        }
    }
}

/// Export metric data
#[derive(Debug, Clone)]
pub struct ExportMetric {
    /// Metric name
    pub name: String,
    /// Help text
    pub help: String,
    /// Metric type
    pub metric_type: MetricType,
    /// Label names
    pub label_names: Vec<String>,
    /// Histogram buckets
    pub buckets: Vec<f64>,
    /// Values
    pub values: ExportValues,
}

/// Export values
#[derive(Debug, Clone)]
pub enum ExportValues {
    /// Counter values
    Counter(BTreeMap<String, f64>),
    /// Gauge values
    Gauge(BTreeMap<String, f64>),
    /// Histogram values
    Histogram(BTreeMap<String, HistogramData>),
    /// Summary values
    Summary(BTreeMap<String, SummaryData>),
}

/// Create a label key from label values
fn make_label_key(labels: &[&str]) -> String {
    labels.join("\x00")
}

/// Parse a label key back to values
pub fn parse_label_key(key: &str) -> Vec<&str> {
    key.split('\x00').collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_registry_counter() {
        let mut registry = MetricsRegistry::new();

        let id = registry
            .register_counter("test_counter", "A test counter", &["label1"])
            .unwrap();

        registry.counter_inc(id, &["value1"], 1.0);
        registry.counter_inc(id, &["value1"], 2.0);
        registry.counter_inc(id, &["value2"], 5.0);

        let metrics = registry.get_all_for_export();
        assert_eq!(metrics.len(), 1);

        if let ExportValues::Counter(values) = &metrics[0].values {
            assert_eq!(*values.get("value1").unwrap(), 3.0);
            assert_eq!(*values.get("value2").unwrap(), 5.0);
        } else {
            panic!("Expected counter values");
        }
    }

    #[test]
    fn test_registry_gauge() {
        let mut registry = MetricsRegistry::new();

        let id = registry
            .register_gauge("test_gauge", "A test gauge", &["label1"])
            .unwrap();

        registry.gauge_set(id, &["value1"], 100.0);
        registry.gauge_add(id, &["value1"], 10.0);

        let metrics = registry.get_all_for_export();

        if let ExportValues::Gauge(values) = &metrics[0].values {
            assert_eq!(*values.get("value1").unwrap(), 110.0);
        }
    }

    #[test]
    fn test_registry_histogram() {
        let mut registry = MetricsRegistry::new();

        let id = registry
            .register_histogram("test_histogram", "A test histogram", &[], &[1.0, 5.0, 10.0])
            .unwrap();

        registry.histogram_observe(id, &[], 0.5);
        registry.histogram_observe(id, &[], 3.0);
        registry.histogram_observe(id, &[], 7.0);

        let metrics = registry.get_all_for_export();

        if let ExportValues::Histogram(values) = &metrics[0].values {
            let hist = values.get("").unwrap();
            assert_eq!(hist.count, 3);
        }
    }

    #[test]
    fn test_duplicate_metric() {
        let mut registry = MetricsRegistry::new();

        registry.register_counter("test", "Help", &[]).unwrap();

        let result = registry.register_counter("test", "Help", &[]);
        assert!(result.is_err());
    }
}