//! CUDA Initialization and Device Detection Tests
//! 
//! Tests CUDA device availability, initialization, and basic functionality

use super::utils::*;
use candle_core::{Device, DType, Tensor};
use log::{info, warn};

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_cuda_device_detection() {
        init_gpu_test_env();
        info!("🔍 Testing CUDA device detection");
        
        // Test CUDA availability detection
        let cuda_available = cuda_available();
        info!("CUDA available: {}", cuda_available);
        
        if cuda_available {
            info!("✅ CUDA device detected successfully");
        } else {
            warn!("⚠️  CUDA device not available - running in CPU-only mode");
        }
        
        // This test always passes, but logs the detection result
        assert!(true, "CUDA detection test completed");
    }
    
    #[test]
    fn test_candle_cuda_device_creation() {
        init_gpu_test_env();
        info!("🔧 Testing Candle CUDA device creation");
        
        match Device::cuda_if_available(0) {
            Ok(device) => {
                info!("✅ Candle CUDA device created: {:?}", device);
                assert!(device.is_cuda(), "Device should be CUDA");
                
                // Test basic tensor creation on CUDA device
                let tensor = Tensor::zeros((2, 2), DType::F32, &device)
                    .expect("Should create tensor on CUDA device");
                
                assert_eq!(tensor.device(), &device);
                assert_eq!(tensor.shape().dims(), &[2, 2]);
                info!("✅ Basic tensor creation on CUDA device successful");
            }
            Err(e) => {
                warn!("⚠️  Candle CUDA device creation failed: {}", e);
                info!("Falling back to CPU device for testing");
                
                let cpu_device = Device::Cpu;
                let tensor = Tensor::zeros((2, 2), DType::F32, &cpu_device)
                    .expect("Should create tensor on CPU device");
                
                assert!(!tensor.device().is_cuda(), "Fallback device should be CPU");
            }
        }
    }
    
    #[test] 
    fn test_cuda_device_properties() {
        init_gpu_test_env();
        
        if !cuda_available() {
            info!("⚠️  Skipping CUDA device properties test - CUDA not available");
            return;
        }
        
        info!("📊 Testing CUDA device properties");
        
        #[cfg(feature = "cuda")]
        {
            match cudarc::driver::CudaDevice::new(0) {
                Ok(device) => {
                    info!("✅ CUDA device initialized successfully");
                    
                    // Test device properties
                    let device_name = device.name().unwrap_or("Unknown".to_string());
                    let total_memory = device.total_memory().unwrap_or(0);
                    
                    info!("Device name: {}", device_name);
                    info!("Total memory: {} bytes ({:.2} GB)", 
                          total_memory, total_memory as f64 / (1024.0 * 1024.0 * 1024.0));
                    
                    // Basic validation
                    assert!(!device_name.is_empty(), "Device should have a name");
                    assert!(total_memory > 0, "Device should have memory");
                    
                    info!("✅ CUDA device properties validation passed");
                }
                Err(e) => {
                    warn!("⚠️  CUDA device initialization failed: {}", e);
                }
            }
        }
        
        #[cfg(not(feature = "cuda"))]
        {
            info!("⚠️  CUDA feature not enabled - skipping device properties test");
        }
    }
    
    #[test]
    fn test_multiple_cuda_devices() {
        init_gpu_test_env();
        info!("🔢 Testing multiple CUDA device detection");
        
        let mut cuda_devices = Vec::new();
        
        // Try to detect up to 8 CUDA devices
        for device_id in 0..8 {
            match Device::cuda_if_available(device_id) {
                Ok(device) => {
                    info!("✅ CUDA device {} available: {:?}", device_id, device);
                    cuda_devices.push((device_id, device));
                }
                Err(_) => {
                    // Stop at first unavailable device
                    break;
                }
            }
        }
        
        info!("Found {} CUDA device(s)", cuda_devices.len());
        
        if cuda_devices.is_empty() {
            warn!("⚠️  No CUDA devices found");
        } else {
            info!("✅ CUDA device enumeration successful");
            
            // Test basic operations on each device
            for (device_id, device) in cuda_devices {
                let test_tensor = Tensor::ones((10, 10), DType::F32, &device)
                    .expect(&format!("Should create tensor on device {}", device_id));
                
                assert!(test_tensor.device().is_cuda());
                info!("✅ Device {} tensor creation test passed", device_id);
            }
        }
    }
    
    #[test]
    fn test_cuda_error_handling() {
        init_gpu_test_env();
        info!("🚨 Testing CUDA error handling");
        
        // Test invalid device ID
        match Device::cuda_if_available(99) {
            Ok(_) => {
                warn!("⚠️  Unexpected: Device 99 should not be available");
            }
            Err(e) => {
                info!("✅ Expected error for invalid device ID: {}", e);
            }
        }
        
        // Test graceful fallback
        let device = get_test_device();
        info!("✅ Graceful device fallback: {:?}", device);
        
        // Ensure we can always create tensors on the fallback device
        let tensor = Tensor::zeros((5, 5), DType::F32, &device)
            .expect("Should create tensor on fallback device");
        
        assert_eq!(tensor.shape().dims(), &[5, 5]);
        info!("✅ Error handling and fallback test passed");
    }
}