prvng_platform/crates/rag/src/monitoring.rs

//! Monitoring and metrics collection for RAG system
//!
//! Provides comprehensive metrics tracking for:
//! - Query latency (P50, P95, P99)
//! - Cache hit rates
//! - Confidence score distribution
//! - Error rates and types
//! - Volume metrics
//!
//! Metrics can be exported in Prometheus format

use std::collections::VecDeque;
use std::sync::Arc;

use chrono::{DateTime, Utc};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};

/// Latency histogram for percentile calculations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LatencyHistogram {
    /// All recorded latencies (milliseconds)
    latencies: VecDeque<f32>,
    /// Maximum number of samples to keep
    max_samples: usize,
}

impl LatencyHistogram {
    /// Create a new latency histogram
    pub fn new(max_samples: usize) -> Self {
        Self {
            latencies: VecDeque::with_capacity(max_samples),
            max_samples,
        }
    }

    /// Record a latency measurement
    pub fn record(&mut self, latency_ms: f32) {
        if self.latencies.len() >= self.max_samples {
            self.latencies.pop_front();
        }
        self.latencies.push_back(latency_ms);
    }

    /// Get P50 (median) latency
    pub fn p50(&self) -> f32 {
        self.percentile(50.0)
    }

    /// Get P95 latency
    pub fn p95(&self) -> f32 {
        self.percentile(95.0)
    }

    /// Get P99 latency
    pub fn p99(&self) -> f32 {
        self.percentile(99.0)
    }

    /// Get percentile latency
    pub fn percentile(&self, p: f32) -> f32 {
        if self.latencies.is_empty() {
            return 0.0;
        }

        let mut sorted: Vec<_> = self.latencies.iter().copied().collect();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let index = ((p / 100.0) * sorted.len() as f32).ceil() as usize;
        let index = index.saturating_sub(1).min(sorted.len() - 1);

        sorted[index]
    }

    /// Get average latency
    pub fn average(&self) -> f32 {
        if self.latencies.is_empty() {
            return 0.0;
        }

        let sum: f32 = self.latencies.iter().sum();
        sum / self.latencies.len() as f32
    }

    /// Get minimum latency
    pub fn min(&self) -> f32 {
        self.latencies.iter().copied().fold(f32::INFINITY, f32::min)
    }

    /// Get maximum latency
    pub fn max(&self) -> f32 {
        self.latencies.iter().copied().fold(0.0, f32::max)
    }
}

/// Confidence score distribution
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ConfidenceDistribution {
    /// Scores in ranges: [0.0-0.25), [0.25-0.5), [0.5-0.75), [0.75-1.0]
    ranges: [u64; 4],
    total: u64,
}

impl ConfidenceDistribution {
    /// Create a new confidence distribution
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a confidence score
    pub fn record(&mut self, score: f32) {
        let score = score.clamp(0.0, 1.0);
        let idx = (score * 4.0).min(3.0) as usize;
        self.ranges[idx] += 1;
        self.total += 1;
    }

    /// Get average confidence score
    pub fn average(&self) -> f32 {
        if self.total == 0 {
            return 0.0;
        }

        let sum = (self.ranges[0] as f32 * 0.125)
            + (self.ranges[1] as f32 * 0.375)
            + (self.ranges[2] as f32 * 0.625)
            + (self.ranges[3] as f32 * 0.875);

        sum / self.total as f32
    }

    /// Get distribution as percentages
    pub fn percentages(&self) -> [f32; 4] {
        if self.total == 0 {
            return [0.0; 4];
        }

        [
            (self.ranges[0] as f32 / self.total as f32) * 100.0,
            (self.ranges[1] as f32 / self.total as f32) * 100.0,
            (self.ranges[2] as f32 / self.total as f32) * 100.0,
            (self.ranges[3] as f32 / self.total as f32) * 100.0,
        ]
    }
}

