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

pub mod capabilities;
pub mod l2_cache;
pub mod memory_profile;
pub mod profile;

use serde::{Deserialize, Serialize};

use crate::MLError;
use crate::device::MlDevice;

/// 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 `MlDevice`.
    ///
    /// - `Cpu` → always `MlDevice::Cpu`
    /// - `Cuda(id)` → `MlDevice::cuda(id)`, errors if unavailable
    /// - `Auto` → CUDA device 0 if available, else error
    pub fn resolve(&self) -> Result<MlDevice, MLError> {
        match self {
            DeviceConfig::Cpu => Ok(MlDevice::Cpu),
            #[cfg(feature = "cuda")]
            DeviceConfig::Cuda(id) => MlDevice::cuda(*id),
            #[cfg(not(feature = "cuda"))]
            DeviceConfig::Cuda(id) => Err(MLError::ConfigError(format!(
                "CUDA device {} requested but cuda feature not enabled",
                id
            ))),
            #[cfg(feature = "cuda")]
            DeviceConfig::Auto => MlDevice::cuda(0),
            #[cfg(not(feature = "cuda"))]
            DeviceConfig::Auto => Err(MLError::ConfigError(
                "CUDA required (DeviceConfig::Auto) but cuda feature not enabled".to_owned(),
            )),
        }
    }

    /// 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_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);
        }
    }
}
