mecha10-core 0.8.0

//! Metrics Collection
//!
//! Standardized metrics collection system for Mecha10 applications.
//!
//! # Features
//!
//! - Multiple metric types (Counter, Gauge, Histogram)
//! - Multiple backends (Prometheus, StatsD, OpenTelemetry)
//! - Automatic node metrics collection
//! - MLflow integration for experiment tracking
//! - Low overhead with buffering
//! - Type-safe metric names and labels
//!
//! # Architecture
//!
//! ```text
//! ┌────────────┐
//! │   Node A   │─┐
//! └────────────┘ │
//!                │  Record Metrics
//! ┌────────────┐ │
//! │   Node B   │─┼──▶  ┌─────────────────┐
//! └────────────┘ │     │ MetricsRegistry │
//!                │     └─────────────────┘
//! ┌────────────┐ │              │
//! │   Node C   │─┘              │
//! └────────────┘                │
//!                               ▼
//!                    ┌──────────────────────┐
//!                    │  Metrics Backends    │
//!                    ├──────────────────────┤
//!                    │ • Prometheus         │
//!                    │ • StatsD             │
//!                    │ • OpenTelemetry      │
//!                    │ • MLflow             │
//!                    └──────────────────────┘
//! ```
//!
//! # Example
//!
//! ```rust
//! use mecha10_core::prelude::*;
//! use mecha10_core::metrics::{MetricsRegistry, Counter, Gauge};
//!
//! # async fn example() -> Result<()> {
//! // Create metrics registry
//! let metrics = MetricsRegistry::new();
//!
//! // Define metrics
//! let messages_processed = Counter::new("messages_processed", "Total messages processed");
//! let queue_size = Gauge::new("queue_size", "Current queue size");
//!
//! // Record metrics
//! metrics.increment(&messages_processed, 1.0);
//! metrics.set(&queue_size, 42.0);
//!
//! // Export to Prometheus
//! let prometheus_text = metrics.export_prometheus();
//! # Ok(())
//! # }
//! ```

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MetricType {
    Counter,
    Gauge,
    Histogram,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MetricName {
    pub name: String,
    pub labels: HashMap<String, String>,
}

impl std::hash::Hash for MetricName {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
        let mut labels: Vec<_> = self.labels.iter().collect();
        labels.sort_by_key(|(k, _)| *k);
        for (k, v) in labels {
            k.hash(state);
            v.hash(state);
        }
    }
}

impl MetricName {
    /// Create a new metric name
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            labels: HashMap::new(),
        }
    }

    crate::builder_setter!(
        /// Add a label
        insert, with_label, labels, impl Into<String>, impl Into<String>
    );

    /// Get Prometheus-style name with labels
    pub fn prometheus_name(&self) -> String {
        if self.labels.is_empty() {
            self.name.clone()
        } else {
            let labels: Vec<String> = self.labels.iter().map(|(k, v)| format!("{}=\"{}\"", k, v)).collect();
            format!("{}{{{}}}", self.name, labels.join(","))
        }
    }
}

#[derive(Debug, Clone)]
pub struct Counter {
    pub name: MetricName,
    pub help: String,
}

impl Counter {
    /// Create a new counter
    pub fn new(name: impl Into<String>, help: impl Into<String>) -> Self {
        let (name, help) = new_metric_name_and_help(name, help);
        Self { name, help }
    }

    /// Create with labels
    pub fn with_labels(name: impl Into<String>, help: impl Into<String>, labels: HashMap<String, String>) -> Self {
        let (name, help) = new_metric_name_and_help_with_labels(name, help, labels);
        Self { name, help }
    }
}

#[derive(Debug, Clone)]
pub struct Gauge {
    pub name: MetricName,
    pub help: String,
}

impl Gauge {
    /// Create a new gauge
    pub fn new(name: impl Into<String>, help: impl Into<String>) -> Self {
        let (name, help) = new_metric_name_and_help(name, help);
        Self { name, help }
    }

    /// Create with labels
    pub fn with_labels(name: impl Into<String>, help: impl Into<String>, labels: HashMap<String, String>) -> Self {
        let (name, help) = new_metric_name_and_help_with_labels(name, help, labels);
        Self { name, help }
    }
}

