//! Device abstraction layer
//!
//! This module provides unified abstractions for GPU devices across different backends
//! (CUDA, ROCm, Metal). The design prioritizes performance, safety, and ease of use.
//!
//! # Architecture
//!
//! - **Device**: Main abstraction representing a GPU device
//! - **Stream**: Asynchronous execution stream for non-blocking operations
//! - **Event**: Synchronization primitive for inter-stream coordination
//! - **Context**: Device context managing resources and state
//!
//! # Safety
//!
//! All device operations are carefully designed to prevent resource leaks and ensure
//! proper cleanup even in error conditions.

use crate::error::{RuntimeError, Result};
use crate::allocator::{GpuAllocator, DevicePtr};
use parking_lot::{RwLock, Mutex};
use std::collections::{BTreeMap, HashMap};
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::fmt;
use tracing::{debug, info, trace};

// CUDA integration - using simple pointer handles for now
// TODO: Use proper cudarc types once API is clarified

/// Device identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct DeviceId(pub u32);

impl fmt::Display for DeviceId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Device({})", self.0)
    }
}

/// Stream identifier for asynchronous execution
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct StreamId(pub u32);

impl fmt::Display for StreamId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Stream({})", self.0)
    }
}

/// Event identifier for synchronization
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EventId(pub u32);

impl fmt::Display for EventId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "Event({})", self.0)
    }
}

/// Stream priority levels for critical path optimization
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Default)]
pub enum StreamPriority {
    /// Highest priority for critical execution paths
    High,
    /// Normal priority for regular operations
    #[default]
    Normal,
    /// Lowest priority for background tasks
    Low,
}


/// GPU backend type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendType {
    /// NVIDIA CUDA backend
    Cuda,
    /// AMD ROCm backend
    Rocm,
    /// Apple Metal backend
    Metal,
    /// CPU fallback backend
    Cpu,
}

impl fmt::Display for BackendType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BackendType::Cuda => write!(f, "CUDA"),
            BackendType::Rocm => write!(f, "ROCm"),
            BackendType::Metal => write!(f, "Metal"),
            BackendType::Cpu => write!(f, "CPU"),
        }
    }
}

/// Device properties and capabilities
#[derive(Debug, Clone)]
pub struct DeviceProperties {
    /// Device name
    pub name: String,
    /// Backend type
    pub backend: BackendType,
    /// Total memory in bytes
    pub total_memory: usize,
    /// Available memory in bytes
    pub available_memory: usize,
    /// Compute capability major version
    pub major: u32,
    /// Compute capability minor version
    pub minor: u32,
    /// Number of multiprocessors
    pub multiprocessor_count: u32,
    /// Maximum threads per block
    pub max_threads_per_block: u32,
    /// Maximum shared memory per block
    pub max_shared_memory_per_block: usize,
    /// Supports unified memory
    pub unified_memory: bool,
    /// PCI bus ID
    pub pci_bus_id: String,
}

/// Stream execution statistics
#[derive(Debug, Clone, Default)]
pub struct StreamStats {
    /// Total kernels executed
    pub kernels_executed: u64,
    /// Total execution time in microseconds
    pub total_execution_time_us: u64,
    /// Number of synchronization events
    pub sync_events: u64,
    /// Number of memory operations
    pub memory_operations: u64,
}

/// GPU stream for asynchronous execution
pub struct Stream {
    /// Stream ID
    pub id: StreamId,
    /// Parent device
    pub device_id: DeviceId,
    /// Backend-specific handle (opaque)
    handle: u64,
    /// Stream statistics
    stats: Arc<Mutex<StreamStats>>,
    /// Active flag
    active: AtomicU64,
    /// Stream priority
    priority: StreamPriority,
    /// Stream callbacks
    callbacks: Arc<Mutex<Vec<Box<dyn FnOnce() + Send + 'static>>>>,
}

/// Event for synchronization between streams
pub struct Event {
    /// Event ID
    pub id: EventId,
    /// Parent device
    pub device_id: DeviceId,
    /// Backend-specific handle (opaque)
    handle: u64,
    /// Creation timestamp
    created_at: AtomicU64,
    /// Completion status
    completed: AtomicU64,
}

/// Device context managing resources
pub struct DeviceContext {
    /// Device ID
    pub device_id: DeviceId,
    /// Active streams
    streams: RwLock<HashMap<StreamId, Arc<Stream>>>,
    /// Active events
    events: RwLock<HashMap<EventId, Arc<Event>>>,
    /// Next stream ID
    next_stream_id: AtomicU32,
    /// Next event ID
    next_event_id: AtomicU32,
    /// Memory allocator for this device
    allocator: Arc<GpuAllocator>,
    /// Mock CUDA device initialization flag
    cuda_device_initialized: AtomicU64, // Use as boolean (0 = false, 1 = true)
}

