/*!
* Integration tests for RustyTorch++ runtime stack
*
* These tests exercise the complete integration between allocator, device,
* streams, scheduler, and kernel launch systems using real GPU resources
* when available.
*/

use rtx_runtime::*;
use std::sync::Arc;
use std::time::Duration;

/// Environment-based backend selection for testing
fn get_test_backend() -> Backend {
    match std::env::var("RTX_BACKEND").as_deref() {
        Ok("cuda") => Backend::Cuda,
        Ok("rocm") => Backend::ROCm,
        Ok("metal") => Backend::Metal,
        _ => Backend::CPU,
    }
}

/// Get device count from environment or default to 1
fn get_device_count() -> usize {
    std::env::var("RTX_DEVICE_COUNT")
    .ok()
    .and_then(|s| s.parse().ok())
    .unwrap_or(1)
}

#[tokio::test]
async fn test_full_stack_initialization() -> Result<()> {
    // Test complete stack initialization with real GPU if available
    let backend = get_test_backend();
    let device_count = get_device_count();

    // Create allocator with production settings
    let allocator_config = AllocatorConfig {
        arena_size_bytes: 256 * 1024 * 1024, // 256MB
        max_fragmentation_percent: 15,
        enable_statistics: true,
    };
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    // Initialize devices
    let mut devices = Vec::new();
    for device_id in 0..device_count {
        let device = Device::new(backend, device_id, allocator.clone())?;
        assert_eq!(device.backend(), backend);
        assert_eq!(device.device_id(), device_id);
        devices.push(device);
    }

    // Create scheduler with all devices
    let scheduler = Scheduler::new(devices)?;
    assert_eq!(scheduler.device_count(), device_count);

    println!("Successfully initialized full stack with {} {} devices",
    device_count,
    match backend {
        Backend::CUDA => "CUDA",
        Backend::ROCm => "ROCm",
        Backend::Metal => "Metal",
        Backend::CPU => "CPU",
    }
    );

    Ok(())
}

#[tokio::test]
async fn test_memory_lifecycle_integration() -> Result<()> {
    let backend = get_test_backend();
    let allocator_config = AllocatorConfig {
        arena_size_bytes: 128 * 1024 * 1024, // 128MB
        max_fragmentation_percent: 20,
        enable_statistics: true,
    };
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    let device = Device::new(backend, 0, allocator.clone())?;
    let scheduler = Scheduler::new(vec![device])?;

    // Test various allocation sizes across multiple size classes
    let test_sizes = vec![
    1024,        // 1KB
    64 * 1024,   // 64KB
    1024 * 1024, // 1MB
    16 * 1024 * 1024, // 16MB
    ];

    let mut handles = Vec::new();

    // Allocate memory blocks
    for &size in &test_sizes {
        let handle = allocator.allocate(size)?;
        handles.push(handle);
    }

    // Check allocator statistics
    let stats = allocator.statistics();
    assert!(stats.total_allocated_bytes > 0);
    assert_eq!(stats.active_allocations, test_sizes.len());

    // Free half the blocks
    for handle in handles.drain(0..test_sizes.len() / 2) {
        allocator.deallocate(handle)?;
    }

    // Verify statistics updated correctly
    let stats_after = allocator.statistics();
    assert_eq!(stats_after.active_allocations, test_sizes.len() / 2);
    assert!(stats_after.total_freed_bytes > 0);

    // Free remaining blocks
    for handle in handles {
        allocator.deallocate(handle)?;
    }

    let final_stats = allocator.statistics();
    assert_eq!(final_stats.active_allocations, 0);

    Ok(())
}