/// RAG system metrics
#[derive(Debug)]
pub struct RagMetrics {
    /// Query latency tracking
    pub query_latency: Arc<RwLock<LatencyHistogram>>,

    /// Retrieval latency tracking
    pub retrieval_latency: Arc<RwLock<LatencyHistogram>>,

    /// LLM latency tracking
    pub llm_latency: Arc<RwLock<LatencyHistogram>>,

    /// Confidence score distribution
    pub confidence: Arc<RwLock<ConfidenceDistribution>>,

    /// Total queries processed
    pub queries_total: u64,

    /// Total cache hits
    pub cache_hits: u64,

    /// Total cache misses
    pub cache_misses: u64,

    /// Total errors
    pub errors_total: u64,

    /// API errors
    pub api_errors: u64,

    /// Average sources per query
    pub avg_sources: f32,

    /// When metrics were created
    pub created_at: DateTime<Utc>,

    /// Last reset time
    pub last_reset: DateTime<Utc>,
}

// Manual Clone implementation for Arc<RwLock<T>>
impl Clone for RagMetrics {
    fn clone(&self) -> Self {
        Self {
            query_latency: Arc::clone(&self.query_latency),
            retrieval_latency: Arc::clone(&self.retrieval_latency),
            llm_latency: Arc::clone(&self.llm_latency),
            confidence: Arc::clone(&self.confidence),
            queries_total: self.queries_total,
            cache_hits: self.cache_hits,
            cache_misses: self.cache_misses,
            errors_total: self.errors_total,
            api_errors: self.api_errors,
            avg_sources: self.avg_sources,
            created_at: self.created_at,
            last_reset: self.last_reset,
        }
    }
}

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

impl RagMetrics {
    /// Create new metrics tracker
    pub fn new() -> Self {
        let now = Utc::now();
        Self {
            query_latency: Arc::new(RwLock::new(LatencyHistogram::new(1000))),
            retrieval_latency: Arc::new(RwLock::new(LatencyHistogram::new(1000))),
            llm_latency: Arc::new(RwLock::new(LatencyHistogram::new(1000))),
            confidence: Arc::new(RwLock::new(ConfidenceDistribution::new())),
            queries_total: 0,
            cache_hits: 0,
            cache_misses: 0,
            errors_total: 0,
            api_errors: 0,
            avg_sources: 0.0,
            created_at: now,
            last_reset: now,
        }
    }

    /// Record a query execution
    pub fn record_query(
        &mut self,
        total_latency_ms: f32,
        retrieval_latency_ms: f32,
        llm_latency_ms: f32,
        source_count: usize,
        confidence: f32,
    ) {
        self.queries_total += 1;

        // Update latency histograms
        self.query_latency.write().record(total_latency_ms);
        self.retrieval_latency.write().record(retrieval_latency_ms);
        self.llm_latency.write().record(llm_latency_ms);

        // Update confidence distribution
        self.confidence.write().record(confidence);

        // Update average sources
        self.avg_sources = (self.avg_sources * (self.queries_total - 1) as f32
            + source_count as f32)
            / self.queries_total as f32;
    }

    /// Record a cache hit
    pub fn record_cache_hit(&mut self, latency_ms: f32) {
        self.cache_hits += 1;
        self.query_latency.write().record(latency_ms);
    }

    /// Record a cache miss
    pub fn record_cache_miss(&mut self) {
        self.cache_misses += 1;
    }

    /// Record an error
    pub fn record_error(&mut self, error_type: ErrorType) {
        self.errors_total += 1;
        if let ErrorType::ApiError = error_type {
            self.api_errors += 1;
        }
    }

    /// Get cache hit rate
    pub fn cache_hit_rate(&self) -> f32 {
        let total = self.cache_hits + self.cache_misses;
        if total == 0 {
            0.0
        } else {
            self.cache_hits as f32 / total as f32
        }
    }