/// GPU device abstraction
pub struct Device {
    /// Device ID
    pub id: DeviceId,
    /// Device properties
    pub properties: DeviceProperties,
    /// Device context
    context: Arc<DeviceContext>,
}

impl Device {
    /// Create a new device
    pub fn new(id: DeviceId, properties: DeviceProperties) -> Result<Self> {
        let context = Arc::new(DeviceContext {
            device_id: id,
            streams: RwLock::new(HashMap::new()),
            events: RwLock::new(HashMap::new()),
            next_stream_id: AtomicU32::new(0),
            next_event_id: AtomicU32::new(0),
            allocator: Arc::new(GpuAllocator::new()?),
            cuda_device_initialized: AtomicU64::new(0),
        });

        Ok(Self {
            id,
            properties,
            context,
        })
    }

    /// Get device properties
    pub fn properties(&self) -> &DeviceProperties {
        &self.properties
    }

    /// Create a new stream on this device
    pub fn create_stream(&self) -> Result<Arc<Stream>> {
        self.create_priority_stream(StreamPriority::Normal)
    }

    /// Create a new stream with specified priority
    pub fn create_priority_stream(&self, priority: StreamPriority) -> Result<Arc<Stream>> {
        let stream_id = StreamId(self.context.next_stream_id.fetch_add(1, Ordering::SeqCst));
        let handle = self.allocate_stream_handle_with_priority(priority)?;

        let stream = Arc::new(Stream {
            id: stream_id,
            device_id: self.id,
            handle,
            stats: Arc::new(Mutex::new(StreamStats::default())),
            active: AtomicU64::new(1),
            priority,
            callbacks: Arc::new(Mutex::new(Vec::new())),
        });

        self.context.streams.write().insert(stream_id, stream.clone());
        info!("Created stream {} on device {} with priority {:?}", stream_id, self.id, priority);
        Ok(stream)
    }

    /// Create a new event on this device
    pub fn create_event(&self) -> Result<Arc<Event>> {
        let event_id = EventId(self.context.next_event_id.fetch_add(1, Ordering::SeqCst));
        let handle = self.allocate_event_handle()?;

        let event = Arc::new(Event {
            id: event_id,
            device_id: self.id,
            handle,
            created_at: AtomicU64::new(current_time_us()),
            completed: AtomicU64::new(0),
        });

        self.context.events.write().insert(event_id, event.clone());
        info!("Created event {} on device {}", event_id, self.id);
        Ok(event)
    }

    /// Allocate device memory
    pub fn allocate(&self, size: usize) -> Result<DevicePtr> {
        self.context.allocator.allocate(size)
    }

    /// Free device memory
    pub fn free(&self, ptr: DevicePtr) -> Result<()> {
        self.context.allocator.free(ptr)
    }

    /// Get memory statistics
    pub fn memory_stats(&self) -> crate::allocator::AllocationStats {
        self.context.allocator.stats()
    }

    /// Synchronize the device (wait for all operations to complete)
    pub fn synchronize(&self) -> Result<()> {
        debug!("Synchronizing device {}", self.id);
        // In real implementation, this would call cuDeviceSynchronize() or equivalent
        self.device_synchronize_impl()
    }

    /// Get the number of active streams
    pub fn active_stream_count(&self) -> usize {
        self.context.streams.read().len()
    }

    /// Get the number of active events
    pub fn active_event_count(&self) -> usize {
        self.context.events.read().len()
    }

    /// Clean up completed resources
    pub fn cleanup(&self) -> Result<usize> {
        let mut cleaned = 0;

        // Clean up completed streams
        let mut streams = self.context.streams.write();
        streams.retain(|_, stream| {
            if stream.active.load(Ordering::SeqCst) == 0 {
                cleaned += 1;
                false
                } else {
                    true
                }
            });

            // Clean up completed events
            let mut events = self.context.events.write();
            events.retain(|_, event| {
                if event.completed.load(Ordering::SeqCst) > 0 {
                    cleaned += 1;
                    false
                    } else {
                        true
                    }
                });

                if cleaned > 0 {
                    debug!("Cleaned up {} resources on device {}", cleaned, self.id);
                }

                Ok(cleaned)
            }

            // Private implementation methods

