//! Comprehensive Training Metrics for ML Training Service
//!
//! Provides detailed Prometheus metrics for model training, GPU utilization,
//! convergence monitoring, and checkpoint management.
//!
//! Metric registration uses [`unwrap_metric`] instead of `.unwrap()` so that
//! `deny(clippy::unwrap_used)` is satisfied. Registration failures with
//! hard-coded literal names indicate a programmer mistake (duplicate name) and
//! are fatal at startup; `abort()` is the appropriate response.

use once_cell::sync::Lazy;
use prometheus::{
    register_counter_vec, register_gauge, register_gauge_vec, register_histogram_vec, CounterVec,
    Gauge, GaugeVec, HistogramVec,
};

/// Unwrap a Prometheus metric registration result.
///
/// Registration with hard-coded literal names is infallible at runtime; a
/// failure means a duplicate metric name (programmer error) caught at first
/// access during `Lazy` init. Aborting is safer than panicking because it
/// avoids unwind-related UB in static initializers.
fn unwrap_metric<T>(result: prometheus::Result<T>) -> T {
    match result {
        Ok(v) => v,
        Err(e) => {
            eprintln!("FATAL: Prometheus metric registration failed: {e}");
            std::process::abort();
        }
    }
}

// ============================================================================
// Training Progress Metrics
// ============================================================================

/// Training loss per epoch (labeled by model_type)
pub static TRAINING_LOSS: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_loss",
        "Current training loss value per model",
        &["model_type", "job_id"]
    ))
});

/// Validation loss per epoch (labeled by model_type)
pub static VALIDATION_LOSS: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_validation_loss",
        "Current validation loss value per model",
        &["model_type", "job_id"]
    ))
});

/// Current training epoch
pub static TRAINING_EPOCH: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_current_epoch",
        "Current training epoch number",
        &["model_type", "job_id"]
    ))
});

/// Training progress percentage (0-100)
pub static TRAINING_PROGRESS: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_progress_percent",
        "Training progress percentage",
        &["model_type", "job_id"]
    ))
});

/// Training speed (epochs per second)
pub static TRAINING_SPEED: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_epochs_per_second",
        "Training speed in epochs per second",
        &["model_type", "job_id"]
    ))
});

/// Training iteration duration histogram
pub static TRAINING_ITERATION_DURATION: Lazy<HistogramVec> = Lazy::new(|| {
    unwrap_metric(register_histogram_vec!(
        "ml_training_iteration_seconds",
        "Training iteration duration in seconds",
        &["model_type", "job_id"],
        vec![1.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0] // 1s to 10min
    ))
});

/// Model convergence rate (loss change per epoch)
pub static CONVERGENCE_RATE: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_convergence_rate",
        "Rate of loss decrease per epoch",
        &["model_type", "job_id"]
    ))
});

// ============================================================================
// NaN Detection and Training Errors
// ============================================================================

/// NaN detection count during training
pub static NAN_DETECTION_COUNT: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_training_nan_count",
        "Number of NaN values detected during training",
        &["model_type", "job_id", "tensor_type"] // tensor_type: loss, gradient, activation
    ))
});

/// Gradient explosion events
pub static GRADIENT_EXPLOSION_COUNT: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_training_gradient_explosion_total",
        "Number of gradient explosion events",
        &["model_type", "job_id"]
    ))
});

/// Training failures by error type
pub static TRAINING_FAILURES: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_training_failures_total",
        "Total training failures by error type",
        &["model_type", "error_type"] // error_type: nan, oom, timeout, config, etc.
    ))
});

// ============================================================================
// GPU Monitoring Metrics
// ============================================================================

/// GPU utilization percentage (0-100)
pub static GPU_UTILIZATION: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_gpu_utilization_percent",
        "GPU utilization percentage",
        &["gpu_id"]
    ))
});

/// GPU memory used in bytes
pub static GPU_MEMORY_USED: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_gpu_memory_used_bytes",
        "GPU memory used in bytes",
        &["gpu_id"]
    ))
});

/// GPU memory total in bytes
pub static GPU_MEMORY_TOTAL: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_gpu_memory_total_bytes",
        "Total GPU memory in bytes",
        &["gpu_id"]
    ))
});

/// GPU temperature in Celsius
pub static GPU_TEMPERATURE: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_gpu_temperature_celsius",
        "GPU temperature in Celsius",
        &["gpu_id"]
    ))
});

/// GPU power usage in watts
pub static GPU_POWER_USAGE: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_gpu_power_watts",
        "GPU power usage in watts",
        &["gpu_id"]
    ))
});

/// GPU errors total
pub static GPU_ERRORS: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_gpu_errors_total",
        "Total GPU errors detected",
        &["gpu_id", "error_type"]
    ))
});

// ============================================================================
// Checkpoint Management Metrics
// ============================================================================