#[tokio::test]
async fn test_stream_scheduling_integration() -> Result<()> {
    let backend = get_test_backend();
    let allocator_config = AllocatorConfig::default();
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    let device = Device::new(backend, 0, allocator.clone())?;
    let mut scheduler = Scheduler::new(vec![device])?;

    // Create a complex dependency graph with multiple operations
    let op1 = ScheduledOperation::MemoryTransfer {
        src: 0x1000,
        dst: 0x2000,
        size: 1024 * 1024,
        device_id: 0
    };
    let op2 = ScheduledOperation::KernelLaunch {
        kernel_id: 1,
        device_id: 0
    };
    let op3 = ScheduledOperation::MemoryTransfer {
        src: 0x2000,
        dst: 0x3000,
        size: 1024 * 1024,
        device_id: 0
    };
    let op4 = ScheduledOperation::KernelLaunch {
        kernel_id: 2,
        device_id: 0
    };

    // Schedule operations with dependencies: op1 -> op2 -> op3 -> op4
    let task1 = scheduler.schedule(op1, Vec::new()).await?;
    let task2 = scheduler.schedule(op2, vec![task1]).await?;
    let task3 = scheduler.schedule(op3, vec![task2]).await?;
    let task4 = scheduler.schedule(op4, vec![task3]).await?;

    // Wait for all operations to complete
    scheduler.wait_for_completion(task4).await?;

    // Verify scheduler statistics
    let stats = scheduler.statistics();
    assert_eq!(stats.total_scheduled, 4);
    assert_eq!(stats.completed_tasks, 4);
    assert!(stats.average_scheduling_time_ns < 1_000_000); // < 1ms

    Ok(())
}

#[tokio::test]
async fn test_kernel_launch_integration() -> Result<()> {
    let backend = get_test_backend();
    let allocator_config = AllocatorConfig::default();
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    let device = Device::new(backend, 0, allocator.clone())?;
    let mut scheduler = Scheduler::new(vec![device])?;

    // Create a kernel launch system
    let mut kernel_system = KernelLaunchSystem::new();

    // Mock PTX source for a simple vector add kernel
    let ptx_source = r#"
    .version 7.0
    .target sm_120
    .address_size 64

    .visible .entry vector_add(
    .param .u64 vector_add_param_0,
    .param .u64 vector_add_param_1,
    .param .u64 vector_add_param_2,
    .param .u32 vector_add_param_3
    )
    {
    .reg .u32 %tid;
    .reg .u64 %rd<4>;

    mov.u32 %tid, %tid.x;

    ld.param.u64 %rd0, [vector_add_param_0];
    ld.param.u64 %rd1, [vector_add_param_1];
    ld.param.u64 %rd2, [vector_add_param_2];

    ret;
    }
    "#;

    // Load and compile the kernel
    let kernel_id = kernel_system.load_kernel("vector_add", ptx_source)?;

    // Allocate input/output arrays
    let n = 1024;
    let size_bytes = n * std::mem::size_of::<f32>();

    let input_a = allocator.allocate(size_bytes)?;
    let input_b = allocator.allocate(size_bytes)?;
    let output = allocator.allocate(size_bytes)?;

    // Configure kernel launch parameters
    let params = KernelParams {
        grid_dim: (16, 1, 1),      // 16 blocks
        block_dim: (64, 1, 1),    // 64 threads per block
        shared_memory_bytes: 0,
        parameters: vec![
        input_a.device_ptr() as u64,
        input_b.device_ptr() as u64,
        output.device_ptr() as u64,
        n as u32,
        ],
    };

    // Launch kernel through scheduler
    let launch_op = ScheduledOperation::KernelLaunch {
        kernel_id,
        device_id: 0
    };

    let start_time = std::time::Instant::now();
    let task = scheduler.schedule(launch_op, Vec::new()).await?;
    scheduler.wait_for_completion(task).await?;
    let elapsed = start_time.elapsed();

    // Verify kernel execution completed
    println!("Kernel execution completed in {:?}", elapsed);
    assert!(elapsed < Duration::from_millis(100)); // Should complete quickly

    // Check kernel system statistics
    let stats = kernel_system.statistics();
    assert_eq!(stats.kernels_loaded, 1);
    assert_eq!(stats.total_launches, 1);
    assert!(stats.average_launch_time_ns > 0);

    // Clean up allocated memory
    allocator.deallocate(input_a)?;
    allocator.deallocate(input_b)?;
    allocator.deallocate(output)?;

    Ok(())
}