            /// CUDA device initialization
            fn ensure_cuda_initialized(&self) -> Result<()> {
                // Use compare_exchange to ensure only one thread initializes
                if self.context.cuda_device_initialized.compare_exchange(
                0, 1, Ordering::SeqCst, Ordering::SeqCst
                ).is_ok() {
                    debug!("Initializing CUDA device {}", self.id.0);
                    
                    // Initialize CUDA device and set current context
                    #[cfg(feature = "cuda")]
                    {
                        use crate::cuda_backend;
                        unsafe {
                            cuda_backend::cuda_device_init(self.id.0 as i32)?;
                            cuda_backend::cuda_context_set_current(self.id.0 as i32)?;
                        }
                    }
                    #[cfg(not(feature = "cuda"))]
                    {
                        warn!("CUDA not available - using CPU fallback for device {}", self.id.0);
                    }
                }
                Ok(())
            }

            fn allocate_stream_handle(&self) -> Result<u64> {
                self.allocate_stream_handle_with_priority(StreamPriority::Normal)
            }

            fn allocate_stream_handle_with_priority(&self, priority: StreamPriority) -> Result<u64> {
                self.ensure_cuda_initialized()?;
                debug!("Creating CUDA stream on device {} with priority {:?}", self.id, priority);

                // Use real CUDA stream creation from cuda_backend
                use crate::cuda_backend;

                let cuda_stream = unsafe {
                    cuda_backend::cuda_stream_create_with_priority(self.id.0 as i32, priority)?
                };

                // Convert stream handle to u64 for storage
                let handle = Box::into_raw(Box::new(cuda_stream)) as u64;
                Ok(handle)
            }

            fn allocate_event_handle(&self) -> Result<u64> {
                self.ensure_cuda_initialized()?;
                debug!("Creating CUDA event on device {}", self.id);

                // Use real CUDA event creation from cuda_backend
                use crate::cuda_backend;

                let cuda_event = unsafe {
                    cuda_backend::cuda_event_create(self.id.0 as i32)?
                };

                // Convert event handle to u64 for storage
                let handle = Box::into_raw(Box::new(cuda_event)) as u64;
                Ok(handle)
            }

            fn device_synchronize_impl(&self) -> Result<()> {
                self.ensure_cuda_initialized()?;
                debug!("Synchronizing CUDA device {}", self.id);

                // Use real CUDA device synchronization
                use crate::cuda_backend;
                unsafe { cuda_backend::cuda_device_synchronize()? };

                Ok(())
            }

