//! GPU device types and inference strategy recommendation.

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum GpuVendor {
    Nvidia,
    Amd,
    Intel,
    Apple,
    Unknown,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuDevice {
    pub name: String,
    pub vendor: GpuVendor,
    pub vram_total_mb: u64,
    pub vram_free_mb: u64,
    pub compute_capability: Option<String>,
    pub driver_version: Option<String>,
    pub pci_bus_id: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum InferenceStrategy {
    CpuOnly,
    SingleGpu,
    TensorParallel,
    PhaseSplit,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuTopology {
    pub gpus: Vec<GpuDevice>,
    pub total_vram_mb: u64,
    pub total_free_vram_mb: u64,
    pub recommended_strategy: InferenceStrategy,
    pub cpu_cores: usize,
    pub ram_total_mb: u64,
    pub ram_free_mb: u64,
}

impl GpuTopology {
    pub fn recommend_strategy(gpus: &[GpuDevice]) -> InferenceStrategy {
        match gpus.len() {
            0 => InferenceStrategy::CpuOnly,
            1 => InferenceStrategy::SingleGpu,
            n if n <= 4 => InferenceStrategy::TensorParallel,
            _ => InferenceStrategy::PhaseSplit,
        }
    }
}

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

    #[test]
    fn test_strategy_cpu_only_when_no_gpus() {
        assert_eq!(GpuTopology::recommend_strategy(&[]), InferenceStrategy::CpuOnly);
    }

    #[test]
    fn test_strategy_single_gpu() {
        let gpu = GpuDevice {
            name: "RTX 4090".into(),
            vendor: GpuVendor::Nvidia,
            vram_total_mb: 24576,
            vram_free_mb: 20000,
            compute_capability: Some("8.9".into()),
            driver_version: Some("535.0".into()),
            pci_bus_id: Some("0000:01:00.0".into()),
        };
        assert_eq!(GpuTopology::recommend_strategy(&[gpu]), InferenceStrategy::SingleGpu);
    }

    #[test]
    fn test_strategy_tensor_parallel() {
        let gpus: Vec<GpuDevice> = (0..2).map(|i| GpuDevice {
            name: format!("GPU {i}"),
            vendor: GpuVendor::Nvidia,
            vram_total_mb: 24576,
            vram_free_mb: 20000,
            compute_capability: None,
            driver_version: None,
            pci_bus_id: None,
        }).collect();
        assert_eq!(GpuTopology::recommend_strategy(&gpus), InferenceStrategy::TensorParallel);
    }

    #[test]
    fn test_strategy_phase_split() {
        let gpus: Vec<GpuDevice> = (0..8).map(|i| GpuDevice {
            name: format!("GPU {i}"),
            vendor: GpuVendor::Nvidia,
            vram_total_mb: 80000,
            vram_free_mb: 75000,
            compute_capability: None,
            driver_version: None,
            pci_bus_id: None,
        }).collect();
        assert_eq!(GpuTopology::recommend_strategy(&gpus), InferenceStrategy::PhaseSplit);
    }
}
