//! GPU hardware profiles for VRAM-aware training configuration.
//!
//! This module provides GPU hardware management for the Foxhunt HFT Trading System benchmarks.
//! It handles device initialization, warmup protocols to eliminate cold-start effects,
//! and thermal monitoring to prevent throttling during benchmarks.
//!
//! # Features
//! - Automatic GPU detection with CPU fallback
//! - 10-pass warmup protocol (1000x1000 matrix multiplication)
//! - Real-time thermal monitoring via nvidia-smi
//! - Graceful error handling for CUDA failures

use ml_core::cuda_autograd::GpuTensor;
use ml_core::native_types::NativeDevice;
use ml_core::device::MlDevice;
use std::process::Command;
use std::time::{Duration, Instant};
use thiserror::Error;
use tracing::{debug, error, info, warn};

/// Errors that can occur during GPU hardware management
#[derive(Error, Debug)]
pub enum GpuHardwareError {
    #[error("CUDA initialization failed: {0}")]
    CudaInitFailed(String),

    #[error("GPU thermal throttling detected: {current}°C (threshold: {threshold}°C)")]
    ThermalThrottling { current: f32, threshold: f32 },

    #[error("GPU temperature warning: {current}°C (warning threshold: {threshold}°C)")]
    ThermalWarning { current: f32, threshold: f32 },

    #[error("Failed to read GPU temperature: {0}")]
    TemperatureReadFailed(String),

    #[error("GPU warmup failed: {0}")]
    WarmupFailed(String),

    #[error("Tensor operation failed: {0}")]
    TensorOpFailed(String),
}

/// Configuration for GPU hardware manager
#[derive(Debug, Clone)]
pub struct GpuHardwareConfig {
    /// Temperature threshold for throttling (°C)
    pub throttle_threshold: f32,
    /// Temperature threshold for warning (°C)
    pub warning_threshold: f32,
    /// Number of warmup passes
    pub warmup_passes: usize,
    /// Size of warmup matrix (NxN)
    pub warmup_matrix_size: usize,
}

impl Default for GpuHardwareConfig {
    fn default() -> Self {
        Self {
            throttle_threshold: 85.0,
            warning_threshold: 75.0,
            warmup_passes: 10,
            warmup_matrix_size: 1000,
        }
    }
}

/// GPU Hardware Manager for benchmark execution
///
/// Manages GPU device initialization, warmup protocol, and thermal monitoring.
/// Automatically falls back to CPU if GPU is unavailable.
#[derive(Debug)]
pub struct GpuHardwareManager {
    device: NativeDevice,
    ml_device: Option<MlDevice>,
    initial_temp_celsius: Option<f32>,
    config: GpuHardwareConfig,
    is_gpu: bool,
}

impl GpuHardwareManager {
    /// Create a new GPU hardware manager with default configuration
    ///
    /// # Errors
    /// Returns error if device initialization fails
    pub fn new() -> Result<Self, GpuHardwareError> {
        Self::with_config(GpuHardwareConfig::default())
    }

    /// Create a new GPU hardware manager with custom configuration
    ///
    /// # Arguments
    /// * `config` - Hardware configuration
    ///
    /// # Errors
    /// Returns error if device initialization fails
    pub fn with_config(config: GpuHardwareConfig) -> Result<Self, GpuHardwareError> {
        info!("Initializing GPU hardware manager...");

        // CUDA device required — no CPU fallback for benchmarks
        let ml_dev = MlDevice::cuda(0)
            .map_err(|e| GpuHardwareError::CudaInitFailed(format!("CUDA GPU required for benchmark: {}", e)))?;
        let device = NativeDevice::Cuda(0);
        let is_gpu = true;
        info!("CUDA device 0 initialized successfully");

        // Read initial GPU temperature if using GPU
        let initial_temp_celsius = if is_gpu {
            match Self::read_gpu_temperature() {
                Ok(temp) => {
                    info!("Initial GPU temperature: {:.1}°C", temp);
                    Some(temp)
                },
                Err(e) => {
                    warn!("Could not read GPU temperature: {}", e);
                    None
                },
            }
        } else {
            None
        };

        Ok(Self {
            device,
            ml_device: Some(ml_dev),
            initial_temp_celsius,
            config,
            is_gpu,
        })
    }

    /// Get reference to the device
    pub fn device(&self) -> &NativeDevice {
        &self.device
    }

    /// Check if using GPU
    pub fn is_gpu(&self) -> bool {
        self.is_gpu
    }

