//! Metrics collection and management.
//!
//! This module provides Prometheus-compatible metrics for monitoring
//! RustyTorch++ ML inference and training workloads.

use crate::{MonitoringError, MonitoringResult};
use prometheus::{
    Counter, CounterVec, Gauge, GaugeVec, Histogram, HistogramOpts, HistogramVec, Opts, Registry,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Metric type enumeration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MetricType {
    /// Monotonically increasing counter
    Counter,
    /// Value that can go up and down
    Gauge,
    /// Distribution of values
    Histogram,
}

/// Custom metric definition.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CustomMetric {
    /// Metric name
    pub name: String,
    /// Help text
    pub help: String,
    /// Type of metric
    pub metric_type: MetricType,
    /// Label key-value pairs
    pub labels: HashMap<String, String>,
}

/// Metric registry for managing Prometheus metrics.
#[derive(Debug)]
pub struct MetricRegistry {
    registry: Registry,
    counters: HashMap<String, Counter>,
    gauges: HashMap<String, Gauge>,
    histograms: HashMap<String, Histogram>,
}

impl MetricRegistry {
    /// Create a new metric registry.
    pub fn new() -> Self {
        Self {
            registry: Registry::new(),
            counters: HashMap::new(),
            gauges: HashMap::new(),
            histograms: HashMap::new(),
        }
    }

    /// Register a counter metric.
    pub fn register_counter(&mut self, name: &str, help: &str) -> MonitoringResult<()> {
        let counter = Counter::new(name, help)?;
        self.registry.register(Box::new(counter.clone()))?;
        self.counters.insert(name.to_string(), counter);
        Ok(())
    }

    /// Register a gauge metric.
    pub fn register_gauge(&mut self, name: &str, help: &str) -> MonitoringResult<()> {
        let gauge = Gauge::new(name, help)?;
        self.registry.register(Box::new(gauge.clone()))?;
        self.gauges.insert(name.to_string(), gauge);
        Ok(())
    }

    /// Register a histogram metric.
    pub fn register_histogram(
        &mut self,
        name: &str,
        help: &str,
        buckets: Vec<f64>,
    ) -> MonitoringResult<()> {
        let opts = HistogramOpts::new(name, help).buckets(buckets);
        let histogram = Histogram::with_opts(opts)?;
        self.registry.register(Box::new(histogram.clone()))?;
        self.histograms.insert(name.to_string(), histogram);
        Ok(())
    }

    /// Increment a counter.
    pub fn increment_counter(&self, name: &str) -> MonitoringResult<()> {
        if let Some(counter) = self.counters.get(name) {
            counter.inc();
            Ok(())
        } else {
            Err(MonitoringError::metrics_error(format!(
                "Counter '{name}' not found"
            )))
        }
    }

    /// Add to a counter.
    pub fn add_counter(&self, name: &str, value: f64) -> MonitoringResult<()> {
        if let Some(counter) = self.counters.get(name) {
            counter.inc_by(value);
            Ok(())
        } else {
            Err(MonitoringError::metrics_error(format!(
                "Counter '{name}' not found"
            )))
        }
    }

    /// Set a gauge value.
    pub fn set_gauge(&self, name: &str, value: f64) -> MonitoringResult<()> {
        if let Some(gauge) = self.gauges.get(name) {
            gauge.set(value);
            Ok(())
        } else {
            Err(MonitoringError::metrics_error(format!(
                "Gauge '{name}' not found"
            )))
        }
    }

