veza/veza-common/src/metrics.rs

//! Metrics utilities for Veza Rust services
//! 
//! This module provides centralized metrics collection and reporting.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use prometheus::{
    Counter, CounterVec, Gauge, GaugeVec, Histogram, HistogramVec,
    HistogramOpts, Registry, Encoder, TextEncoder, Opts,
};

use crate::{VezaError, VezaResult};

/// Metrics collector for Veza services
pub struct MetricsCollector {
    registry: Registry,
    counters: HashMap<String, Counter>,
    counter_vecs: HashMap<String, CounterVec>,
    gauges: HashMap<String, Gauge>,
    gauge_vecs: HashMap<String, GaugeVec>,
    histograms: HashMap<String, Histogram>,
    histogram_vecs: HashMap<String, HistogramVec>,
}

impl MetricsCollector {
    /// Create a new metrics collector
    pub fn new() -> VezaResult<Self> {
        let registry = Registry::new();
        
        Ok(Self {
            registry,
            counters: HashMap::new(),
            counter_vecs: HashMap::new(),
            gauges: HashMap::new(),
            gauge_vecs: HashMap::new(),
            histograms: HashMap::new(),
            histogram_vecs: HashMap::new(),
        })
    }
    
    /// Register a counter
    pub fn register_counter(&mut self, name: &str, help: &str) -> VezaResult<()> {
        let counter = Counter::with_opts(Opts::new(name, help))
            .map_err(|e| VezaError::Internal(format!("Failed to create counter {}: {}", name, e)))?;
        
        self.registry.register(Box::new(counter.clone()))
            .map_err(|e| VezaError::Internal(format!("Failed to register counter {}: {}", name, e)))?;
        
        self.counters.insert(name.to_string(), counter);
        Ok(())
    }
    
    /// Register a counter vector
    pub fn register_counter_vec(&mut self, name: &str, help: &str, labels: &[&str]) -> VezaResult<()> {
        let counter_vec = CounterVec::new(Opts::new(name, help), labels)
            .map_err(|e| VezaError::Internal(format!("Failed to create counter vec {}: {}", name, e)))?;
        
        self.registry.register(Box::new(counter_vec.clone()))
            .map_err(|e| VezaError::Internal(format!("Failed to register counter vec {}: {}", name, e)))?;
        
        self.counter_vecs.insert(name.to_string(), counter_vec);
        Ok(())
    }
    
    /// Register a gauge
    pub fn register_gauge(&mut self, name: &str, help: &str) -> VezaResult<()> {
        let gauge = Gauge::with_opts(Opts::new(name, help))
            .map_err(|e| VezaError::Internal(format!("Failed to create gauge {}: {}", name, e)))?;
        
        self.registry.register(Box::new(gauge.clone()))
            .map_err(|e| VezaError::Internal(format!("Failed to register gauge {}: {}", name, e)))?;
        
        self.gauges.insert(name.to_string(), gauge);
        Ok(())
    }
    
    /// Register a gauge vector
    pub fn register_gauge_vec(&mut self, name: &str, help: &str, labels: &[&str]) -> VezaResult<()> {
        let gauge_vec = GaugeVec::new(Opts::new(name, help), labels)
            .map_err(|e| VezaError::Internal(format!("Failed to create gauge vec {}: {}", name, e)))?;
        
        self.registry.register(Box::new(gauge_vec.clone()))
            .map_err(|e| VezaError::Internal(format!("Failed to register gauge vec {}: {}", name, e)))?;
        
        self.gauge_vecs.insert(name.to_string(), gauge_vec);
        Ok(())
    }
    
    /// Register a histogram
    pub fn register_histogram(&mut self, name: &str, help: &str, buckets: Option<Vec<f64>>) -> VezaResult<()> {
        let mut opts = HistogramOpts::new(name, help);
        if let Some(buckets) = buckets {
            opts = opts.buckets(buckets);
        }
        let histogram = Histogram::with_opts(opts)
            .map_err(|e| VezaError::Internal(format!("Failed to create histogram {}: {}", name, e)))?;
        
        self.registry.register(Box::new(histogram.clone()))
            .map_err(|e| VezaError::Internal(format!("Failed to register histogram {}: {}", name, e)))?;
        
        self.histograms.insert(name.to_string(), histogram);
        Ok(())
    }
    
