//! GPU device management for ML training
//!
//! Provides centralized device selection, GPU capability detection,
//! and per-model batch size optimization.

pub mod capabilities;
pub mod memory_profile;

use candle_core::Device;
use serde::{Deserialize, Serialize};

use crate::MLError;

/// Device configuration for ML training and inference.
///
/// Use `Auto` for production -- it detects CUDA if available, falls back to CPU.
/// Use `Cuda(id)` to pin to a specific GPU. Use `Cpu` to force CPU mode.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum DeviceConfig {
    Cpu,
    Cuda(usize),
    /// Auto-detect: CUDA device 0 if available, else CPU (silent fallback -- never errors).
    Auto,
}

impl DeviceConfig {
    /// Resolve this config into a concrete candle `Device`.
    ///
    /// - `Cpu` → always `Device::Cpu`
    /// - `Cuda(id)` → `Device::new_cuda(id)`, errors if unavailable
    /// - `Auto` → CUDA device 0 if available, else CPU
    pub fn resolve(&self) -> Result<Device, MLError> {
        match self {
            DeviceConfig::Cpu => Ok(Device::Cpu),
            DeviceConfig::Cuda(id) => Device::new_cuda(*id).map_err(|e| {
                MLError::ConfigurationError(format!(
                    "CUDA device {} required but unavailable: {}",
                    id, e
                ))
            }),
            DeviceConfig::Auto => {
                match Device::new_cuda(0) {
                    Ok(dev) => Ok(dev),
                    Err(_) => Ok(Device::Cpu),
                }
            }
        }
    }

    /// Returns true if this config will attempt to use a GPU.
    pub fn is_gpu(&self) -> bool {
        matches!(self, DeviceConfig::Cuda(_) | DeviceConfig::Auto)
    }
}

impl Default for DeviceConfig {
    fn default() -> Self {
        DeviceConfig::Auto
    }
}

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

    #[test]
    fn test_cpu_always_resolves() {
        let device = DeviceConfig::Cpu.resolve().unwrap();
        assert!(!device.is_cuda());
    }

    #[test]
    fn test_auto_resolves_without_error() {
        let device = DeviceConfig::Auto.resolve().unwrap();
        let _ = device;
    }

    #[test]
    fn test_is_gpu() {
        assert!(!DeviceConfig::Cpu.is_gpu());
        assert!(DeviceConfig::Cuda(0).is_gpu());
        assert!(DeviceConfig::Auto.is_gpu());
    }

    #[test]
    fn test_default_is_auto() {
        assert_eq!(DeviceConfig::default(), DeviceConfig::Auto);
    }

    #[test]
    fn test_serde_roundtrip() {
        let configs = vec![DeviceConfig::Cpu, DeviceConfig::Cuda(0), DeviceConfig::Auto];
        for config in configs {
            let json = serde_json::to_string(&config).unwrap();
            let deserialized: DeviceConfig = serde_json::from_str(&json).unwrap();
            assert_eq!(config, deserialized);
        }
    }
}