/// Shared body behind `Counter::new`/`Gauge::new` — both wrap a plain `(MetricName, help)`
/// pair, differing only in the metric-kind struct the caller assembles them into.
fn new_metric_name_and_help(name: impl Into<String>, help: impl Into<String>) -> (MetricName, String) {
    (MetricName::new(name), help.into())
}

/// Shared body behind `Counter::with_labels`/`Gauge::with_labels`.
fn new_metric_name_and_help_with_labels(
    name: impl Into<String>,
    help: impl Into<String>,
    labels: HashMap<String, String>,
) -> (MetricName, String) {
    (
        MetricName {
            name: name.into(),
            labels,
        },
        help.into(),
    )
}

#[derive(Debug, Clone)]
pub struct Histogram {
    pub name: MetricName,
    pub help: String,
    pub buckets: Vec<f64>,
}

impl Histogram {
    /// Create a new histogram with default buckets
    pub fn new(name: impl Into<String>, help: impl Into<String>) -> Self {
        Self {
            name: MetricName::new(name),
            help: help.into(),
            buckets: default_buckets(),
        }
    }

    /// Create with custom buckets
    pub fn with_buckets(name: impl Into<String>, help: impl Into<String>, buckets: Vec<f64>) -> Self {
        Self {
            name: MetricName::new(name),
            help: help.into(),
            buckets,
        }
    }
}

fn default_buckets() -> Vec<f64> {
    vec![0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0]
}

#[derive(Debug, Clone)]
struct CounterValue {
    value: f64,
}

#[derive(Debug, Clone)]
struct GaugeValue {
    value: f64,
}

#[derive(Debug, Clone)]
struct HistogramValue {
    count: u64,
    sum: f64,
    buckets: HashMap<usize, u64>,
}

pub struct MetricsRegistry {
    counters: Arc<RwLock<HashMap<MetricName, CounterValue>>>,
    gauges: Arc<RwLock<HashMap<MetricName, GaugeValue>>>,
    histograms: Arc<RwLock<HashMap<MetricName, HistogramValue>>>,
    metadata: Arc<RwLock<HashMap<String, (MetricType, String)>>>,
}

