use serde::{Deserialize, Serialize};
use std::process::Command;
use tracing::{debug, info, warn};

/// Represents a hardware-accelerated GPU endpoint on an edge relay.
/// Tracks VRAM allocation and device utilization for AI stream processors.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuDevice {
    pub id: String,
    pub name: String,
    pub total_memory_mb: u32,
    pub utilized_memory_mb: u32,
    pub compute_utilization: f32, // 0.0 to 1.0
}

pub struct GpuManager {
    devices: Vec<GpuDevice>,
}

impl GpuManager {
    pub fn new() -> Self {
        info!("[GPU] Initializing Edge Hardware Acceleration...");

        // Attempt to detect actual hardware
        let mut devices = Self::detect_nvidia();
        if devices.is_empty() {
            devices = Self::detect_apple_silicon();
        }

        // Fallback if no dedicated hardware detected
        if devices.is_empty() {
            warn!(
                "[GPU] No hardware accelerators found. AI models will fall back to CPU allocation constraints."
            );
            // We expose an honest "CPU Fallback" device to track memory allocation uniformly.
            // Estimate safe memory bounded by some default (e.g. 4096 MB)
            devices.push(GpuDevice {
                id: "cpu-fallback".to_string(),
                name: "CPU Software Fallback".to_string(),
                total_memory_mb: 4096,
                utilized_memory_mb: 0,
                compute_utilization: 0.0,
            });
        }

        info!(
            "[GPU] Discovered {} accelerator(s) / memory endpoints.",
            devices.len()
        );

        Self { devices }
    }

    #[cfg(test)]
    pub fn new_for_tests() -> Self {
        Self {
            devices: vec![GpuDevice {
                id: "cpu-fallback".to_string(),
                name: "CPU Software Fallback".to_string(),
                total_memory_mb: 4096,
                utilized_memory_mb: 0,
                compute_utilization: 0.0,
            }],
        }
    }

    /// Try to allocate VRAM for an AI model (e.g., Whisper, YOLO)
    pub fn allocate_memory(&mut self, required_mb: u32) -> Result<String, &'static str> {
        for dev in &mut self.devices {
            if dev.total_memory_mb >= dev.utilized_memory_mb + required_mb {
                dev.utilized_memory_mb += required_mb;
                info!("[GPU] Allocated {} MB on device '{}'.", required_mb, dev.id);
                return Ok(dev.id.clone());
            }
        }

        warn!(
            "[GPU] Out of Memory! Could not allocate {} MB.",
            required_mb
        );
        Err("OOM: No hardware accelerator has enough capacity.")
    }

    pub fn free_memory(&mut self, gpu_id: &str, amount_mb: u32) {
        if let Some(dev) = self.devices.iter_mut().find(|d| d.id == gpu_id) {
            dev.utilized_memory_mb = dev.utilized_memory_mb.saturating_sub(amount_mb);
            info!("[GPU] Freed {} MB on device '{}'.", amount_mb, gpu_id);
        }
    }

    pub fn get_devices(&self) -> Vec<GpuDevice> {
        self.devices.clone()
    }

    /// Run nvidia-smi to detect installed NVIDIA GPUs
    fn detect_nvidia() -> Vec<GpuDevice> {
        let mut gpus = Vec::new();
        let out = Command::new("nvidia-smi")
            .args(&[
                "--query-gpu=index,name,memory.total",
                "--format=csv,noheader,nounits",
            ])
            .output();

        if let Ok(output) = out {
            let stdout = String::from_utf8_lossy(&output.stdout);
            for line in stdout.lines() {
                let parts: Vec<&str> = line.split(',').collect();
                if parts.len() >= 3 {
                    debug!("Found NVIDIA GPU: {}", line);
                    let id = format!("nvidia-{}", parts[0].trim());
                    let name = parts[1].trim().to_string();
                    let mem_str = parts[2].trim();
                    if let Ok(total_memory_mb) = mem_str.parse::<u32>() {
                        gpus.push(GpuDevice {
                            id,
                            name,
                            total_memory_mb,
                            utilized_memory_mb: 0,
                            compute_utilization: 0.0,
                        });
                    }
                }
            }
        }
        gpus
    }

    /// Check for Apple Silicon / Metal GPUs on macOS
    fn detect_apple_silicon() -> Vec<GpuDevice> {
        let mut gpus = Vec::new();
        // Uses `system_profiler SPDisplaysDataType` to grab the chipset and VRAM amount heuristically
        let out = Command::new("system_profiler")
            .arg("SPDisplaysDataType")
            .output();

        if let Ok(output) = out {
            let stdout = String::from_utf8_lossy(&output.stdout);
            let mut name = String::new();
            // On Apple Silicon, unified memory isn't strictly given in SPDisplaysDataType usually,
            // but we can parse the chipset model and assign a nominal memory limit (e.g., 8192 MB)
            for line in stdout.lines() {
                if line.contains("Chipset Model:") {
                    name = line.split(':').nth(1).unwrap_or("").trim().to_string();
                }
            }
            if !name.is_empty() {
                gpus.push(GpuDevice {
                    id: "apple-metal-0".to_string(),
                    name,
                    total_memory_mb: 8192, // Heuristic limit for Unified Memory assigned to Metal
                    utilized_memory_mb: 0,
                    compute_utilization: 0.0,
                });
            }
        }
        gpus
    }
}

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

    #[test]
    fn test_gpu_manager_fallback() {
        let mut manager = GpuManager::new(); // will hit CPU fallback or system GPU
        assert!(!manager.get_devices().is_empty());

        let dev_id = manager
            .allocate_memory(1024)
            .expect("Should allocate 1024MB");
        let dev = manager
            .get_devices()
            .into_iter()
            .find(|d| d.id == dev_id)
            .unwrap();
        assert_eq!(dev.utilized_memory_mb, 1024);

        manager.free_memory(&dev_id, 512);
        let dev = manager
            .get_devices()
            .into_iter()
            .find(|d| d.id == dev_id)
            .unwrap();
        assert_eq!(dev.utilized_memory_mb, 512);
    }

    #[test]
    fn test_gpu_manager_oom() {
        let mut manager = GpuManager::new();
        // Exhaust the fallback/metal GPU memory
        let huge_req = 999999;
        assert!(manager.allocate_memory(huge_req).is_err());
    }
}