/// Checkpoint save operations (success)
pub static CHECKPOINT_SAVES_SUCCESS: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_checkpoint_saves_total",
        "Total successful checkpoint saves",
        &["model_type", "job_id"]
    ))
});

/// Checkpoint save failures
pub static CHECKPOINT_SAVE_FAILURES: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_checkpoint_save_failures_total",
        "Total checkpoint save failures",
        &["model_type", "job_id", "error_type"]
    ))
});

/// Checkpoint save duration
pub static CHECKPOINT_SAVE_DURATION: Lazy<HistogramVec> = Lazy::new(|| {
    unwrap_metric(register_histogram_vec!(
        "ml_checkpoint_save_duration_seconds",
        "Checkpoint save duration in seconds",
        &["model_type", "job_id"],
        vec![0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0] // 100ms to 1min
    ))
});

/// Checkpoint size in bytes
pub static CHECKPOINT_SIZE: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_checkpoint_size_bytes",
        "Checkpoint file size in bytes",
        &["model_type", "job_id"]
    ))
});

// ============================================================================
// Model Performance Metrics
// ============================================================================

/// Model accuracy (0-1)
pub static MODEL_ACCURACY: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_model_accuracy",
        "Model accuracy on validation set",
        &["model_type", "job_id"]
    ))
});

/// Model F1 score (0-1)
pub static MODEL_F1_SCORE: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_model_f1_score",
        "Model F1 score on validation set",
        &["model_type", "job_id"]
    ))
});

/// Model precision (0-1)
pub static MODEL_PRECISION: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_model_precision",
        "Model precision on validation set",
        &["model_type", "job_id"]
    ))
});

/// Model recall (0-1)
pub static MODEL_RECALL: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_model_recall",
        "Model recall on validation set",
        &["model_type", "job_id"]
    ))
});

// ============================================================================
// Resource Usage Metrics
// ============================================================================

/// System memory used during training (bytes)
pub static SYSTEM_MEMORY_USED: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_memory_used_bytes",
        "System memory used during training",
        &["job_id"]
    ))
});

/// System CPU usage during training (percentage)
pub static SYSTEM_CPU_USAGE: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_cpu_usage_percent",
        "CPU usage percentage during training",
        &["job_id"]
    ))
});

/// Active training workers
pub static ACTIVE_WORKERS: Lazy<Gauge> = Lazy::new(|| {
    unwrap_metric(register_gauge!(
        "ml_training_active_workers",
        "Number of active training workers"
    ))
});

// ============================================================================
// Data Pipeline Metrics
// ============================================================================

/// Training data batches processed
pub static DATA_BATCHES_PROCESSED: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_training_data_batches_total",
        "Total data batches processed",
        &["model_type", "job_id"]
    ))
});

/// Data loading time histogram
pub static DATA_LOADING_DURATION: Lazy<HistogramVec> = Lazy::new(|| {
    unwrap_metric(register_histogram_vec!(
        "ml_training_data_loading_seconds",
        "Data loading duration per batch",
        &["model_type", "job_id"],
        vec![0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0] // 1ms to 1s
    ))
});

/// Feature engineering errors
pub static FEATURE_ENGINEERING_ERRORS: Lazy<CounterVec> = Lazy::new(|| {
    unwrap_metric(register_counter_vec!(
        "ml_feature_engineering_errors_total",
        "Total feature engineering errors",
        &["error_type"]
    ))
});

// ============================================================================
// Training Job Lifecycle Metrics
// ============================================================================

/// Training jobs by status
pub static TRAINING_JOBS_BY_STATUS: Lazy<GaugeVec> = Lazy::new(|| {
    unwrap_metric(register_gauge_vec!(
        "ml_training_jobs_by_status",
        "Number of training jobs by status",
        &["status"] // pending, running, completed, failed, stopped, paused
    ))
});

/// Training job duration histogram
pub static TRAINING_JOB_DURATION: Lazy<HistogramVec> = Lazy::new(|| {
    unwrap_metric(register_histogram_vec!(
        "ml_training_job_duration_seconds",
        "Total training job duration",
        &["model_type", "status"],
        vec![60.0, 300.0, 600.0, 1800.0, 3600.0, 7200.0, 14400.0] // 1min to 4h
    ))
});

// ============================================================================
// Helper Functions
// ============================================================================