    /// Get error rate
    pub fn error_rate(&self) -> f32 {
        if self.queries_total == 0 {
            0.0
        } else {
            self.errors_total as f32 / self.queries_total as f32
        }
    }

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

        // Query latency metrics
        let ql = self.query_latency.read();
        output.push_str("# HELP rag_query_duration_seconds Query duration in seconds\n");
        output.push_str("# TYPE rag_query_duration_seconds histogram\n");
        output.push_str(&format!(
            "rag_query_duration_seconds{{quantile=\"0.5\"}} {}\n",
            ql.p50() / 1000.0
        ));
        output.push_str(&format!(
            "rag_query_duration_seconds{{quantile=\"0.95\"}} {}\n",
            ql.p95() / 1000.0
        ));
        output.push_str(&format!(
            "rag_query_duration_seconds{{quantile=\"0.99\"}} {}\n",
            ql.p99() / 1000.0
        ));

        // Retrieval latency
        let rl = self.retrieval_latency.read();
        output.push_str("# HELP rag_retrieval_duration_ms Retrieval duration in milliseconds\n");
        output.push_str("# TYPE rag_retrieval_duration_ms histogram\n");
        output.push_str(&format!(
            "rag_retrieval_duration_ms{{quantile=\"0.95\"}} {}\n",
            rl.p95()
        ));

        // LLM latency
        let ll = self.llm_latency.read();
        output.push_str("# HELP rag_llm_duration_ms LLM duration in milliseconds\n");
        output.push_str("# TYPE rag_llm_duration_ms histogram\n");
        output.push_str(&format!(
            "rag_llm_duration_ms{{quantile=\"0.95\"}} {}\n",
            ll.p95()
        ));

        // Cache metrics
        output.push_str("# HELP rag_cache_hits Total cache hits\n");
        output.push_str("# TYPE rag_cache_hits counter\n");
        output.push_str(&format!("rag_cache_hits {}\n", self.cache_hits));

        output.push_str("# HELP rag_cache_misses Total cache misses\n");
        output.push_str("# TYPE rag_cache_misses counter\n");
        output.push_str(&format!("rag_cache_misses {}\n", self.cache_misses));

        output.push_str("# HELP rag_cache_hit_rate Cache hit rate (0-1)\n");
        output.push_str("# TYPE rag_cache_hit_rate gauge\n");
        output.push_str(&format!("rag_cache_hit_rate {}\n", self.cache_hit_rate()));

        // Query metrics
        output.push_str("# HELP rag_queries_total Total queries processed\n");
        output.push_str("# TYPE rag_queries_total counter\n");
        output.push_str(&format!("rag_queries_total {}\n", self.queries_total));

        // Error metrics
        output.push_str("# HELP rag_errors_total Total errors\n");
        output.push_str("# TYPE rag_errors_total counter\n");
        output.push_str(&format!("rag_errors_total {}\n", self.errors_total));

        output.push_str("# HELP rag_error_rate Error rate (0-1)\n");
        output.push_str("# TYPE rag_error_rate gauge\n");
        output.push_str(&format!("rag_error_rate {}\n", self.error_rate()));

        // Confidence metrics
        let conf = self.confidence.read();
        output.push_str("# HELP rag_confidence_avg Average confidence score\n");
        output.push_str("# TYPE rag_confidence_avg gauge\n");
        output.push_str(&format!("rag_confidence_avg {}\n", conf.average()));

        // Sources metrics
        output.push_str("# HELP rag_avg_sources Average number of sources\n");
        output.push_str("# TYPE rag_avg_sources gauge\n");
        output.push_str(&format!("rag_avg_sources {}\n", self.avg_sources));