    /// Observe a histogram value.
    pub fn observe_histogram(&self, name: &str, value: f64) -> MonitoringResult<()> {
        if let Some(histogram) = self.histograms.get(name) {
            histogram.observe(value);
            Ok(())
        } else {
            Err(MonitoringError::metrics_error(format!(
                "Histogram '{name}' not found"
            )))
        }
    }

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

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

impl From<prometheus::Error> for MonitoringError {
    fn from(err: prometheus::Error) -> Self {
        Self::metrics_error(err.to_string())
    }
}

// ============================================================================
// ML-Specific Metrics
// ============================================================================

/// Default latency buckets for inference (in seconds)
pub const INFERENCE_LATENCY_BUCKETS: &[f64] = &[
    0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0,
];

/// Default batch size buckets
pub const BATCH_SIZE_BUCKETS: &[f64] = &[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0];

/// ML inference metrics collection
#[derive(Debug)]
pub struct InferenceMetrics {
    /// Total inference requests
    pub requests_total: CounterVec,
    /// Total inference errors
    pub errors_total: CounterVec,
    /// Inference latency histogram
    pub latency_seconds: HistogramVec,
    /// Tokens generated (for LLMs)
    pub tokens_generated: CounterVec,
    /// Tokens per second (throughput)
    pub tokens_per_second: GaugeVec,
    /// Batch size distribution
    pub batch_size: HistogramVec,
    /// Active requests gauge
    pub active_requests: GaugeVec,
    /// Queue depth
    pub queue_depth: GaugeVec,
    /// Cache hit rate
    pub cache_hits: CounterVec,
    /// Cache misses
    pub cache_misses: CounterVec,
    /// GPU memory usage
    pub gpu_memory_bytes: GaugeVec,
    /// GPU utilization percentage
    pub gpu_utilization: GaugeVec,
    /// Model load time
    pub model_load_seconds: HistogramVec,
    /// First token latency (time to first token for streaming)
    pub time_to_first_token: HistogramVec,
    /// Registry reference
    registry: Registry,
}

impl InferenceMetrics {
    /// Create a new inference metrics instance with default configuration
    pub fn new() -> MonitoringResult<Self> {
        let registry = Registry::new();

        // Request counter with model and status labels
        let requests_total = CounterVec::new(
            Opts::new("rtx_inference_requests_total", "Total inference requests"),
            &["model", "status"],
        )?;
        registry.register(Box::new(requests_total.clone()))?;

        // Error counter with model and error_type labels
        let errors_total = CounterVec::new(
            Opts::new("rtx_inference_errors_total", "Total inference errors"),
            &["model", "error_type"],
        )?;
        registry.register(Box::new(errors_total.clone()))?;

        // Latency histogram
        let latency_seconds = HistogramVec::new(
            HistogramOpts::new(
                "rtx_inference_latency_seconds",
                "Inference latency in seconds",
            )
            .buckets(INFERENCE_LATENCY_BUCKETS.to_vec()),
            &["model", "batch_size"],
        )?;
        registry.register(Box::new(latency_seconds.clone()))?;

        // Tokens generated
        let tokens_generated = CounterVec::new(
            Opts::new("rtx_tokens_generated_total", "Total tokens generated"),
            &["model"],
        )?;
        registry.register(Box::new(tokens_generated.clone()))?;

        // Tokens per second gauge
        let tokens_per_second = GaugeVec::new(
            Opts::new("rtx_tokens_per_second", "Token generation throughput"),
            &["model"],
        )?;
        registry.register(Box::new(tokens_per_second.clone()))?;

        // Batch size histogram
        let batch_size = HistogramVec::new(
            HistogramOpts::new("rtx_batch_size", "Inference batch size distribution")
                .buckets(BATCH_SIZE_BUCKETS.to_vec()),
            &["model"],
        )?;
        registry.register(Box::new(batch_size.clone()))?;

        // Active requests gauge
        let active_requests = GaugeVec::new(
            Opts::new("rtx_active_requests", "Currently active inference requests"),
            &["model"],
        )?;
        registry.register(Box::new(active_requests.clone()))?;

        // Queue depth gauge
        let queue_depth = GaugeVec::new(
            Opts::new("rtx_queue_depth", "Number of requests waiting in queue"),
            &["model"],
        )?;
        registry.register(Box::new(queue_depth.clone()))?;

        // Cache metrics
        let cache_hits = CounterVec::new(
            Opts::new("rtx_cache_hits_total", "KV cache hits"),
            &["model", "cache_type"],
        )?;
        registry.register(Box::new(cache_hits.clone()))?;

        let cache_misses = CounterVec::new(
            Opts::new("rtx_cache_misses_total", "KV cache misses"),
            &["model", "cache_type"],
        )?;
        registry.register(Box::new(cache_misses.clone()))?;

        // GPU metrics
        let gpu_memory_bytes = GaugeVec::new(
            Opts::new("rtx_gpu_memory_bytes", "GPU memory usage in bytes"),
            &["device", "type"],
        )?;
        registry.register(Box::new(gpu_memory_bytes.clone()))?;

        let gpu_utilization = GaugeVec::new(
            Opts::new("rtx_gpu_utilization_percent", "GPU utilization percentage"),
            &["device"],
        )?;
        registry.register(Box::new(gpu_utilization.clone()))?;

        // Model load time
        let model_load_seconds = HistogramVec::new(
            HistogramOpts::new("rtx_model_load_seconds", "Time to load model")
                .buckets(vec![0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0]),
            &["model"],
        )?;
        registry.register(Box::new(model_load_seconds.clone()))?;

        // Time to first token
        let time_to_first_token = HistogramVec::new(
            HistogramOpts::new(
                "rtx_time_to_first_token_seconds",
                "Time to generate first token",
            )
            .buckets(vec![0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0]),
            &["model"],
        )?;
        registry.register(Box::new(time_to_first_token.clone()))?;

        Ok(Self {
            requests_total,
            errors_total,
            latency_seconds,
            tokens_generated,
            tokens_per_second,
            batch_size,
            active_requests,
            queue_depth,
            cache_hits,
            cache_misses,
            gpu_memory_bytes,
            gpu_utilization,
            model_load_seconds,
            time_to_first_token,
            registry,
        })
    }

