//! Windows-specific CUDA support.
//!
//! This module provides Windows platform detection and CUDA library discovery.

use std::path::PathBuf;

/// Windows CUDA library paths and detection.
#[derive(Debug, Clone)]
pub struct WindowsCudaConfig {
    /// CUDA installation path (e.g., C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.x)
    pub cuda_path: Option<PathBuf>,
    /// cuDNN installation path
    pub cudnn_path: Option<PathBuf>,
    /// Detected CUDA version
    pub cuda_version: Option<String>,
    /// Available CUDA DLLs
    pub available_dlls: Vec<String>,
}

impl WindowsCudaConfig {
    /// Detect CUDA installation on Windows.
    #[cfg(target_os = "windows")]
    pub fn detect() -> Self {
        let cuda_path = Self::find_cuda_path();
        let cudnn_path = Self::find_cudnn_path();
        let cuda_version = Self::detect_cuda_version(&cuda_path);
        let available_dlls = Self::find_cuda_dlls(&cuda_path);

        Self {
            cuda_path,
            cudnn_path,
            cuda_version,
            available_dlls,
        }
    }

    /// Detect CUDA installation (stub for non-Windows).
    #[cfg(not(target_os = "windows"))]
    pub fn detect() -> Self {
        Self {
            cuda_path: None,
            cudnn_path: None,
            cuda_version: None,
            available_dlls: Vec::new(),
        }
    }

    /// Find CUDA installation path from environment or default locations.
    #[cfg(target_os = "windows")]
    fn find_cuda_path() -> Option<PathBuf> {
        // Check CUDA_PATH environment variable first
        if let Ok(path) = std::env::var("CUDA_PATH") {
            let p = PathBuf::from(&path);
            if p.exists() {
                return Some(p);
            }
        }

        // Check CUDA_HOME (alternate environment variable)
        if let Ok(path) = std::env::var("CUDA_HOME") {
            let p = PathBuf::from(&path);
            if p.exists() {
                return Some(p);
            }
        }

        // Check default installation locations
        let default_paths = [
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.5",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.3",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.2",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.1",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.0",
            r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8",
        ];

        for path_str in &default_paths {
            let path = PathBuf::from(path_str);
            if path.exists() {
                return Some(path);
            }
        }

        None
    }

    #[cfg(not(target_os = "windows"))]
    fn find_cuda_path() -> Option<PathBuf> {
        None
    }

    /// Find cuDNN installation path.
    #[cfg(target_os = "windows")]
    fn find_cudnn_path() -> Option<PathBuf> {
        // cuDNN is usually installed alongside CUDA or in CUDNN_PATH
        if let Ok(path) = std::env::var("CUDNN_PATH") {
            let p = PathBuf::from(&path);
            if p.exists() {
                return Some(p);
            }
        }

        // Check if cuDNN is in CUDA path
        if let Some(cuda_path) = Self::find_cuda_path() {
            let cudnn_dll = cuda_path.join("bin").join("cudnn64_9.dll");
            if cudnn_dll.exists() {
                return Some(cuda_path);
            }
            // Check older versions
            let cudnn_dll = cuda_path.join("bin").join("cudnn64_8.dll");
            if cudnn_dll.exists() {
                return Some(cuda_path);
            }
        }

        None
    }

    #[cfg(not(target_os = "windows"))]
    fn find_cudnn_path() -> Option<PathBuf> {
        None
    }

    /// Detect CUDA version from installation.
    #[cfg(target_os = "windows")]
    fn detect_cuda_version(cuda_path: &Option<PathBuf>) -> Option<String> {
        if let Some(path) = cuda_path {
            // Try to read version from version.txt
            let version_file = path.join("version.txt");
            if let Ok(content) = std::fs::read_to_string(&version_file) {
                // Parse "CUDA Version 12.x.y" format
                if let Some(version) = content.lines().next() {
                    if let Some(v) = version.strip_prefix("CUDA Version ") {
                        return Some(v.trim().to_string());
                    }
                }
            }

            // Fallback: extract from path
            if let Some(name) = path.file_name() {
                if let Some(name_str) = name.to_str() {
                    if let Some(version) = name_str.strip_prefix("v") {
                        return Some(version.to_string());
                    }
                }
            }
        }
        None
    }

    #[cfg(not(target_os = "windows"))]
    fn detect_cuda_version(_cuda_path: &Option<PathBuf>) -> Option<String> {
        None
    }

    /// Find available CUDA DLLs.
    #[cfg(target_os = "windows")]
    fn find_cuda_dlls(cuda_path: &Option<PathBuf>) -> Vec<String> {
        let mut dlls = Vec::new();

        if let Some(path) = cuda_path {
            let bin_path = path.join("bin");
            if bin_path.exists() {
                if let Ok(entries) = std::fs::read_dir(&bin_path) {
                    for entry in entries.flatten() {
                        if let Some(name) = entry.file_name().to_str() {
                            if name.ends_with(".dll") {
                                dlls.push(name.to_string());
                            }
                        }
                    }
                }
            }
        }

        // Also check System32 for driver DLLs
        let system32 = PathBuf::from(r"C:\Windows\System32");
        let driver_dlls = ["nvcuda.dll", "nvml.dll", "nvapi64.dll"];
        for dll in &driver_dlls {
            if system32.join(dll).exists() {
                dlls.push((*dll).to_string());
            }
        }

        dlls
    }

