//! Production GPU Integration Tests
//! 
//! Tests GPU integration in production scenarios including environment handling,
//! configuration management, and error recovery

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

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_production_gpu_initialization() {
        init_gpu_test_env();
        info!("🏭 Testing production GPU initialization");
        
        // Simulate production environment variables
        env::set_var("FOXHUNT_GPU_ENABLED", "true");
        env::set_var("FOXHUNT_CUDA_DEVICE_ID", "0");
        env::set_var("FOXHUNT_GPU_MEMORY_LIMIT", "8192"); // 8GB in MB
        
        // Test production-style device initialization
        let gpu_enabled = env::var("FOXHUNT_GPU_ENABLED")
            .unwrap_or_else(|_| "false".to_string())
            .parse::<bool>()
            .unwrap_or(false);
        
        let device_id = env::var("FOXHUNT_CUDA_DEVICE_ID")
            .unwrap_or_else(|_| "0".to_string())
            .parse::<usize>()
            .unwrap_or(0);
        
        let memory_limit_mb = env::var("FOXHUNT_GPU_MEMORY_LIMIT")
            .unwrap_or_else(|_| "4096".to_string())
            .parse::<usize>()
            .unwrap_or(4096);
        
        info!("Production config - GPU enabled: {}, Device ID: {}, Memory limit: {}MB", 
              gpu_enabled, device_id, memory_limit_mb);
        
        if gpu_enabled && cuda_available() {
            match Device::cuda_if_available(device_id) {
                Ok(device) => {
                    info!("✅ Production GPU device initialized: {:?}", device);
                    
                    // Test production-style memory allocation within limits
                    let max_tensor_size = (memory_limit_mb * 1024 * 1024) / (4 * 4); // Rough estimate for f32 matrix
                    let test_size = (max_tensor_size as f64).sqrt() as usize / 10; // Conservative size
                    
                    match Tensor::zeros((test_size, test_size), DType::F32, &device) {
                        Ok(_tensor) => {
                            info!("✅ Production memory allocation successful: {}x{}", test_size, test_size);
                        }
                        Err(e) => {
                            warn!("❌ Production memory allocation failed: {}", e);
                        }
                    }
                }
                Err(e) => {
                    warn!("❌ Production GPU initialization failed: {}", e);
                    info!("Falling back to CPU for production");
                }
            }
        } else {
            info!("💻 Production running in CPU-only mode");
            let cpu_device = Device::Cpu;
            
            // Test CPU production path
            match Tensor::zeros((1000, 1000), DType::F32, &cpu_device) {
                Ok(_tensor) => {
                    info!("✅ Production CPU fallback successful");
                }
                Err(e) => {
                    warn!("❌ Production CPU fallback failed: {}", e);
                }
            }
        }
        
        // Clean up environment variables
        env::remove_var("FOXHUNT_GPU_ENABLED");
        env::remove_var("FOXHUNT_CUDA_DEVICE_ID");
        env::remove_var("FOXHUNT_GPU_MEMORY_LIMIT");
        
        info!("✅ Production GPU initialization test completed");
    }
    
    #[test]
    fn test_cuda_visible_devices_handling() {
        init_gpu_test_env();
        info!("👁️  Testing CUDA_VISIBLE_DEVICES handling");
        
        // Save original CUDA_VISIBLE_DEVICES
        let original_cuda_devices = env::var("CUDA_VISIBLE_DEVICES").ok();
        
        // Test with restricted devices
        env::set_var("CUDA_VISIBLE_DEVICES", "0");
        
        let device = get_test_device();
        info!("Device with CUDA_VISIBLE_DEVICES=0: {:?}", device);
        
        if device.is_cuda() {
            info!("✅ CUDA device accessible with restricted visibility");
        } else {
            info!("💻 Using CPU device (CUDA may not be available)");
        }
        
        // Test with no visible devices
        env::set_var("CUDA_VISIBLE_DEVICES", "");
        
        let restricted_device = Device::cuda_if_available(0);
        match restricted_device {
            Ok(_) => {
                warn!("⚠️  CUDA device available when none should be visible");
            }
            Err(_) => {
                info!("✅ CUDA correctly restricted when CUDA_VISIBLE_DEVICES is empty");
            }
        }
        
        // Restore original setting
        if let Some(original) = original_cuda_devices {
            env::set_var("CUDA_VISIBLE_DEVICES", original);
        } else {
            env::remove_var("CUDA_VISIBLE_DEVICES");
        }
        
        info!("✅ CUDA_VISIBLE_DEVICES handling test completed");
    }
    
    #[test]
    fn test_gpu_memory_limits_production() {
        init_gpu_test_env();
        info!("📏 Testing GPU memory limits in production scenarios");
        
        if !cuda_available() {
            info!("⚠️  Skipping GPU memory limits test - CUDA not available");
            return;
        }
        
        let device = get_test_device();
        
        // Test progressive memory allocation to find limits
        let mut allocated_tensors = Vec::new();
        let chunk_size = 1024 * 1024; // 1M elements = 4MB
        let max_chunks = 100; // Max 400MB total
        
        info!("Testing progressive memory allocation...");
        
        for i in 0..max_chunks {
            match Tensor::zeros((chunk_size,), DType::F32, &device) {
                Ok(tensor) => {
                    allocated_tensors.push(tensor);
                    let total_mb = (i + 1) * 4; // Each chunk is 4MB
                    
                    if i % 10 == 0 {
                        info!("✅ Allocated {}MB in {} chunks", total_mb, i + 1);
                    }
                }
                Err(e) => {
                    let failed_at_mb = (i + 1) * 4;
                    warn!("❌ Memory allocation failed at {}MB: {}", failed_at_mb, e);
                    break;
                }
            }
        }
        
        let total_allocated = allocated_tensors.len() * 4;
        info!("Total GPU memory allocated: {}MB in {} chunks", total_allocated, allocated_tensors.len());
        
        // Test memory cleanup
        let chunks_to_free = allocated_tensors.len() / 2;
        allocated_tensors.truncate(chunks_to_free);
        
        info!("Freed {}MB, {} chunks remaining", (allocated_tensors.len()) * 4, allocated_tensors.len());
        
        // Test that we can allocate more after freeing
        match Tensor::zeros((chunk_size,), DType::F32, &device) {
            Ok(_tensor) => {
                info!("✅ Memory reallocation successful after cleanup");
            }
            Err(e) => {
                warn!("❌ Memory reallocation failed after cleanup: {}", e);
            }
        }
        
        info!("✅ GPU memory limits production test completed");
    }
    
    #[tokio::test]
    async fn test_production_ml_inference_pipeline() {
        init_gpu_test_env();
        info!("🧠 Testing production ML inference pipeline with GPU");
        
        let device = get_test_device();
        
        // Simulate production ML inference workload
        let batch_sizes = vec![1, 8, 16, 32, 64];
        let feature_size = 256;
        
        for batch_size in batch_sizes {
            info!("Testing batch size: {}", batch_size);
            
            // Create realistic feature data
            let features = match Tensor::randn(0.0, 1.0, (batch_size, feature_size), &device) {
                Ok(tensor) => tensor,
                Err(e) => {
                    warn!("❌ Failed to create features for batch size {}: {}", batch_size, e);
                    continue;
                }
            };
            
            // Simulate neural network layers
            let layer_sizes = vec![feature_size, 512, 256, 128, 64, 1];
            
            let start_time = Instant::now();
            let mut x = features;
            
            // Forward pass through layers
            for i in 0..layer_sizes.len() - 1 {
                let weight = match Tensor::randn(0.0, 0.01, (layer_sizes[i], layer_sizes[i + 1]), &device) {
                    Ok(tensor) => tensor,
                    Err(e) => {
                        warn!("❌ Failed to create weight matrix: {}", e);
                        break;
                    }
                };
                
                x = match x.matmul(&weight) {
                    Ok(result) => result,
                    Err(e) => {
                        warn!("❌ Matrix multiplication failed: {}", e);
                        break;
                    }
                };
                
                // Apply activation (except for output layer)
                if i < layer_sizes.len() - 2 {
                    x = match x.relu() {
                        Ok(result) => result,
                        Err(e) => {
                            warn!("❌ ReLU activation failed: {}", e);
                            break;
                        }
                    };
                }
            }
            
            let inference_time = start_time.elapsed();
            
            // Validate output shape
            if x.shape().dims() == &[batch_size, 1] {
                let latency_us = inference_time.as_micros();
                let latency_per_sample_us = latency_us / batch_size as u128;
                
                info!("✅ Batch {} inference: {:?} total, {}μs per sample", 
                      batch_size, inference_time, latency_per_sample_us);
                
                // Check HFT requirements
                if latency_per_sample_us < 50 {
                    info!("   🚀 Meets sub-50μs HFT requirement!");
                } else if latency_per_sample_us < 100 {
                    info!("   ✅ Good latency for production");
                } else {
                    warn!("   ⚠️  High latency - may need optimization");
                }
            } else {
                warn!("❌ Unexpected output shape for batch size {}: {:?}", batch_size, x.shape());
            }
        }
        
        info!("✅ Production ML inference pipeline test completed");
    }
    
    #[test]
    fn test_gpu_error_recovery_production() {
        init_gpu_test_env();
        info!("🚨 Testing GPU error recovery in production scenarios");
        
        let device = get_test_device();
        
        // Test recovery from out-of-memory errors
        info!("Testing OOM recovery...");
        
        // Try to allocate a very large tensor (should fail)
        let large_size = 100_000;
        let oom_result = Tensor::zeros((large_size, large_size), DType::F32, &device);
        
        match oom_result {
            Ok(_) => {
                warn!("⚠️  Large tensor allocation succeeded (unexpected)");
            }
            Err(e) => {
                info!("✅ Expected OOM error: {}", e);
                
                // Test that we can still allocate normal-sized tensors after OOM
                match Tensor::zeros((100, 100), DType::F32, &device) {
                    Ok(_) => {
                        info!("✅ Recovery successful - can allocate normal tensors after OOM");
                    }
                    Err(recovery_error) => {
                        warn!("❌ Recovery failed: {}", recovery_error);
                    }
                }
            }
        }
        
        // Test recovery from invalid operations
        info!("Testing invalid operation recovery...");
        
        let tensor_a = Tensor::ones((10, 5), DType::F32, &device).expect("Should create tensor A");
        let tensor_b = Tensor::ones((3, 7), DType::F32, &device).expect("Should create tensor B");
        
        // Try invalid matrix multiplication (incompatible dimensions)
        let invalid_result = tensor_a.matmul(&tensor_b);
        
        match invalid_result {
            Ok(_) => {
                warn!("⚠️  Invalid matrix multiplication succeeded (unexpected)");
            }
            Err(e) => {
                info!("✅ Expected dimension error: {}", e);
                
                // Test that we can still do valid operations
                let tensor_c = Tensor::ones((5, 7), DType::F32, &device).expect("Should create tensor C");
                match tensor_a.matmul(&tensor_c) {
                    Ok(_) => {
                        info!("✅ Recovery successful - valid operations work after error");
                    }
                    Err(recovery_error) => {
                        warn!("❌ Recovery failed: {}", recovery_error);
                    }
                }
            }
        }
        
        info!("✅ GPU error recovery production test completed");
    }
    
    #[test]
    fn test_production_gpu_monitoring() {
        init_gpu_test_env();
        info!("📊 Testing production GPU monitoring and telemetry");
        
        if !cuda_available() {
            info!("⚠️  Testing CPU monitoring (CUDA not available)");
            
            // Test CPU monitoring fallback
            let cpu_device = Device::Cpu;
            let start_time = Instant::now();
            
            let _tensor = Tensor::zeros((1000, 1000), DType::F32, &cpu_device)
                .expect("Should create CPU tensor");
            
            let cpu_time = start_time.elapsed();
            
            info!("CPU operation monitoring:");
            info!("  Operation time: {:?}", cpu_time);
            info!("  Device type: CPU");
            info!("  Memory usage: Estimated 4MB");
            
            return;
        }
        
        let device = get_test_device();
        
        // Test GPU monitoring
        #[cfg(feature = "cuda")]
        {
            match cudarc::driver::CudaDevice::new(0) {
                Ok(cuda_device) => {
                    if let Ok(total_memory) = cuda_device.total_memory() {
                        info!("GPU monitoring data:");
                        info!("  Total GPU memory: {} bytes ({:.2} GB)", 
                              total_memory, total_memory as f64 / (1024.0 * 1024.0 * 1024.0));
                        
                        // Test memory usage monitoring during allocation
                        let start_time = Instant::now();
                        let tensor_result = Tensor::zeros((2000, 2000), DType::F32, &device);
                        let allocation_time = start_time.elapsed();
                        
                        match tensor_result {
                            Ok(_tensor) => {
                                let estimated_usage = 2000 * 2000 * 4; // 4 bytes per f32
                                
                                info!("  Allocation time: {:?}", allocation_time);
                                info!("  Estimated memory usage: {} bytes ({:.2} MB)", 
                                      estimated_usage, estimated_usage as f64 / (1024.0 * 1024.0));
                                info!("  ✅ GPU monitoring data collected successfully");
                            }
                            Err(e) => {
                                warn!("  ❌ GPU allocation failed during monitoring: {}", e);
                            }
                        }
                    }
                }
                Err(e) => {
                    warn!("❌ Failed to initialize CUDA device for monitoring: {}", e);
                }
            }
        }
        
        #[cfg(not(feature = "cuda"))]
        {
            info!("GPU monitoring fallback - CUDA feature not enabled");
        }
        
        info!("✅ Production GPU monitoring test completed");
    }
    
    #[test]
    fn test_production_configuration_validation() {
        init_gpu_test_env();
        info!("⚙️  Testing production GPU configuration validation");
        
        // Test various production configurations
        let test_configs = vec![
            ("development", false, 2048),  // Development: GPU disabled, low memory
            ("staging", true, 4096),       // Staging: GPU enabled, medium memory
            ("production", true, 8192),    // Production: GPU enabled, high memory
        ];
        
        for (env_name, gpu_enabled, memory_limit_mb) in test_configs {
            info!("Testing {} configuration", env_name);
            
            // Set environment
            env::set_var("FOXHUNT_ENVIRONMENT", env_name);
            env::set_var("FOXHUNT_GPU_ENABLED", gpu_enabled.to_string());
            env::set_var("FOXHUNT_GPU_MEMORY_LIMIT", memory_limit_mb.to_string());
            
            // Validate configuration
            let actual_gpu_enabled = gpu_enabled && cuda_available();
            let device = if actual_gpu_enabled {
                Device::cuda_if_available(0).unwrap_or(Device::Cpu)
            } else {
                Device::Cpu
            };
            
            info!("  Environment: {}", env_name);
            info!("  GPU enabled: {} (available: {})", gpu_enabled, cuda_available());
            info!("  Memory limit: {}MB", memory_limit_mb);
            info!("  Actual device: {:?}", device);
            
            // Test memory allocation within limits
            let test_tensor_size = (memory_limit_mb * 256) / 4; // Conservative test size
            let side_length = (test_tensor_size as f64).sqrt() as usize;
            
            match Tensor::zeros((side_length, side_length), DType::F32, &device) {
                Ok(_tensor) => {
                    info!("  ✅ Memory allocation within limits successful");
                }
                Err(e) => {
                    warn!("  ❌ Memory allocation failed: {}", e);
                }
            }
            
            // Clean up
            env::remove_var("FOXHUNT_ENVIRONMENT");
            env::remove_var("FOXHUNT_GPU_ENABLED");
            env::remove_var("FOXHUNT_GPU_MEMORY_LIMIT");
        }
        
        info!("✅ Production configuration validation test completed");
    }
}