        output
    }

    /// Reset all metrics
    pub fn reset(&mut self) {
        self.query_latency = Arc::new(RwLock::new(LatencyHistogram::new(1000)));
        self.retrieval_latency = Arc::new(RwLock::new(LatencyHistogram::new(1000)));
        self.llm_latency = Arc::new(RwLock::new(LatencyHistogram::new(1000)));
        self.confidence = Arc::new(RwLock::new(ConfidenceDistribution::new()));
        self.queries_total = 0;
        self.cache_hits = 0;
        self.cache_misses = 0;
        self.errors_total = 0;
        self.api_errors = 0;
        self.avg_sources = 0.0;
        self.last_reset = Utc::now();
    }
}

/// Error types for metrics
#[derive(Debug, Clone, Copy)]
pub enum ErrorType {
    /// API error (network, rate limit, etc)
    ApiError,
    /// Retrieval error
    RetrievalError,
    /// LLM error
    LlmError,
    /// Other error
    Other,
}

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

    #[test]
    fn test_latency_histogram_percentiles() {
        let mut hist = LatencyHistogram::new(100);

        // Record some latencies
        for i in 1..=100 {
            hist.record(i as f32);
        }

        assert!(hist.p50() > 0.0);
        assert!(hist.p95() > hist.p50());
        assert!(hist.p99() > hist.p95());
    }

    #[test]
    fn test_latency_histogram_min_max() {
        let mut hist = LatencyHistogram::new(100);
        hist.record(10.0);
        hist.record(50.0);
        hist.record(100.0);

        assert_eq!(hist.min(), 10.0);
        assert_eq!(hist.max(), 100.0);
    }

    #[test]
    fn test_confidence_distribution() {
        let mut dist = ConfidenceDistribution::new();

        dist.record(0.1); // Low confidence
        dist.record(0.5); // Medium confidence
        dist.record(0.9); // High confidence

        let pcts = dist.percentages();
        assert!(pcts[0] > 0.0); // Low bucket has 1/3
        assert!(pcts[2] > 0.0); // High bucket has 1/3
    }

    #[test]
    fn test_rag_metrics_creation() {
        let metrics = RagMetrics::new();
        assert_eq!(metrics.queries_total, 0);
        assert_eq!(metrics.cache_hits, 0);
        assert_eq!(metrics.cache_hit_rate(), 0.0);
    }

    #[test]
    fn test_rag_metrics_record_query() {
        let mut metrics = RagMetrics::new();
        metrics.record_query(450.0, 100.0, 300.0, 3, 0.85);

        assert_eq!(metrics.queries_total, 1);
        assert_eq!(metrics.avg_sources, 3.0);
    }

    #[test]
    fn test_rag_metrics_cache_hit_rate() {
        let mut metrics = RagMetrics::new();
        metrics.record_cache_hit(50.0);
        metrics.record_cache_miss();
        metrics.record_cache_hit(45.0);

        assert_eq!(metrics.cache_hits, 2);
        assert_eq!(metrics.cache_misses, 1);
        assert!((metrics.cache_hit_rate() - 2.0 / 3.0).abs() < 0.01);
    }

    #[test]
    fn test_rag_metrics_prometheus_export() {
        let mut metrics = RagMetrics::new();
        metrics.record_query(450.0, 100.0, 300.0, 3, 0.85);
        metrics.record_cache_hit(50.0);

        let prometheus_output = metrics.export_prometheus();
        assert!(prometheus_output.contains("rag_query_duration_seconds"));
        assert!(prometheus_output.contains("rag_cache_hit_rate"));
        assert!(prometheus_output.contains("rag_queries_total"));
    }

    #[test]
    fn test_rag_metrics_reset() {
        let mut metrics = RagMetrics::new();
        metrics.record_query(450.0, 100.0, 300.0, 3, 0.85);
        assert_eq!(metrics.queries_total, 1);

        metrics.reset();
        assert_eq!(metrics.queries_total, 0);
    }

    #[test]
    fn test_latency_histogram_overflow() {
        let mut hist = LatencyHistogram::new(10);

        // Record more than capacity
        for i in 0..20 {
            hist.record(i as f32);
        }

        // Should only have 10 items
        assert!(hist.latencies.len() <= 10);
    }
}