            /// Create a mock device for testing
            pub fn mock_device(device_id: u32) -> Result<Self> {
                let props = DeviceProperties {
                    name: "Mock GPU".to_string(),
                    backend: BackendType::Cuda,
                    total_memory: 8 * 1024 * 1024 * 1024,
                    available_memory: 7 * 1024 * 1024 * 1024,
                    major: 8,
                    minor: 0,
                    multiprocessor_count: 80,
                    max_threads_per_block: 1024,
                    max_shared_memory_per_block: 48 * 1024,
                    unified_memory: false,
                    pci_bus_id: format!("0000:0{}:00.0", device_id),
                };

                Self::new(DeviceId(device_id), props)
            }
        }

        impl Stream {
            /// Get stream statistics
            pub fn stats(&self) -> StreamStats {
                self.stats.lock().clone()
            }

            /// Get stream priority
            pub fn priority(&self) -> StreamPriority {
                self.priority
            }

            /// Add callback to be executed after stream completion
            pub fn add_callback(&self, callback: Box<dyn FnOnce() + Send + 'static>) -> Result<()> {
                self.callbacks.lock().push(callback);
                Ok(())
            }

            /// Query stream completion status without blocking
            pub fn query(&self) -> Result<bool> {
                use crate::cuda_backend;

                // Get CUDA stream handle
                let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };

                // Query CUDA stream
                unsafe { cuda_backend::cuda_stream_query(cuda_stream) }
            }

            /// CUDA stream operations
            fn cuda_stream_op(&self, operation: &str) -> Result<()> {
                trace!("{} operation on stream {}", operation, self.id);
                #[cfg(feature = "cuda")]
                {
                    // In a real implementation, this would perform the actual CUDA operation
                    // For now, just ensure the stream is valid
                    self.ensure_cuda_initialized()?;
                }
                Ok(())
            }

            /// Launch a kernel on this stream
            pub fn launch_kernel(&self, kernel_name: &str, grid_size: (u32, u32, u32), block_size: (u32, u32, u32)) -> Result<()> {
                trace!("Launching kernel '{}' on stream {} with grid {:?}, block {:?}",
                kernel_name, self.id, grid_size, block_size);

                // Update statistics
                let mut stats = self.stats.lock();
                stats.kernels_executed += 1;

                #[cfg(feature = "cuda")]
                {
                    use crate::cuda_backend;
                    // Launch kernel using CUDA backend
                    unsafe {
                        let stream = self.get_cuda_stream()?;
                        cuda_backend::cuda_kernel_launch(
                            kernel_name,
                            grid_size,
                            block_size,
                            stream
                        )?;
                    }
                }
                #[cfg(not(feature = "cuda"))]
                {
                    trace!("CUDA not available - kernel '{}' execution simulated", kernel_name);
                }

                Ok(())
            }

                /// Copy memory asynchronously on this stream
                pub fn copy_async(&self, dst: DevicePtr, src: DevicePtr, size: usize) -> Result<()> {
                    trace!("Async copy {} bytes from {:?} to {:?} on stream {}",
                    size, src, dst, self.id);

                    // Validate pointers
                    if src.is_null() || dst.is_null() {
                        return Err(RuntimeError::device_error(
                        self.device_id.0,
                        "Cannot copy from/to null pointer".to_string()
                        ));
                    }

                    use crate::cuda_backend;

                    // Get CUDA stream handle
                    let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };

                    // Perform real CUDA async memory copy
                    unsafe {
                        cuda_backend::cuda_memcpy_device_to_device_async(dst, src, size, cuda_stream)?
                    };

                    let mut stats = self.stats.lock();
                    stats.memory_operations += 1;

                    Ok(())
                }

                /// Record an event on this stream
                pub fn record_event(&self, event: &Event) -> Result<()> {
                    if event.device_id != self.device_id {
                        return Err(RuntimeError::device_error(
                        self.device_id.0,
                        format!("Event {} belongs to device {}, cannot record on stream {} of device {}",
                        event.id, event.device_id, self.id, self.device_id)
                        ));
                    }

                    trace!("Recording event {} on stream {}", event.id, self.id);

                    use crate::cuda_backend;

                    // Get handles
                    let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
                    let cuda_event = unsafe { &*(event.handle as *const crate::cuda_backend::CudaEventHandle) };

                    // Record event on CUDA stream
                    unsafe { cuda_backend::cuda_event_record(cuda_event, cuda_stream)? };

                    let mut stats = self.stats.lock();
                    stats.sync_events += 1;

                    Ok(())
                }

                /// Wait for an event on this stream
                pub fn wait_event(&self, event: &Event) -> Result<()> {
                    if event.device_id != self.device_id {
                        return Err(RuntimeError::device_error(
                        self.device_id.0,
                        format!("Event {} belongs to device {}, cannot wait on stream {} of device {}",
                        event.id, event.device_id, self.id, self.device_id)
                        ));
                    }

                    trace!("Waiting for event {} on stream {}", event.id, self.id);

                    use crate::cuda_backend;

                    // Get handles
                    let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
                    let cuda_event = unsafe { &*(event.handle as *const crate::cuda_backend::CudaEventHandle) };

                    // Wait for event on CUDA stream
                    unsafe { cuda_backend::cuda_stream_wait_event(cuda_stream, cuda_event)? };

                    Ok(())
                }

                /// Synchronize this stream (wait for completion)
                pub fn synchronize(&self) -> Result<()> {
                    debug!("Synchronizing stream {}", self.id);

                    use crate::cuda_backend;

                    // Get CUDA stream handle
                    let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };

                    // Synchronize CUDA stream
                    unsafe { cuda_backend::cuda_stream_synchronize(cuda_stream)? };

                    // Execute any pending callbacks
                    let mut callbacks = self.callbacks.lock();
                    for callback in callbacks.drain(..) {
                        callback();
                    }

                    Ok(())
                }

                /// Check if stream is active
                pub fn is_active(&self) -> bool {
                    self.active.load(Ordering::SeqCst) > 0
                }

                /// Mark stream as inactive (internal use)
                pub fn deactivate(&self) {
                    self.active.store(0, Ordering::SeqCst);
                }

                /// Get the raw stream handle (for internal use by backends)
                pub fn raw_handle(&self) -> u64 {
                    self.handle
                }

                /// Create a simple stream with default values (for testing/simplified usage)
                pub fn new(id: StreamId) -> Self {
                    use std::sync::atomic::AtomicU64;
                    use parking_lot::Mutex;

                    Self {
                        id,
                        device_id: DeviceId(0), // Default device
                        handle: 0, // Default handle
                        stats: Arc::new(Mutex::new(StreamStats {
                            kernels_executed: 0,
                            total_execution_time_us: 0,
                            sync_events: 0,
                            memory_operations: 0,
                        })),
                        active: AtomicU64::new(0),
                        priority: StreamPriority::Normal,
                        callbacks: Arc::new(Mutex::new(Vec::new())),
                    }
                }
            }

            // Manual Debug implementation for Stream
            impl std::fmt::Debug for Stream {
                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                    use std::sync::atomic::Ordering;

                    f.debug_struct("Stream")
                    .field("id", &self.id)
                    .field("device_id", &self.device_id)
                    .field("handle", &self.handle)
                    .field("stats", &self.stats)
                    .field("active", &self.active.load(Ordering::SeqCst))
                    .field("priority", &self.priority)
                    .field("callbacks", &format!("{} callbacks", self.callbacks.lock().len()))
                    .finish()
                }
            }

            impl Event {
                /// Check if event has completed
                pub fn is_completed(&self) -> bool {
                    // First check our cached completion status
                    if self.completed.load(Ordering::SeqCst) > 0 {
                        return true;
                    }

                    // Query real CUDA event completion
                    self.query().unwrap_or(false)
                }

                /// Query event completion status (non-blocking)
                pub fn query(&self) -> Result<bool> {
                    use crate::cuda_backend;

                    // Get CUDA event handle
                    let cuda_event = unsafe { &*(self.handle as *const crate::cuda_backend::CudaEventHandle) };

                    // Query CUDA event
                    let is_completed = unsafe { cuda_backend::cuda_event_query(cuda_event)? };

                    if is_completed {
                        self.completed.store(current_time_us(), Ordering::SeqCst);
                    }

                    Ok(is_completed)
                }

                /// Wait for event completion (blocking)
                pub fn wait(&self) -> Result<()> {
                    debug!("Waiting for event {} completion", self.id);

                    use crate::cuda_backend;

                    // Get CUDA event handle
                    let cuda_event = unsafe { &*(self.handle as *const crate::cuda_backend::CudaEventHandle) };

                    // Synchronize CUDA event (blocking)
                    unsafe { cuda_backend::cuda_event_synchronize(cuda_event)? };

                    // Mark as completed
                    self.completed.store(current_time_us(), Ordering::SeqCst);

                    Ok(())
                }

                /// Calculate elapsed time between events in milliseconds
                pub fn elapsed_time_ms(&self, start_event: &Event) -> Result<f32> {
                    use crate::cuda_backend;

                    // Get CUDA event handles
                    let start_cuda_event = unsafe { &*(start_event.handle as *const crate::cuda_backend::CudaEventHandle) };
                    let end_cuda_event = unsafe { &*(self.handle as *const crate::cuda_backend::CudaEventHandle) };

                    // Calculate elapsed time using CUDA events
                    unsafe { cuda_backend::cuda_event_elapsed_time(start_cuda_event, end_cuda_event) }
                }

                /// Get elapsed time since event creation
                pub fn elapsed_time_us(&self) -> u64 {
                    let current = current_time_us();
                    let created = self.created_at.load(Ordering::SeqCst);
                    current.saturating_sub(created)
                }
            }

            /// Get current time in microseconds (mock implementation)
            fn current_time_us() -> u64 {
                use std::time::{SystemTime, UNIX_EPOCH};
                SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_micros() as u64
            }

            /// Discover available devices
            pub fn discover_devices(devices: &mut BTreeMap<DeviceId, Device>) -> Result<usize> {
                // Use real CUDA device discovery
                crate::cuda_backend::discover_cuda_devices(devices)
                .map_err(|e| RuntimeError::device_error(0, format!("CUDA device discovery failed: {}", e)))
            }

            // Comprehensive failing tests (TDD approach)
            #[cfg(test)]
            mod tests {
                use super::*;
                use std::thread;
                use std::time::Duration;

                // Basic device tests

                #[test]
                fn test_device_creation() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).expect("Device creation should succeed");
                    assert_eq!(device.id, DeviceId(0));
                    assert_eq!(device.properties.name, "Test GPU");
                    assert_eq!(device.properties.backend, BackendType::Cuda);
                    assert_eq!(device.active_stream_count(), 0);
                    assert_eq!(device.active_event_count(), 0);
                }

                #[test]
                fn test_device_discovery() {
                    let mut devices = BTreeMap::new();
                    let count = discover_devices(&mut devices).expect("Device discovery should succeed");
                    assert_eq!(count, 2);
                    assert!(devices.contains_key(&DeviceId(0)));
                    assert!(devices.contains_key(&DeviceId(1)));

                    let device_0 = devices.get(&DeviceId(0)).unwrap();
                    assert_eq!(device_0.properties.name, "RTX 5090");
                    assert_eq!(device_0.properties.backend, BackendType::Cuda);
                }

                // Stream tests (should fail initially due to unimplemented methods)

                #[test]
                fn test_stream_creation() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    let stream = device.create_stream().expect("Stream creation should succeed");

                    assert_eq!(stream.device_id, DeviceId(0));
                    assert_eq!(stream.id, StreamId(0));
                    assert!(stream.is_active());
                    assert_eq!(device.active_stream_count(), 1);
                }

                #[test]
                fn test_multiple_streams() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();

                    let stream_1 = device.create_stream().unwrap();
                    let stream_2 = device.create_stream().unwrap();
                    let stream_3 = device.create_stream().unwrap();

                    assert_eq!(device.active_stream_count(), 3);
                    assert_eq!(stream_1.id, StreamId(0));
                    assert_eq!(stream_2.id, StreamId(1));
                    assert_eq!(stream_3.id, StreamId(2));

                    // All streams should belong to the same device
                    assert_eq!(stream_1.device_id, DeviceId(0));
                    assert_eq!(stream_2.device_id, DeviceId(0));
                    assert_eq!(stream_3.device_id, DeviceId(0));
                }

                #[test]
                fn test_stream_kernel_launch() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    let stream = device.create_stream().unwrap();

                    stream.launch_kernel("test_kernel", (1, 1, 1), (256, 1, 1))
                    .expect("Kernel launch should succeed");

                    let stats = stream.stats();
                    assert_eq!(stats.kernels_executed, 1);
                }

                // Event tests (should fail initially)

                #[test]
                fn test_event_creation() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    let event = device.create_event().expect("Event creation should succeed");

                    assert_eq!(event.device_id, DeviceId(0));
                    assert_eq!(event.id, EventId(0));
                    assert!(!event.is_completed());
                    assert_eq!(device.active_event_count(), 1);
                }

                #[test]
                fn test_stream_event_interaction() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    let stream = device.create_stream().unwrap();
                    let event = device.create_event().unwrap();

                    // Record event on stream
                    stream.record_event(&event).expect("Event recording should succeed");

                    // Wait for event on another stream
                    let stream2 = device.create_stream().unwrap();
                    stream2.wait_event(&event).expect("Event waiting should succeed");
                }

                #[test]
                fn test_cross_device_event_error() {
                    let props1 = DeviceProperties {
                        name: "GPU 1".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let props2 = DeviceProperties {
                        name: "GPU 2".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:02:00.0".to_string(),
                    };

                    let device1 = Device::new(DeviceId(0), props1).unwrap();
                    let device2 = Device::new(DeviceId(1), props2).unwrap();

                    // This should work without panicking
                    let _stream1 = device1.create_stream();
                    let _event2 = device2.create_event();

                    // Cross-device operations should return errors, not panic
                    // We can't test this fully until stream creation is implemented
                }

                // Memory allocation tests

                #[test]
                fn test_device_memory_allocation() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();

                    // Test allocation
                    let ptr = device.allocate(1024).expect("Memory allocation should succeed");
                    assert!(!ptr.is_null());

                    // Test stats tracking
                    let stats = device.memory_stats();
                    assert!(stats.active_allocations > 0);
                    assert!(stats.total_allocated > 0);

                    // Test deallocation
                    device.free(ptr).expect("Memory deallocation should succeed");
                }

                #[test]
                fn test_device_synchronization() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    device.synchronize().expect("Device synchronization should succeed");
                }

                // Cleanup and resource management tests

                #[test]
                fn test_resource_cleanup() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();

                    // Create some resources
                    let stream = device.create_stream().unwrap();
                    let event = device.create_event().unwrap();

                    assert_eq!(device.active_stream_count(), 1);
                    assert_eq!(device.active_event_count(), 1);

                    // Deactivate resources
                    stream.deactivate();

                    // Cleanup should remove inactive resources
                    let cleaned = device.cleanup().expect("Cleanup should succeed");
                    assert!(cleaned > 0);
                }

                // Performance and stress tests

                #[test]
                fn test_many_streams() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    let mut streams = Vec::new();

                    // Create many streams
                    for _ in 0..100 {
                        streams.push(device.create_stream().unwrap());
                    }

                    assert_eq!(device.active_stream_count(), 100);

                    // Launch kernels on all streams
                    for stream in &streams {
                        stream.launch_kernel("test_kernel", (1, 1, 1), (256, 1, 1)).unwrap();
                    }

                    // Synchronize all streams
                    for stream in &streams {
                        stream.synchronize().unwrap();
                    }
                }

                // TDD tests for real CUDA streams (should fail initially)

                #[test]
                fn test_real_cuda_stream_creation() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let stream = device.create_stream().expect("Real CUDA stream creation should succeed");

                    assert!(stream.is_active());
                    assert_eq!(stream.device_id, device.id);

                    // Stream should have real CUDA handle (not mock)
                    assert_ne!(stream.handle, 0);
                }

                #[test]
                fn test_real_cuda_stream_synchronization() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let stream = device.create_stream().expect("Stream creation should succeed");

                    // Real synchronization should work
                    stream.synchronize().expect("Real CUDA stream synchronization should succeed");
                }

                #[test]
                fn test_real_cuda_stream_priority() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");

                    // Create high priority stream for critical paths
                    let high_priority_stream = device.create_priority_stream(StreamPriority::High)
                    .expect("High priority stream creation should succeed");

                    // Create low priority stream for background tasks
                    let low_priority_stream = device.create_priority_stream(StreamPriority::Low)
                    .expect("Low priority stream creation should succeed");

                    // Priority streams should have different handles
                    assert_ne!(high_priority_stream.handle, low_priority_stream.handle);
                    assert_eq!(high_priority_stream.priority(), StreamPriority::High);
                    assert_eq!(low_priority_stream.priority(), StreamPriority::Low);
                }

                #[test]
                fn test_real_cuda_stream_callbacks() {
                    use std::sync::atomic::{AtomicBool, Ordering};
                    use std::sync::Arc;

                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let stream = device.create_stream().expect("Stream creation should succeed");

                    // Test callback execution
                    let callback_executed = Arc::new(AtomicBool::new(false));
                    let callback_flag = callback_executed.clone();

                    stream.add_callback(Box::new(move || {
                        callback_flag.store(true, Ordering::SeqCst);
                    })).expect("Adding callback should succeed");

                    // Synchronize to ensure callback is executed
                    stream.synchronize().expect("Stream sync should succeed");

                    // Callback should have been executed
                    assert!(callback_executed.load(Ordering::SeqCst));
                }

                #[test]
                fn test_real_cuda_memory_copy_async() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let stream = device.create_stream().expect("Stream creation should succeed");

                    // Allocate source and destination memory
                    let src_ptr = device.allocate(1024).expect("Source allocation should succeed");
                    let dst_ptr = device.allocate(1024).expect("Destination allocation should succeed");

                    // Real async copy should work
                    stream.copy_async(dst_ptr, src_ptr, 1024)
                    .expect("Real async copy should succeed");

                    // Synchronize to ensure copy completes
                    stream.synchronize().expect("Stream sync should succeed");

                    // Clean up
                    device.free(src_ptr).expect("Source deallocation should succeed");
                    device.free(dst_ptr).expect("Destination deallocation should succeed");
                }

                #[test]
                fn test_real_cuda_stream_query_status() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let stream = device.create_stream().expect("Stream creation should succeed");

                    // Query should work on real stream
                    let is_ready = stream.query().expect("Stream query should succeed");

                    // Empty stream should be ready
                    assert!(is_ready);
                }

                // TDD tests for real CUDA events (should fail initially)

                #[test]
                fn test_real_cuda_event_creation() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let event = device.create_event().expect("Real CUDA event creation should succeed");

                    assert_eq!(event.device_id, device.id);
                    // Event should have real CUDA handle (not mock)
                    assert_ne!(event.handle, 0);
                }

                #[test]
                fn test_real_cuda_event_timing() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");

                    // Create events for timing
                    let start_event = device.create_event().expect("Start event creation should succeed");
                    let end_event = device.create_event().expect("End event creation should succeed");

                    let stream = device.create_stream().expect("Stream creation should succeed");

                    // Record start event
                    stream.record_event(&start_event).expect("Recording start event should succeed");

                    // Simulate some work (memory allocation/deallocation)
                    let ptr = device.allocate(1024 * 1024).expect("Allocation should succeed");
                    device.free(ptr).expect("Deallocation should succeed");

                    // Record end event
                    stream.record_event(&end_event).expect("Recording end event should succeed");

                    // Wait for completion
                    stream.synchronize().expect("Stream sync should succeed");

                    // Calculate elapsed time
                    let elapsed_ms = end_event.elapsed_time_ms(&start_event)
                    .expect("Elapsed time calculation should succeed");

                    // Should have some measurable time
                    assert!(elapsed_ms >= 0.0);
                }

                #[test]
                fn test_real_cuda_event_record_and_wait() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");

                    let stream1 = device.create_stream().expect("Stream1 creation should succeed");
                    let stream2 = device.create_stream().expect("Stream2 creation should succeed");
                    let event = device.create_event().expect("Event creation should succeed");

                    // Record event on stream1
                    stream1.record_event(&event).expect("Recording event should succeed");

                    // Wait for event on stream2 (inter-stream dependency)
                    stream2.wait_event(&event).expect("Waiting for event should succeed");

                    // Both streams should synchronize without error
                    stream1.synchronize().expect("Stream1 sync should succeed");
                    stream2.synchronize().expect("Stream2 sync should succeed");
                }

                #[test]
                fn test_real_cuda_event_query() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");
                    let event = device.create_event().expect("Event creation should succeed");
                    let stream = device.create_stream().expect("Stream creation should succeed");

                    // Record event on stream
                    stream.record_event(&event).expect("Recording event should succeed");

                    // Query event status - might be completed or not
                    let is_completed = event.query().expect("Event query should succeed");

                    // After synchronization, event should definitely be completed
                    stream.synchronize().expect("Stream sync should succeed");
                    let is_completed_after_sync = event.query().expect("Event query after sync should succeed");
                    assert!(is_completed_after_sync);
                }

                // TDD tests for CUDA device synchronization

                #[test]
                fn test_real_cuda_device_synchronize() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");

                    // Create multiple streams with work
                    let stream1 = device.create_stream().expect("Stream1 creation should succeed");
                    let stream2 = device.create_stream().expect("Stream2 creation should succeed");

                    // Launch work on both streams
                    let ptr1 = device.allocate(1024).expect("Allocation1 should succeed");
                    let ptr2 = device.allocate(1024).expect("Allocation2 should succeed");

                    stream1.copy_async(ptr1, ptr1, 1024).expect("Copy on stream1 should succeed");
                    stream2.copy_async(ptr2, ptr2, 1024).expect("Copy on stream2 should succeed");

                    // Device sync should wait for all streams
                    device.synchronize().expect("Real CUDA device synchronization should succeed");

                    // Clean up
                    device.free(ptr1).expect("Deallocation1 should succeed");
                    device.free(ptr2).expect("Deallocation2 should succeed");
                }

                // TDD tests for advanced stream features

                #[test]
                fn test_cuda_stream_dependencies() {
                    let mut devices = BTreeMap::new();
                    discover_devices(&mut devices).expect("Should discover CUDA devices");

                    let device = devices.values().next().expect("Should have at least one device");

                    // Create dependency chain: stream1 -> event -> stream2
                    let stream1 = device.create_stream().expect("Stream1 creation should succeed");
                    let stream2 = device.create_stream().expect("Stream2 creation should succeed");
                    let event = device.create_event().expect("Event creation should succeed");

                    // Work on stream1
                    let ptr = device.allocate(1024).expect("Allocation should succeed");
                    stream1.copy_async(ptr, ptr, 1024).expect("Copy on stream1 should succeed");

                    // Record completion event
                    stream1.record_event(&event).expect("Recording event should succeed");

                    // Stream2 waits for stream1 completion
                    stream2.wait_event(&event).expect("Waiting for event should succeed");

                    // More work on stream2 (depends on stream1)
                    stream2.copy_async(ptr, ptr, 1024).expect("Copy on stream2 should succeed");

                    // Synchronize all
                    device.synchronize().expect("Device sync should succeed");

                    device.free(ptr).expect("Deallocation should succeed");
                }

                // Edge cases and error conditions

                #[test]
                fn test_backend_type_display() {
                    assert_eq!(format!("{}", BackendType::Cuda), "CUDA");
                    assert_eq!(format!("{}", BackendType::Rocm), "ROCm");
                    assert_eq!(format!("{}", BackendType::Metal), "Metal");
                    assert_eq!(format!("{}", BackendType::Cpu), "CPU");
                }

                #[test]
                fn test_id_display() {
                    assert_eq!(format!("{}", DeviceId(42)), "Device(42)");
                    assert_eq!(format!("{}", StreamId(123)), "Stream(123)");
                    assert_eq!(format!("{}", EventId(456)), "Event(456)");
                }

                #[test]
                fn test_event_elapsed_time() {
                    let props = DeviceProperties {
                        name: "Test GPU".to_string(),
                        backend: BackendType::Cuda,
                        total_memory: 8 * 1024 * 1024 * 1024,
                        available_memory: 7 * 1024 * 1024 * 1024,
                        major: 8,
                        minor: 0,
                        multiprocessor_count: 80,
                        max_threads_per_block: 1024,
                        max_shared_memory_per_block: 48 * 1024,
                        unified_memory: false,
                        pci_bus_id: "0000:01:00.0".to_string(),
                    };

                    let device = Device::new(DeviceId(0), props).unwrap();
                    // Can't fully test without event creation working, but we can test the time function
                    let start_time = current_time_us();
                    thread::sleep(Duration::from_millis(1));
                    let end_time = current_time_us();
                    assert!(end_time > start_time);
                }
            }