    /// Get current GPU temperature in Celsius
    ///
    /// # Errors
    /// Returns error if GPU unavailable or temperature read fails
    pub fn get_temperature(&self) -> Result<f32, GpuHardwareError> {
        if !self.is_gpu {
            return Err(GpuHardwareError::TemperatureReadFailed(
                "Not using GPU".to_owned(),
            ));
        }
        Self::read_gpu_temperature()
    }

    /// Check for thermal throttling
    ///
    /// # Returns
    /// - `Ok(true)` if throttling detected (temperature > threshold)
    /// - `Ok(false)` if temperature is safe
    /// - `Err` if GPU unavailable or temperature read fails
    pub fn check_thermal_throttling(&self) -> Result<bool, GpuHardwareError> {
        if !self.is_gpu {
            // CPU doesn't throttle in the same way
            return Ok(false);
        }

        let temp = Self::read_gpu_temperature()?;

        if temp >= self.config.throttle_threshold {
            error!(
                "⚠️  THERMAL THROTTLING: {:.1}°C >= {:.1}°C threshold",
                temp, self.config.throttle_threshold
            );
            Err(GpuHardwareError::ThermalThrottling {
                current: temp,
                threshold: self.config.throttle_threshold,
            })
        } else if temp >= self.config.warning_threshold {
            warn!(
                "⚠️  Temperature warning: {:.1}°C >= {:.1}°C (throttle at {:.1}°C)",
                temp, self.config.warning_threshold, self.config.throttle_threshold
            );
            Err(GpuHardwareError::ThermalWarning {
                current: temp,
                threshold: self.config.warning_threshold,
            })
        } else {
            debug!("✓ Temperature OK: {:.1}°C", temp);
            Ok(false)
        }
    }

    /// Execute warmup protocol to eliminate cold-start effects
    ///
    /// Performs N dummy forward passes with random tensors to warm up the GPU.
    /// This eliminates first-epoch variance caused by GPU initialization.
    ///
    /// # Returns
    /// Duration of warmup execution (not including in benchmark time)
    ///
    /// # Errors
    /// Returns error if warmup operations fail
    pub fn warmup(&self) -> Result<Duration, GpuHardwareError> {
        info!(
            "Starting GPU warmup: {} passes with {}x{} matrices",
            self.config.warmup_passes,
            self.config.warmup_matrix_size,
            self.config.warmup_matrix_size
        );

        let start = Instant::now();

        for pass in 0..self.config.warmup_passes {
            // Create random matrices on device to warm GPU memory subsystem
            let _a = self.create_random_matrix()?;
            let _b = self.create_random_matrix()?;

            // GPU is warmed by allocation + H2D transfer (matmul not needed for warmup)

            debug!(
                "Warmup pass {}/{} completed",
                pass + 1,
                self.config.warmup_passes
            );
        }

        let warmup_duration = start.elapsed();
        info!(
            "✓ Warmup completed in {:.2}ms ({} passes)",
            warmup_duration.as_secs_f64() * 1000.0,
            self.config.warmup_passes
        );

        Ok(warmup_duration)
    }

    /// Create a random matrix for warmup
    fn create_random_matrix(&self) -> Result<GpuTensor, GpuHardwareError> {
        let size = self.config.warmup_matrix_size;
        let data: Vec<f32> = (0..size * size).map(|_| fastrand::f32()).collect();

        let ml_dev = self.ml_device.as_ref().ok_or_else(|| {
            GpuHardwareError::TensorOpFailed("MlDevice not initialized".to_owned())
        })?;
        let stream = ml_dev.cuda_stream().map_err(|e| {
            GpuHardwareError::TensorOpFailed(format!("No CUDA stream: {}", e))
        })?;
        GpuTensor::from_host(&data, vec![size, size], stream)
            .map_err(|e| GpuHardwareError::TensorOpFailed(e.to_string()))
    }

    /// Read GPU temperature from nvidia-smi
    fn read_gpu_temperature() -> Result<f32, GpuHardwareError> {
        let output = Command::new("nvidia-smi")
            .args([
                "--query-gpu=temperature.gpu",
                "--format=csv,noheader,nounits",
            ])
            .output()
            .map_err(|e| {
                GpuHardwareError::TemperatureReadFailed(format!("nvidia-smi failed: {}", e))
            })?;

        if !output.status.success() {
            return Err(GpuHardwareError::TemperatureReadFailed(
                "nvidia-smi command failed".to_owned(),
            ));
        }

        let temp_str = String::from_utf8_lossy(&output.stdout);
        let temp = temp_str.trim().parse::<f32>().map_err(|e| {
            GpuHardwareError::TemperatureReadFailed(format!(
                "Failed to parse temperature '{}': {}",
                temp_str, e
            ))
        })?;

        Ok(temp)
    }

