Rustelo/server/src/metrics.rs

//! Metrics collection module with basic Prometheus integration
//!
//! This module provides basic metrics collection for monitoring:
//! - HTTP request metrics (count, duration, status codes)
//! - Database connection metrics
//! - System resource metrics

#![allow(dead_code)]

use axum::{Router, extract::State, http::StatusCode, response::Response, routing::get};
use prometheus::{Encoder, IntCounter, IntGauge, Registry, TextEncoder};
use std::sync::Arc;
use std::time::Instant;
use tracing::{debug, error, warn};

use crate::AppState;

/// Basic metrics registry
pub struct MetricsRegistry {
    registry: Arc<Registry>,
    http_requests_total: IntCounter,
    http_requests_in_flight: IntGauge,
    db_connections_active: IntGauge,
    auth_requests_total: IntCounter,
    content_requests_total: IntCounter,
    email_sent_total: IntCounter,
}

impl MetricsRegistry {
    /// Create a new metrics registry with basic metrics
    pub fn new() -> Result<Self, prometheus::Error> {
        let registry = Arc::new(Registry::new());

        // HTTP metrics
        let http_requests_total =
            IntCounter::new("http_requests_total", "Total number of HTTP requests")?;

        let http_requests_in_flight = IntGauge::new(
            "http_requests_in_flight",
            "Current number of HTTP requests being processed",
        )?;

        // Database metrics
        let db_connections_active = IntGauge::new(
            "db_connections_active",
            "Number of active database connections",
        )?;

        // Authentication metrics
        let auth_requests_total = IntCounter::new(
            "auth_requests_total",
            "Total number of authentication requests",
        )?;

        // Content service metrics
        let content_requests_total =
            IntCounter::new("content_requests_total", "Total number of content requests")?;

        // Email service metrics
        let email_sent_total = IntCounter::new("email_sent_total", "Total number of emails sent")?;

        // Register all metrics
        registry.register(Box::new(http_requests_total.clone()))?;
        registry.register(Box::new(http_requests_in_flight.clone()))?;
        registry.register(Box::new(db_connections_active.clone()))?;
        registry.register(Box::new(auth_requests_total.clone()))?;
        registry.register(Box::new(content_requests_total.clone()))?;
        registry.register(Box::new(email_sent_total.clone()))?;

        Ok(Self {
            registry,
            http_requests_total,
            http_requests_in_flight,
            db_connections_active,
            auth_requests_total,
            content_requests_total,
            email_sent_total,
        })
    }

    /// Get the underlying registry
    pub fn registry(&self) -> &Registry {
        &self.registry
    }

    /// Increment HTTP request counter
    pub fn inc_http_requests(&self) {
        self.http_requests_total.inc();
    }

    /// Increment HTTP requests in flight
    pub fn inc_http_in_flight(&self) {
        self.http_requests_in_flight.inc();
    }

    /// Decrement HTTP requests in flight
    pub fn dec_http_in_flight(&self) {
        self.http_requests_in_flight.dec();
    }

    /// Set database connections active
    pub fn set_db_connections_active(&self, count: i64) {
        self.db_connections_active.set(count);
    }

    /// Increment auth requests
    pub fn inc_auth_requests(&self) {
        self.auth_requests_total.inc();
    }

    /// Increment content requests
    pub fn inc_content_requests(&self) {
        self.content_requests_total.inc();
    }

    /// Increment emails sent
    pub fn inc_emails_sent(&self) {
        self.email_sent_total.inc();
    }
}

/// Metrics service for collecting application metrics
pub struct MetricsService {
    registry: Option<Arc<MetricsRegistry>>,
}

impl MetricsService {
    /// Create a new metrics service
    pub fn new() -> Result<Self, prometheus::Error> {
        let registry = MetricsRegistry::new()?;
        Ok(Self {
            registry: Some(Arc::new(registry)),
        })
    }

    /// Get metrics registry
    pub fn registry(&self) -> Option<&Arc<MetricsRegistry>> {
        self.registry.as_ref()
    }

    /// Record HTTP request
    #[allow(dead_code)]
    pub fn record_http_request(&self) {
        if let Some(registry) = &self.registry {
            registry.inc_http_requests();
        }
    }

    /// Record HTTP request start
    #[allow(dead_code)]
    pub fn record_http_request_start(&self) {
        if let Some(registry) = &self.registry {
            registry.inc_http_in_flight();
        }
    }

    /// Record HTTP request end
    #[allow(dead_code)]
    pub fn record_http_request_end(&self) {
        if let Some(registry) = &self.registry {
            registry.dec_http_in_flight();
        }
    }

    /// Record database connection count
    #[allow(dead_code)]
    pub fn record_db_connections(&self, count: i64) {
        if let Some(registry) = &self.registry {
            registry.set_db_connections_active(count);
        }
    }

    /// Record authentication request
    #[allow(dead_code)]
    pub fn record_auth_request(&self) {
        if let Some(registry) = &self.registry {
            registry.inc_auth_requests();
        }
    }