#[tokio::test]
async fn test_multi_device_coordination() -> Result<()> {
    let backend = get_test_backend();
    let device_count = get_device_count().max(1); // Use at least 1 device

    if device_count == 1 {
        println!("Skipping multi-device test - only 1 device available");
        return Ok(());
    }

    let allocator_config = AllocatorConfig::default();
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    // Create multiple devices
    let mut devices = Vec::new();
    for device_id in 0..device_count {
        let device = Device::new(backend, device_id, allocator.clone())?;
        devices.push(device);
    }

    let mut scheduler = Scheduler::new(devices)?;

    // Schedule operations across different devices
    let mut tasks = Vec::new();
    for device_id in 0..device_count {
        let op = ScheduledOperation::MemoryTransfer {
            src: 0x1000 + (device_id as u64 * 0x1000),
            dst: 0x2000 + (device_id as u64 * 0x1000),
            size: 1024 * 1024,
            device_id,
        };

        let task = scheduler.schedule(op, Vec::new()).await?;
        tasks.push(task);
    }

    // Wait for all cross-device operations to complete
    for task in tasks {
        scheduler.wait_for_completion(task).await?;
    }

    let stats = scheduler.statistics();
    assert_eq!(stats.completed_tasks, device_count);

    println!("Successfully coordinated operations across {} devices", device_count);

    Ok(())
}

#[tokio::test]
async fn test_error_recovery_integration() -> Result<()> {
    let backend = get_test_backend();
    let allocator_config = AllocatorConfig {
        arena_size_bytes: 64 * 1024 * 1024, // Smaller arena to trigger OOM
        max_fragmentation_percent: 15,
        enable_statistics: true,
    };
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    let device = Device::new(backend, 0, allocator.clone())?;
    let mut scheduler = Scheduler::new(vec![device])?;

    // Try to allocate more memory than available to test error handling
    let large_size = 128 * 1024 * 1024; // 128MB (larger than 64MB arena)

    let result = allocator.allocate(large_size);
    match result {
        Ok(handle) => {
            // If allocation succeeded, clean up
            allocator.deallocate(handle)?;
            println!("Large allocation succeeded (arena may be larger than expected)");
        }
        Err(Error::OutOfMemory) => {
            println!("Expected OOM error occurred - error recovery working correctly");
        }
        Err(e) => {
            panic!("Unexpected error type: {:?}", e);
        }
    }

    // Verify that smaller allocations still work after OOM
    let small_handle = allocator.allocate(1024)?;
    allocator.deallocate(small_handle)?;

    // Test scheduler with invalid device ID
    let invalid_op = ScheduledOperation::KernelLaunch {
        kernel_id: 999,
        device_id: 999 // Invalid device
    };

    let schedule_result = scheduler.schedule(invalid_op, Vec::new()).await;
    assert!(schedule_result.is_err(), "Expected error for invalid device ID");

    println!("Error recovery integration tests passed");

    Ok(())
}

#[tokio::test]
async fn test_performance_baseline() -> Result<()> {
    let backend = get_test_backend();
    let allocator_config = AllocatorConfig::default();
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    let device = Device::new(backend, 0, allocator.clone())?;
    let mut scheduler = Scheduler::new(vec![device])?;

    // Performance test: rapid allocation/deallocation
    let iterations = 1000;
    let allocation_size = 64 * 1024; // 64KB

    let start_time = std::time::Instant::now();

    for _ in 0..iterations {
        let handle = allocator.allocate(allocation_size)?;
        allocator.deallocate(handle)?;
    }

    let elapsed = start_time.elapsed();
    let avg_per_cycle = elapsed / iterations;

    println!("Allocation/deallocation performance: {:?} per cycle", avg_per_cycle);

    // Should be sub-millisecond per allocation cycle
    assert!(avg_per_cycle < Duration::from_micros(100),
    "Allocation cycle too slow: {:?}", avg_per_cycle);

    // Performance test: scheduling overhead
    let scheduling_iterations = 100;
    let start_time = std::time::Instant::now();

    for i in 0..scheduling_iterations {
        let op = ScheduledOperation::MemoryTransfer {
            src: 0x1000 + (i as u64 * 1024),
            dst: 0x2000 + (i as u64 * 1024),
            size: 1024,
            device_id: 0,
        };

        let task = scheduler.schedule(op, Vec::new()).await?;
        scheduler.wait_for_completion(task).await?;
    }

    let scheduling_elapsed = start_time.elapsed();
    let avg_schedule_time = scheduling_elapsed / scheduling_iterations;

    println!("Average scheduling time: {:?}", avg_schedule_time);

    // Target: sub-microsecond scheduling overhead
    assert!(avg_schedule_time < Duration::from_micros(10),
    "Scheduling overhead too high: {:?}", avg_schedule_time);

    let final_stats = scheduler.statistics();
    assert!(final_stats.average_scheduling_time_ns < 10_000); // < 10μs

    Ok(())
}