    /// Get initial temperature reading
    pub fn initial_temperature(&self) -> Option<f32> {
        self.initial_temp_celsius
    }

    /// Get configuration
    pub fn config(&self) -> &GpuHardwareConfig {
        &self.config
    }
}

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

    #[test]

    fn test_gpu_hardware_manager_creation() {
        let manager = GpuHardwareManager::new();
        assert!(
            manager.is_ok(),
            "Manager creation failed: {:?}",
            manager.err()
        );

        let manager = manager.unwrap();
        assert!(manager.is_gpu());
    }

    #[test]

    fn test_custom_config() {
        let config = GpuHardwareConfig {
            throttle_threshold: 90.0,
            warning_threshold: 80.0,
            warmup_passes: 5,
            warmup_matrix_size: 500,
        };

        let manager = GpuHardwareManager::with_config(config.clone());
        assert!(manager.is_ok());

        let manager = manager.unwrap();
        assert_eq!(manager.config().throttle_threshold, 90.0);
        assert_eq!(manager.config().warning_threshold, 80.0);
        assert_eq!(manager.config().warmup_passes, 5);
        assert_eq!(manager.config().warmup_matrix_size, 500);
    }

    #[test]

    fn test_device_access() {
        let manager = GpuHardwareManager::new().unwrap();
        let device = manager.device();

        // Should be either CPU or CUDA - just verify device is valid
        let _ = device; // Device is valid if we got here
    }

    #[test]

    fn test_warmup_protocol() {
        let manager = GpuHardwareManager::new().unwrap();

        if !manager.is_gpu() {
            info!("Skipping warmup test - no GPU available");
            return;
        }

        let warmup_result = manager.warmup();
        assert!(
            warmup_result.is_ok(),
            "Warmup failed: {:?}",
            warmup_result.err()
        );

        let duration = warmup_result.unwrap();
        assert!(
            duration.as_secs() < 60,
            "Warmup took too long: {:?}",
            duration
        );
    }

    #[test]

    fn test_temperature_reading() {
        let manager = GpuHardwareManager::new().unwrap();

        if !manager.is_gpu() {
            info!("Skipping temperature test - no GPU available");
            return;
        }

        // Should have initial temperature
        assert!(
            manager.initial_temperature().is_some(),
            "Initial temperature not recorded"
        );

        // Should be able to read current temperature
        let temp_result = manager.get_temperature();
        assert!(
            temp_result.is_ok(),
            "Temperature read failed: {:?}",
            temp_result.err()
        );

        let temp = temp_result.unwrap();
        assert!(
            temp > 0.0 && temp < 120.0,
            "Temperature out of reasonable range: {}°C",
            temp
        );
    }

    #[test]

    fn test_thermal_monitoring() {
        let manager = GpuHardwareManager::new().unwrap();

        if !manager.is_gpu() {
            info!("Skipping thermal test - no GPU available");
            return;
        }

        // Check thermal throttling
        let throttling_result = manager.check_thermal_throttling();

        match throttling_result {
            Ok(false) => {
                // Normal - temperature below warning threshold
                info!("Temperature OK");
            },
            Err(GpuHardwareError::ThermalWarning { current, threshold }) => {
                // Warning - temperature above warning threshold
                info!(current_temp = current, threshold = threshold, "Temperature warning");
            },
            Err(GpuHardwareError::ThermalThrottling { current, threshold }) => {
                // Critical - temperature above throttle threshold
                panic!(
                    "⚠️  THERMAL THROTTLING: {}°C >= {}°C - Cannot run benchmarks",
                    current, threshold
                );
            },
            Err(e) => {
                panic!("Thermal check failed: {:?}", e);
            },
            Ok(true) => {
                panic!("Unexpected throttling state");
            },
        }
    }

    #[test]

    fn test_warmup_with_custom_size() {
        let config = GpuHardwareConfig {
            warmup_passes: 3,
            warmup_matrix_size: 100,
            ..Default::default()
        };

        let manager = GpuHardwareManager::with_config(config).unwrap();

        // Warmup should succeed even with small matrices
        let warmup_result = manager.warmup();
        assert!(
            warmup_result.is_ok(),
            "Warmup with custom size failed: {:?}",
            warmup_result.err()
        );
    }

    #[test]
    fn test_error_display() {
        let err = GpuHardwareError::ThermalThrottling {
            current: 87.5,
            threshold: 85.0,
        };
        let msg = format!("{}", err);
        assert!(msg.contains("87.5"));
        assert!(msg.contains("85"));

        let err = GpuHardwareError::WarmupFailed("test error".to_owned());
        let msg = format!("{}", err);
        assert!(msg.contains("test error"));
    }
}