    /// Register a histogram vector
    pub fn register_histogram_vec(&mut self, name: &str, help: &str, labels: &[&str], buckets: Option<Vec<f64>>) -> VezaResult<()> {
        let mut opts = HistogramOpts::new(name, help);
        if let Some(buckets) = buckets {
            opts = opts.buckets(buckets);
        }
        let histogram_vec = HistogramVec::new(opts, labels)
            .map_err(|e| VezaError::Internal(format!("Failed to create histogram vec {}: {}", name, e)))?;
        
        self.registry.register(Box::new(histogram_vec.clone()))
            .map_err(|e| VezaError::Internal(format!("Failed to register histogram vec {}: {}", name, e)))?;
        
        self.histogram_vecs.insert(name.to_string(), histogram_vec);
        Ok(())
    }
    
    /// Increment a counter
    pub fn increment_counter(&self, name: &str) -> VezaResult<()> {
        if let Some(counter) = self.counters.get(name) {
            counter.inc();
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Counter {} not found", name)))
        }
    }
    
    /// Increment a counter by value
    pub fn increment_counter_by(&self, name: &str, value: f64) -> VezaResult<()> {
        if let Some(counter) = self.counters.get(name) {
            counter.inc_by(value);
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Counter {} not found", name)))
        }
    }
    
    /// Increment a counter vector
    pub fn increment_counter_vec(&self, name: &str, labels: &[&str], value: f64) -> VezaResult<()> {
        if let Some(counter_vec) = self.counter_vecs.get(name) {
            counter_vec.with_label_values(labels).inc_by(value);
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Counter vec {} not found", name)))
        }
    }
    
    /// Set a gauge value
    pub fn set_gauge(&self, name: &str, value: f64) -> VezaResult<()> {
        if let Some(gauge) = self.gauges.get(name) {
            gauge.set(value);
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Gauge {} not found", name)))
        }
    }
    
    /// Set a gauge vector value
    pub fn set_gauge_vec(&self, name: &str, labels: &[&str], value: f64) -> VezaResult<()> {
        if let Some(gauge_vec) = self.gauge_vecs.get(name) {
            gauge_vec.with_label_values(labels).set(value);
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Gauge vec {} not found", name)))
        }
    }
    
    /// Observe a histogram value
    pub fn observe_histogram(&self, name: &str, value: f64) -> VezaResult<()> {
        if let Some(histogram) = self.histograms.get(name) {
            histogram.observe(value);
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Histogram {} not found", name)))
        }
    }
    
    /// Observe a histogram vector value
    pub fn observe_histogram_vec(&self, name: &str, labels: &[&str], value: f64) -> VezaResult<()> {
        if let Some(histogram_vec) = self.histogram_vecs.get(name) {
            histogram_vec.with_label_values(labels).observe(value);
            Ok(())
        } else {
            Err(VezaError::Internal(format!("Histogram vec {} not found", name)))
        }
    }
    
    /// Get metrics as string
    pub fn get_metrics(&self) -> VezaResult<String> {
        let encoder = TextEncoder::new();
        let metric_families = self.registry.gather();
        
        encoder.encode_to_string(&metric_families)
            .map_err(|e| VezaError::Internal(format!("Failed to encode metrics: {}", e)))
    }
}

impl Default for MetricsCollector {
    fn default() -> Self {
        Self::new().expect("Failed to create metrics collector")
    }
}

/// Global metrics collector
pub static mut GLOBAL_METRICS: Option<Arc<MetricsCollector>> = None;