/// Initialize all training metrics (registers them with global registry)
pub fn init_metrics() {
    // Force lazy initialization by accessing each metric
    let _ = &*TRAINING_LOSS;
    let _ = &*VALIDATION_LOSS;
    let _ = &*TRAINING_EPOCH;
    let _ = &*TRAINING_PROGRESS;
    let _ = &*TRAINING_SPEED;
    let _ = &*TRAINING_ITERATION_DURATION;
    let _ = &*CONVERGENCE_RATE;
    let _ = &*NAN_DETECTION_COUNT;
    let _ = &*GRADIENT_EXPLOSION_COUNT;
    let _ = &*TRAINING_FAILURES;
    let _ = &*GPU_UTILIZATION;
    let _ = &*GPU_MEMORY_USED;
    let _ = &*GPU_MEMORY_TOTAL;
    let _ = &*GPU_TEMPERATURE;
    let _ = &*GPU_POWER_USAGE;
    let _ = &*GPU_ERRORS;
    let _ = &*CHECKPOINT_SAVES_SUCCESS;
    let _ = &*CHECKPOINT_SAVE_FAILURES;
    let _ = &*CHECKPOINT_SAVE_DURATION;
    let _ = &*CHECKPOINT_SIZE;
    let _ = &*MODEL_ACCURACY;
    let _ = &*MODEL_F1_SCORE;
    let _ = &*MODEL_PRECISION;
    let _ = &*MODEL_RECALL;
    let _ = &*SYSTEM_MEMORY_USED;
    let _ = &*SYSTEM_CPU_USAGE;
    let _ = &*ACTIVE_WORKERS;
    let _ = &*DATA_BATCHES_PROCESSED;
    let _ = &*DATA_LOADING_DURATION;
    let _ = &*FEATURE_ENGINEERING_ERRORS;
    let _ = &*TRAINING_JOBS_BY_STATUS;
    let _ = &*TRAINING_JOB_DURATION;
}

/// Helper: Record training iteration metrics
pub fn record_training_iteration(
    model_type: &str,
    job_id: &str,
    epoch: u32,
    train_loss: f64,
    val_loss: f64,
    duration_secs: f64,
) {
    TRAINING_EPOCH
        .with_label_values(&[model_type, job_id])
        .set(epoch as f64);
    TRAINING_LOSS
        .with_label_values(&[model_type, job_id])
        .set(train_loss);
    VALIDATION_LOSS
        .with_label_values(&[model_type, job_id])
        .set(val_loss);
    TRAINING_ITERATION_DURATION
        .with_label_values(&[model_type, job_id])
        .observe(duration_secs);
}

/// Helper: Record NaN detection
pub fn record_nan_detection(model_type: &str, job_id: &str, tensor_type: &str) {
    NAN_DETECTION_COUNT
        .with_label_values(&[model_type, job_id, tensor_type])
        .inc();
}

/// Helper: Record GPU metrics
pub fn record_gpu_metrics(
    gpu_id: &str,
    utilization_percent: f64,
    memory_used_bytes: f64,
    memory_total_bytes: f64,
    temperature_celsius: f64,
) {
    GPU_UTILIZATION
        .with_label_values(&[gpu_id])
        .set(utilization_percent);
    GPU_MEMORY_USED
        .with_label_values(&[gpu_id])
        .set(memory_used_bytes);
    GPU_MEMORY_TOTAL
        .with_label_values(&[gpu_id])
        .set(memory_total_bytes);
    GPU_TEMPERATURE
        .with_label_values(&[gpu_id])
        .set(temperature_celsius);
}

/// Helper: Record checkpoint save
pub fn record_checkpoint_save(
    model_type: &str,
    job_id: &str,
    success: bool,
    duration_secs: f64,
    size_bytes: u64,
    error_type: Option<&str>,
) {
    if success {
        CHECKPOINT_SAVES_SUCCESS
            .with_label_values(&[model_type, job_id])
            .inc();
        CHECKPOINT_SAVE_DURATION
            .with_label_values(&[model_type, job_id])
            .observe(duration_secs);
        CHECKPOINT_SIZE
            .with_label_values(&[model_type, job_id])
            .set(size_bytes as f64);
    } else {
        let error_label = error_type.unwrap_or("unknown");
        CHECKPOINT_SAVE_FAILURES
            .with_label_values(&[model_type, job_id, error_label])
            .inc();
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[test]
    fn test_metrics_initialization() {
        init_metrics();
        // All metrics should be registered without panic
    }

    #[test]
    fn test_record_training_iteration() {
        init_metrics();
        record_training_iteration("dqn", "test-job-1", 5, 0.25, 0.28, 12.5);
        // Should not panic
    }

    #[test]
    fn test_record_nan_detection() {
        init_metrics();
        record_nan_detection("mamba", "test-job-2", "loss");
        // Should not panic
    }

    #[test]
    fn test_record_gpu_metrics() {
        init_metrics();
        record_gpu_metrics("0", 85.5, 6.5e9, 8.0e9, 72.0);
        // Should not panic
    }

    #[test]
    fn test_record_checkpoint_save() {
        init_metrics();
        record_checkpoint_save("tft", "test-job-3", true, 5.2, 1024 * 1024 * 512, None);
        record_checkpoint_save("ppo", "test-job-4", false, 0.0, 0, Some("disk_full"));
        // Should not panic
    }
}