    /// Record a successful inference request
    pub fn record_request(&self, model: &str, latency_secs: f64, batch_size_val: usize) {
        self.requests_total
            .with_label_values(&[model, "success"])
            .inc();
        self.latency_seconds
            .with_label_values(&[model, &batch_size_val.to_string()])
            .observe(latency_secs);
        self.batch_size
            .with_label_values(&[model])
            .observe(batch_size_val as f64);
    }

    /// Record a failed inference request
    pub fn record_error(&self, model: &str, error_type: &str) {
        self.requests_total
            .with_label_values(&[model, "error"])
            .inc();
        self.errors_total
            .with_label_values(&[model, error_type])
            .inc();
    }

    /// Record tokens generated
    pub fn record_tokens(&self, model: &str, count: u64) {
        self.tokens_generated
            .with_label_values(&[model])
            .inc_by(count as f64);
    }

    /// Update tokens per second
    pub fn update_throughput(&self, model: &str, tokens_per_sec: f64) {
        self.tokens_per_second
            .with_label_values(&[model])
            .set(tokens_per_sec);
    }

    /// Update active request count
    pub fn set_active_requests(&self, model: &str, count: i64) {
        self.active_requests
            .with_label_values(&[model])
            .set(count as f64);
    }

    /// Update queue depth
    pub fn set_queue_depth(&self, model: &str, depth: i64) {
        self.queue_depth
            .with_label_values(&[model])
            .set(depth as f64);
    }

    /// Record cache hit
    pub fn record_cache_hit(&self, model: &str, cache_type: &str) {
        self.cache_hits
            .with_label_values(&[model, cache_type])
            .inc();
    }

    /// Record cache miss
    pub fn record_cache_miss(&self, model: &str, cache_type: &str) {
        self.cache_misses
            .with_label_values(&[model, cache_type])
            .inc();
    }

    /// Update GPU memory usage
    pub fn set_gpu_memory(&self, device: &str, used_bytes: u64, total_bytes: u64) {
        self.gpu_memory_bytes
            .with_label_values(&[device, "used"])
            .set(used_bytes as f64);
        self.gpu_memory_bytes
            .with_label_values(&[device, "total"])
            .set(total_bytes as f64);
    }

    /// Update GPU utilization
    pub fn set_gpu_utilization(&self, device: &str, utilization_percent: f64) {
        self.gpu_utilization
            .with_label_values(&[device])
            .set(utilization_percent);
    }

    /// Record model load time
    pub fn record_model_load(&self, model: &str, load_time_secs: f64) {
        self.model_load_seconds
            .with_label_values(&[model])
            .observe(load_time_secs);
    }

    /// Record time to first token
    pub fn record_ttft(&self, model: &str, ttft_secs: f64) {
        self.time_to_first_token
            .with_label_values(&[model])
            .observe(ttft_secs);
    }

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