/// Initialize global metrics
pub fn init_global_metrics() -> VezaResult<()> {
    let mut collector = MetricsCollector::new()?;
    
    // Register common metrics
    collector.register_counter("veza_requests_total", "Total number of requests")?;
    collector.register_counter_vec("veza_requests_by_method", "Requests by HTTP method", &["method"])?;
    collector.register_counter_vec("veza_requests_by_status", "Requests by HTTP status", &["status"])?;
    collector.register_histogram("veza_request_duration_seconds", "Request duration in seconds", None)?;
    collector.register_gauge("veza_active_connections", "Number of active connections")?;
    collector.register_gauge("veza_active_users", "Number of active users")?;
    collector.register_counter("veza_errors_total", "Total number of errors")?;
    collector.register_counter_vec("veza_errors_by_type", "Errors by type", &["type"])?;
    collector.register_counter("veza_messages_total", "Total number of messages")?;
    collector.register_counter("veza_tracks_streamed_total", "Total number of tracks streamed")?;
    collector.register_histogram("veza_message_size_bytes", "Message size in bytes", None)?;
    collector.register_histogram("veza_track_duration_seconds", "Track duration in seconds", None)?;
    
    unsafe {
        GLOBAL_METRICS = Some(Arc::new(collector));
    }
    
    Ok(())
}

/// Get global metrics collector
pub fn get_global_metrics() -> VezaResult<Arc<MetricsCollector>> {
    unsafe {
        GLOBAL_METRICS
            .as_ref()
            .ok_or_else(|| VezaError::Internal("Global metrics not initialized".to_string()))
            .map(|m| m.clone())
    }
}

/// Increment request counter
pub fn increment_request_counter() -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter("veza_requests_total")
}

/// Increment request counter by method
pub fn increment_request_counter_by_method(method: &str) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter_vec("veza_requests_by_method", &[method], 1.0)
}

/// Increment request counter by status
pub fn increment_request_counter_by_status(status: u16) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter_vec("veza_requests_by_status", &[&status.to_string()], 1.0)
}

/// Observe request duration
pub fn observe_request_duration(duration: Duration) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.observe_histogram("veza_request_duration_seconds", duration.as_secs_f64())
}

/// Set active connections gauge
pub fn set_active_connections(count: f64) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.set_gauge("veza_active_connections", count)
}

/// Set active users gauge
pub fn set_active_users(count: f64) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.set_gauge("veza_active_users", count)
}

/// Increment error counter
pub fn increment_error_counter() -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter("veza_errors_total")
}

/// Increment error counter by type
pub fn increment_error_counter_by_type(error_type: &str) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter_vec("veza_errors_by_type", &[error_type], 1.0)
}

/// Increment message counter
pub fn increment_message_counter() -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter("veza_messages_total")
}

/// Increment tracks streamed counter
pub fn increment_tracks_streamed_counter() -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.increment_counter("veza_tracks_streamed_total")
}

/// Observe message size
pub fn observe_message_size(size: usize) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.observe_histogram("veza_message_size_bytes", size as f64)
}

/// Observe track duration
pub fn observe_track_duration(duration: Duration) -> VezaResult<()> {
    let metrics = get_global_metrics()?;
    metrics.observe_histogram("veza_track_duration_seconds", duration.as_secs_f64())
}

/// Get all metrics as string
pub fn get_all_metrics() -> VezaResult<String> {
    let metrics = get_global_metrics()?;
    metrics.get_metrics()
}

/// Timer for measuring duration
pub struct Timer {
    start: Instant,
    metric_name: String,
}

impl Timer {
    /// Start a new timer
    pub fn start(metric_name: &str) -> Self {
        Self {
            start: Instant::now(),
            metric_name: metric_name.to_string(),
        }
    }
    
    /// Stop the timer and record the metric
    pub fn stop(self) -> VezaResult<()> {
        let duration = self.start.elapsed();
        let metrics = get_global_metrics()?;
        metrics.observe_histogram(&self.metric_name, duration.as_secs_f64())
    }
}

/// Macro for timing operations
#[macro_export]
macro_rules! time_operation {
    ($metric_name:expr, $operation:expr) => {{
        let timer = crate::metrics::Timer::start($metric_name);
        let result = $operation;
        timer.stop()?;
        result
    }};
}