/// Helper function to run a comprehensive system stress test
#[tokio::test]
async fn test_system_stress() -> Result<()> {
    let backend = get_test_backend();
    let allocator_config = AllocatorConfig {
        arena_size_bytes: 512 * 1024 * 1024, // 512MB for stress test
        max_fragmentation_percent: 20,
        enable_statistics: true,
    };
    let allocator = Arc::new(Allocator::new(allocator_config)?);

    let device = Device::new(backend, 0, allocator.clone())?;
    let mut scheduler = Scheduler::new(vec![device])?;

    println!("Starting system stress test...");

    // Create multiple concurrent workloads
    let num_concurrent_tasks = 50;
    let mut tasks = Vec::new();

    for i in 0..num_concurrent_tasks {
        let op = if i % 2 == 0 {
            ScheduledOperation::MemoryTransfer {
                src: 0x1000 + (i as u64 * 4096),
                dst: 0x10000 + (i as u64 * 4096),
                size: 4096 * (1 + i % 10), // Variable sizes
                device_id: 0,
            }
            } else {
                ScheduledOperation::KernelLaunch {
                    kernel_id: (i % 5) as u32, // Cycle through 5 kernel IDs
                    device_id: 0,
                }
            };

            // Create dependency chains
            let deps = if i > 0 && i % 10 == 0 {
                vec![tasks[i - 1]] // Every 10th task depends on previous
                } else {
                    Vec::new()
                };

                let task = scheduler.schedule(op, deps).await?;
                tasks.push(task);
            }

            // Wait for all tasks to complete
            let start_wait = std::time::Instant::now();
            for task in tasks {
                scheduler.wait_for_completion(task).await?;
            }
            let wait_duration = start_wait.elapsed();

            println!("Stress test completed in {:?}", wait_duration);

            // Verify system statistics
            let scheduler_stats = scheduler.statistics();
            assert_eq!(scheduler_stats.completed_tasks, num_concurrent_tasks);

            let allocator_stats = allocator.statistics();
            println!("Final allocator stats: {:?}", allocator_stats);

            // Should complete stress test within reasonable time
            assert!(wait_duration < Duration::from_secs(30),
            "Stress test took too long: {:?}", wait_duration);

            Ok(())
        }

        /// Test that verifies numerical determinism across runs
        #[tokio::test]
        async fn test_deterministic_execution() -> Result<()> {
            let backend = get_test_backend();
            let allocator_config = AllocatorConfig::default();
            let allocator = Arc::new(Allocator::new(allocator_config)?);

            let device = Device::new(backend, 0, allocator.clone())?;
            let scheduler = Scheduler::new(vec![device])?;

            // For now, just verify that identical operations produce identical results
            // (This would be expanded with actual kernel execution and result comparison)

            let op1 = ScheduledOperation::MemoryTransfer {
                src: 0x1000,
                dst: 0x2000,
                size: 1024,
                device_id: 0,
            };
            let op2 = ScheduledOperation::MemoryTransfer {
                src: 0x1000,
                dst: 0x2000,
                size: 1024,
                device_id: 0,
            };

            // Both operations should behave identically
            assert_eq!(op1, op2);

            println!("Deterministic execution test passed (basic)");

            Ok(())
        }