    #[cfg(not(target_os = "windows"))]
    fn find_cuda_dlls(_cuda_path: &Option<PathBuf>) -> Vec<String> {
        Vec::new()
    }

    /// Check if CUDA is available on Windows.
    pub fn is_available(&self) -> bool {
        self.cuda_path.is_some() && self.available_dlls.contains(&"nvcuda.dll".to_string())
    }

    /// Get the library search paths for CUDA.
    pub fn library_paths(&self) -> Vec<PathBuf> {
        let mut paths = Vec::new();

        if let Some(cuda_path) = &self.cuda_path {
            paths.push(cuda_path.join("bin"));
            paths.push(cuda_path.join("lib").join("x64"));
        }

        if let Some(cudnn_path) = &self.cudnn_path {
            if cudnn_path != self.cuda_path.as_ref().unwrap_or(&PathBuf::new()) {
                paths.push(cudnn_path.join("bin"));
                paths.push(cudnn_path.join("lib").join("x64"));
            }
        }

        paths
    }

    /// Get required DLLs for CUDA operations.
    pub fn required_dlls() -> &'static [&'static str] {
        &[
            "nvcuda.dll",        // CUDA driver
            "nvrtc64_120_0.dll", // NVRTC (runtime compilation)
            "cublas64_12.dll",   // cuBLAS
            "cublasLt64_12.dll", // cuBLAS Lt
        ]
    }

    /// Get optional DLLs for enhanced functionality.
    pub fn optional_dlls() -> &'static [&'static str] {
        &[
            "cudnn64_9.dll",     // cuDNN
            "cusparse64_12.dll", // cuSPARSE
            "cusolver64_11.dll", // cuSOLVER
            "cufft64_11.dll",    // cuFFT
            "curand64_10.dll",   // cuRAND
        ]
    }

    /// Check if all required DLLs are available.
    pub fn has_required_dlls(&self) -> bool {
        // Only nvcuda.dll is strictly required (driver)
        self.available_dlls.contains(&"nvcuda.dll".to_string())
    }
}

/// Windows GPU information from DXGI.
#[derive(Debug, Clone)]
pub struct WindowsGpuInfo {
    /// GPU name
    pub name: String,
    /// Vendor ID
    pub vendor_id: u32,
    /// Device ID
    pub device_id: u32,
    /// Dedicated video memory in bytes
    pub dedicated_memory: usize,
    /// Is NVIDIA GPU
    pub is_nvidia: bool,
    /// Is AMD GPU
    pub is_amd: bool,
    /// Is Intel GPU
    pub is_intel: bool,
}

impl WindowsGpuInfo {
    /// Enumerate all GPUs on Windows using DXGI.
    #[cfg(target_os = "windows")]
    pub fn enumerate_gpus() -> Vec<Self> {
        // This would use DXGI to enumerate adapters
        // For now, return empty (actual implementation requires windows-sys)
        Vec::new()
    }

    #[cfg(not(target_os = "windows"))]
    pub fn enumerate_gpus() -> Vec<Self> {
        Vec::new()
    }

    /// NVIDIA vendor ID
    pub const NVIDIA_VENDOR_ID: u32 = 0x10DE;
    /// AMD vendor ID
    pub const AMD_VENDOR_ID: u32 = 0x1002;
    /// Intel vendor ID
    pub const INTEL_VENDOR_ID: u32 = 0x8086;
}

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

    #[test]
    fn test_cuda_config_detection() {
        let config = WindowsCudaConfig::detect();
        println!("CUDA path: {:?}", config.cuda_path);
        println!("CUDA version: {:?}", config.cuda_version);
        println!("Available DLLs: {:?}", config.available_dlls);

        // On non-Windows, these should be None/empty
        #[cfg(not(target_os = "windows"))]
        {
            assert!(config.cuda_path.is_none());
            assert!(config.available_dlls.is_empty());
        }
    }

    #[test]
    fn test_required_dlls() {
        let required = WindowsCudaConfig::required_dlls();
        assert!(required.contains(&"nvcuda.dll"));
    }

    #[test]
    fn test_gpu_enumeration() {
        let gpus = WindowsGpuInfo::enumerate_gpus();
        println!("Found {} GPUs", gpus.len());
        for gpu in &gpus {
            println!(
                "  - {} (NVIDIA: {}, AMD: {}, Intel: {})",
                gpu.name, gpu.is_nvidia, gpu.is_amd, gpu.is_intel
            );
        }
    }
}