impl MetricsRegistry {
    /// Create a new metrics registry
    pub fn new() -> Self {
        Self {
            counters: Arc::new(RwLock::new(HashMap::new())),
            gauges: Arc::new(RwLock::new(HashMap::new())),
            histograms: Arc::new(RwLock::new(HashMap::new())),
            metadata: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Record a metric's type/help in the metadata table the first time it's seen.
    ///
    /// Shared by `increment`/`set`/`observe`, which each register their own metric's metadata
    /// before touching the value map — only the metric kind and the name/help source differ.
    async fn ensure_metadata(&self, name: &str, metric_type: MetricType, help: &str) {
        let mut metadata = self.metadata.write().await;
        metadata
            .entry(name.to_string())
            .or_insert((metric_type, help.to_string()));
    }

    /// Increment a counter
    pub async fn increment(&self, counter: &Counter, value: f64) {
        self.ensure_metadata(&counter.name.name, MetricType::Counter, &counter.help)
            .await;

        let mut counters = self.counters.write().await;
        counters
            .entry(counter.name.clone())
            .and_modify(|v| v.value += value)
            .or_insert(CounterValue { value });
    }

    /// Set a gauge value
    pub async fn set(&self, gauge: &Gauge, value: f64) {
        self.ensure_metadata(&gauge.name.name, MetricType::Gauge, &gauge.help)
            .await;

        let mut gauges = self.gauges.write().await;
        gauges
            .entry(gauge.name.clone())
            .and_modify(|v| v.value = value)
            .or_insert(GaugeValue { value });
    }

    /// Increment a gauge (delta)
    pub async fn increment_gauge(&self, gauge: &Gauge, delta: f64) {
        let mut gauges = self.gauges.write().await;
        gauges
            .entry(gauge.name.clone())
            .and_modify(|v| v.value += delta)
            .or_insert(GaugeValue { value: delta });
    }

    /// Observe a value in a histogram
    pub async fn observe(&self, histogram: &Histogram, value: f64) {
        self.ensure_metadata(&histogram.name.name, MetricType::Histogram, &histogram.help)
            .await;

        let bucket_idx = histogram
            .buckets
            .iter()
            .position(|&b| value <= b)
            .unwrap_or(histogram.buckets.len());

        let mut histograms = self.histograms.write().await;
        histograms
            .entry(histogram.name.clone())
            .and_modify(|h| {
                h.count += 1;
                h.sum += value;
                *h.buckets.entry(bucket_idx).or_insert(0) += 1;
            })
            .or_insert_with(|| {
                let mut buckets = HashMap::new();
                buckets.insert(bucket_idx, 1);
                HistogramValue {
                    count: 1,
                    sum: value,
                    buckets,
                }
            });
    }

    /// Export metrics in Prometheus text format
    pub async fn export_prometheus(&self) -> String {
        let mut output = String::new();

        {
            let counters = self.counters.read().await;
            let metadata = self.metadata.read().await;

            let mut seen_names = std::collections::HashSet::new();

            for (name, value) in counters.iter() {
                if !seen_names.contains(&name.name) {
                    if let Some((_, help)) = metadata.get(&name.name) {
                        output.push_str(&format!("# HELP {} {}\n", name.name, help));
                        output.push_str(&format!("# TYPE {} counter\n", name.name));
                    }
                    seen_names.insert(name.name.clone());
                }

                output.push_str(&format!("{} {}\n", name.prometheus_name(), value.value));
            }
        }

        {
            let gauges = self.gauges.read().await;
            let metadata = self.metadata.read().await;

            let mut seen_names = std::collections::HashSet::new();

            for (name, value) in gauges.iter() {
                if !seen_names.contains(&name.name) {
                    if let Some((_, help)) = metadata.get(&name.name) {
                        output.push_str(&format!("# HELP {} {}\n", name.name, help));
                        output.push_str(&format!("# TYPE {} gauge\n", name.name));
                    }
                    seen_names.insert(name.name.clone());
                }

                output.push_str(&format!("{} {}\n", name.prometheus_name(), value.value));
            }
        }

        {
            let histograms = self.histograms.read().await;
            let metadata = self.metadata.read().await;

            for (name, value) in histograms.iter() {
                if let Some((_, help)) = metadata.get(&name.name) {
                    output.push_str(&format!("# HELP {} {}\n", name.name, help));
                    output.push_str(&format!("# TYPE {} histogram\n", name.name));
                }

                for (bucket_idx, count) in &value.buckets {
                    output.push_str(&format!(
                        "{}_bucket{{le=\"{}\"}} {}\n",
                        name.prometheus_name(),
                        bucket_idx,
                        count
                    ));
                }

                output.push_str(&format!("{}_sum {}\n", name.prometheus_name(), value.sum));
                output.push_str(&format!("{}_count {}\n", name.prometheus_name(), value.count));
            }
        }

        output
    }

    /// Clear all metrics (useful for testing)
    pub async fn clear(&self) {
        self.counters.write().await.clear();
        self.gauges.write().await.clear();
        self.histograms.write().await.clear();
        self.metadata.write().await.clear();
    }
}

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

pub struct NodeMetrics {
    pub messages_received: Counter,
    pub messages_sent: Counter,
    pub errors_total: Counter,
    pub processing_duration: Histogram,
    pub active_subscriptions: Gauge,
    pub queue_size: Gauge,
}

impl NodeMetrics {
    /// Create standard node metrics for a given node
    pub fn new(node_id: &str) -> Self {
        let mut labels = HashMap::new();
        labels.insert("node".to_string(), node_id.to_string());

        Self {
            messages_received: Counter::with_labels(
                "messages_received_total",
                "Total messages received",
                labels.clone(),
            ),
            messages_sent: Counter::with_labels("messages_sent_total", "Total messages sent", labels.clone()),
            errors_total: Counter::with_labels("errors_total", "Total errors", labels.clone()),
            processing_duration: Histogram::with_buckets(
                "processing_duration_seconds",
                "Message processing duration",
                default_buckets(),
            ),
            active_subscriptions: Gauge::with_labels(
                "active_subscriptions",
                "Number of active subscriptions",
                labels.clone(),
            ),
            queue_size: Gauge::with_labels("queue_size", "Current queue size", labels),
        }
    }
}