    /// Record content request
    #[allow(dead_code)]
    pub fn record_content_request(&self) {
        if let Some(registry) = &self.registry {
            registry.inc_content_requests();
        }
    }

    /// Record email sent
    #[allow(dead_code)]
    pub fn record_email_sent(&self) {
        if let Some(registry) = &self.registry {
            registry.inc_emails_sent();
        }
    }
}

/// HTTP middleware for collecting request metrics
pub async fn metrics_middleware(
    State(state): State<AppState>,
    request: axum::extract::Request,
    next: axum::middleware::Next,
) -> Response {
    let start = Instant::now();

    // Record request start
    if let Some(metrics) = state.metrics_registry.as_ref() {
        metrics.inc_http_requests();
        metrics.inc_http_in_flight();
    }

    let response = next.run(request).await;

    // Record request end
    if let Some(metrics) = state.metrics_registry.as_ref() {
        metrics.dec_http_in_flight();
    }

    let duration = start.elapsed();
    debug!("Request completed in {:?}", duration);

    response
}

/// Handlers for metrics endpoints
pub mod handlers {
    use super::*;
    use axum::body::Body;
    use axum::http::header;
    use axum::response::Response;

    /// Prometheus metrics endpoint
    pub async fn metrics(State(state): State<AppState>) -> Result<Response<Body>, StatusCode> {
        if let Some(metrics) = state.metrics_registry.as_ref() {
            let encoder = TextEncoder::new();
            let metric_families = metrics.registry().gather();

            match encoder.encode_to_string(&metric_families) {
                Ok(output) => {
                    debug!("Serving metrics endpoint");
                    Ok(Response::builder()
                        .header(header::CONTENT_TYPE, encoder.format_type())
                        .body(Body::from(output))
                        .unwrap())
                }
                Err(e) => {
                    error!("Failed to encode metrics: {}", e);
                    Err(StatusCode::INTERNAL_SERVER_ERROR)
                }
            }
        } else {
            warn!("Metrics registry not available");
            Err(StatusCode::SERVICE_UNAVAILABLE)
        }
    }

    /// Health metrics endpoint (JSON format)
    pub async fn health_metrics(
        State(state): State<AppState>,
    ) -> Result<axum::Json<serde_json::Value>, StatusCode> {
        if let Some(metrics) = state.metrics_registry.as_ref() {
            let metric_families = metrics.registry().gather();

            let mut json_metrics = serde_json::Map::new();
            json_metrics.insert(
                "status".to_string(),
                serde_json::Value::String("healthy".to_string()),
            );
            json_metrics.insert(
                "metrics_count".to_string(),
                serde_json::Value::Number(metric_families.len().into()),
            );

            Ok(axum::Json(serde_json::Value::Object(json_metrics)))
        } else {
            warn!("Metrics registry not available");
            Err(StatusCode::SERVICE_UNAVAILABLE)
        }
    }

    /// Readiness probe endpoint
    pub async fn readiness_probe(
        State(state): State<AppState>,
    ) -> Result<axum::Json<serde_json::Value>, StatusCode> {
        // Basic readiness check
        let ready = state.metrics_registry.is_some();

        let mut response = serde_json::Map::new();
        response.insert("ready".to_string(), serde_json::Value::Bool(ready));

        if ready {
            Ok(axum::Json(serde_json::Value::Object(response)))
        } else {
            Err(StatusCode::SERVICE_UNAVAILABLE)
        }
    }

    /// Liveness probe endpoint
    pub async fn liveness_probe() -> axum::Json<serde_json::Value> {
        let mut response = serde_json::Map::new();
        response.insert("alive".to_string(), serde_json::Value::Bool(true));
        axum::Json(serde_json::Value::Object(response))
    }
}

/// Create metrics routes
pub fn create_metrics_routes() -> Router<AppState> {
    Router::new()
        .route("/metrics", get(handlers::metrics))
        .route("/metrics/health", get(handlers::health_metrics))
        .route("/health/ready", get(handlers::readiness_probe))
        .route("/health/live", get(handlers::liveness_probe))
}

/// Extension trait for AppState to include metrics
impl AppState {
    pub fn metrics_registry(&self) -> Option<&Arc<MetricsRegistry>> {
        self.metrics_registry.as_ref()
    }
}

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

    #[test]
    fn test_metrics_registry_creation() {
        let registry = MetricsRegistry::new();
        assert!(registry.is_ok());
    }

    #[test]
    fn test_metrics_service_creation() {
        let service = MetricsService::new();
        assert!(service.is_ok());
    }

    #[test]
    fn test_metrics_operations() {
        let service = MetricsService::new().unwrap();

        // Test recording various metrics
        service.record_http_request();
        service.record_auth_request();
        service.record_content_request();
        service.record_email_sent();
        service.record_db_connections(5);

        // Basic smoke test - if we get here without panicking, the operations work
        assert!(service.registry().is_some());
    }
}