    /// Export metrics in Prometheus format
    pub fn export(&self) -> String {
        use prometheus::Encoder;
        let encoder = prometheus::TextEncoder::new();
        let metric_families = self.registry.gather();
        let mut buffer = Vec::new();
        encoder
            .encode(&metric_families, &mut buffer)
            .unwrap_or_default();
        String::from_utf8(buffer).unwrap_or_default()
    }
}

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

/// Training metrics collection
#[derive(Debug)]
pub struct TrainingMetrics {
    /// Training steps completed
    pub steps_total: CounterVec,
    /// Training loss
    pub loss: GaugeVec,
    /// Learning rate
    pub learning_rate: GaugeVec,
    /// Gradient norm
    pub gradient_norm: GaugeVec,
    /// Samples processed per second
    pub samples_per_second: GaugeVec,
    /// Epoch progress
    pub epoch: GaugeVec,
    /// Checkpoint save time
    pub checkpoint_save_seconds: HistogramVec,
    /// Registry
    registry: Registry,
}

impl TrainingMetrics {
    /// Create new training metrics
    pub fn new() -> MonitoringResult<Self> {
        let registry = Registry::new();

        let steps_total = CounterVec::new(
            Opts::new("rtx_training_steps_total", "Total training steps"),
            &["model", "phase"],
        )?;
        registry.register(Box::new(steps_total.clone()))?;

        let loss = GaugeVec::new(
            Opts::new("rtx_training_loss", "Current training loss"),
            &["model", "loss_type"],
        )?;
        registry.register(Box::new(loss.clone()))?;

        let learning_rate = GaugeVec::new(
            Opts::new("rtx_learning_rate", "Current learning rate"),
            &["model"],
        )?;
        registry.register(Box::new(learning_rate.clone()))?;

        let gradient_norm =
            GaugeVec::new(Opts::new("rtx_gradient_norm", "Gradient norm"), &["model"])?;
        registry.register(Box::new(gradient_norm.clone()))?;

        let samples_per_second = GaugeVec::new(
            Opts::new("rtx_samples_per_second", "Training throughput"),
            &["model"],
        )?;
        registry.register(Box::new(samples_per_second.clone()))?;

        let epoch = GaugeVec::new(Opts::new("rtx_training_epoch", "Current epoch"), &["model"])?;
        registry.register(Box::new(epoch.clone()))?;

        let checkpoint_save_seconds = HistogramVec::new(
            HistogramOpts::new("rtx_checkpoint_save_seconds", "Checkpoint save time")
                .buckets(vec![1.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0]),
            &["model"],
        )?;
        registry.register(Box::new(checkpoint_save_seconds.clone()))?;

        Ok(Self {
            steps_total,
            loss,
            learning_rate,
            gradient_norm,
            samples_per_second,
            epoch,
            checkpoint_save_seconds,
            registry,
        })
    }

    /// Record a training step
    pub fn record_step(&self, model: &str, phase: &str, loss_val: f64, lr: f64, grad_norm: f64) {
        self.steps_total.with_label_values(&[model, phase]).inc();
        self.loss.with_label_values(&[model, "total"]).set(loss_val);
        self.learning_rate.with_label_values(&[model]).set(lr);
        self.gradient_norm
            .with_label_values(&[model])
            .set(grad_norm);
    }

    /// Update throughput
    pub fn update_throughput(&self, model: &str, samples_per_sec: f64) {
        self.samples_per_second
            .with_label_values(&[model])
            .set(samples_per_sec);
    }

    /// Set current epoch
    pub fn set_epoch(&self, model: &str, epoch_num: f64) {
        self.epoch.with_label_values(&[model]).set(epoch_num);
    }

    /// Record checkpoint save time
    pub fn record_checkpoint(&self, model: &str, save_time_secs: f64) {
        self.checkpoint_save_seconds
            .with_label_values(&[model])
            .observe(save_time_secs);
    }

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

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

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

    #[test]
    fn test_metric_registry() {
        let mut registry = MetricRegistry::new();
        assert!(
            registry
                .register_counter("test_counter", "Test counter")
                .is_ok()
        );
        assert!(registry.register_gauge("test_gauge", "Test gauge").is_ok());
        assert!(registry.increment_counter("test_counter").is_ok());
        assert!(registry.set_gauge("test_gauge", 42.0).is_ok());
    }

    #[test]
    fn test_inference_metrics() {
        let metrics = InferenceMetrics::new().unwrap();

        // Record some metrics
        metrics.record_request("gpt-2", 0.1, 8);
        metrics.record_tokens("gpt-2", 100);
        metrics.update_throughput("gpt-2", 500.0);
        metrics.set_active_requests("gpt-2", 5);
        metrics.record_cache_hit("gpt-2", "kv");
        metrics.set_gpu_memory("cuda:0", 4_000_000_000, 8_000_000_000);
        metrics.set_gpu_utilization("cuda:0", 75.0);

        // Export should contain our metrics
        let output = metrics.export();
        assert!(output.contains("rtx_inference_requests_total"));
        assert!(output.contains("rtx_tokens_generated_total"));
    }

    #[test]
    fn test_training_metrics() {
        let metrics = TrainingMetrics::new().unwrap();

        metrics.record_step("bert", "train", 0.5, 0.001, 1.5);
        metrics.update_throughput("bert", 1000.0);
        metrics.set_epoch("bert", 3.0);
        metrics.record_checkpoint("bert", 10.5);
    }
}