//! CUDA device management.
//!
//! Provides the [`CudaDevice`] type for managing CUDA contexts and streams.

#[cfg(feature = "cuda")]
use crate::CudaBackend;
#[cfg(feature = "cuda")]
use cudarc::cublas::CudaBlas;
#[cfg(feature = "cuda")]
use cudarc::driver::{CudaContext, CudaStream};
#[cfg(feature = "cuda")]
use once_cell::sync::{Lazy, OnceCell};
#[cfg(feature = "cuda")]
use parking_lot::RwLock;
#[cfg(feature = "cuda")]
use rtx_backend::{DeviceId, DeviceOps};
#[cfg(feature = "cuda")]
use std::collections::HashMap;
#[cfg(feature = "cuda")]
use std::sync::Arc;

use crate::{CudaError, CudaResult};

#[cfg(feature = "cuda")]
/// Global cache of CUDA contexts for device reuse.
static CONTEXT_CACHE: Lazy<RwLock<HashMap<usize, Arc<CudaContext>>>> =
    Lazy::new(|| RwLock::new(HashMap::new()));

#[cfg(feature = "cuda")]
/// CUDA device wrapper with context and stream management.
///
/// Each `CudaDevice` represents a specific GPU and holds:
/// - A CUDA context (shared across the device)
/// - A default stream for operations
/// - Device properties (memory, compute capability)
/// - Lazy-initialized cuBLAS handle
#[derive(Clone)]
pub struct CudaDevice {
    /// Device index (0, 1, 2, ...)
    pub(crate) index: usize,
    /// CUDA context handle
    pub(crate) context: Arc<CudaContext>,
    /// Default stream for operations
    pub(crate) stream: Arc<CudaStream>,
    /// Total memory in bytes
    pub(crate) memory_total: usize,
    /// Compute capability (major, minor)
    pub(crate) compute_capability: (u32, u32),
    /// cuBLAS handle (lazy-initialized, shared across clones)
    pub(crate) cublas: Arc<OnceCell<CudaBlas>>,
}

#[cfg(not(feature = "cuda"))]
/// CUDA device wrapper (stub when CUDA is not available).
#[derive(Clone)]
pub struct CudaDevice {
    /// Device index (0, 1, 2, ...)
    pub(crate) index: usize,
}

#[cfg(feature = "cuda")]
impl CudaDevice {
    /// Create a new CUDA device.
    ///
    /// # Arguments
    /// - `index`: GPU device index (0 for first GPU)
    ///
    /// # Returns
    /// CUDA device if available, error otherwise.
    pub fn new(index: usize) -> CudaResult<Self> {
        // Check cache first
        {
            let cache = CONTEXT_CACHE.read();
            if let Some(context) = cache.get(&index) {
                let stream = context.default_stream();
                return Ok(Self {
                    index,
                    context: context.clone(),
                    stream,
                    memory_total: 24 * 1024 * 1024 * 1024, // 24GB default
                    compute_capability: (9, 0),            // sm_90 (RTX 5090)
                    cublas: Arc::new(OnceCell::new()),
                });
            }
        }

        // Initialize CUDA
        cudarc::driver::result::init()
            .map_err(|e| CudaError::Initialization(format!("Failed to init CUDA: {:?}", e)))?;

        // Get device count
        let count = CudaContext::device_count().map_err(|e| {
            CudaError::Initialization(format!("Failed to get device count: {:?}", e))
        })?;

        if index >= count as usize {
            return Err(CudaError::DeviceNotFound(index, count as usize));
        }

        // Create context (CudaContext::new returns Arc<CudaContext>)
        let context = CudaContext::new(index)
            .map_err(|e| CudaError::Initialization(format!("Failed to create context: {:?}", e)))?;

        // Get default stream
        let stream = context.default_stream();

        // Cache the context
        {
            let mut cache = CONTEXT_CACHE.write();
            cache.insert(index, context.clone());
        }

        // Get device properties
        let memory_total = context
            .attribute(
                cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_TOTAL_CONSTANT_MEMORY,
            )
            .unwrap_or(24_i32 * 1024 * 1024) as usize
            * 1024; // 24GB default

        // Get compute capability
        let compute_capability = context
            .compute_capability()
            .map(|(major, minor)| (major as u32, minor as u32))
            .unwrap_or((9, 0));

        Ok(Self {
            index,
            context,
            stream,
            memory_total,
            compute_capability,
            cublas: Arc::new(OnceCell::new()),
        })
    }

    /// Get the CUDA context.
    pub fn context(&self) -> &Arc<CudaContext> {
        &self.context
    }

    /// Get the default CUDA stream.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Get or create a cuBLAS handle for this device.
    ///
    /// The handle is lazily initialized and cached for reuse.
    pub fn cublas(&self) -> CudaResult<&CudaBlas> {
        self.cublas.get_or_try_init(|| {
            CudaBlas::new(self.stream.clone())
                .map_err(|e| CudaError::CuBlas(format!("Failed to create cuBLAS handle: {:?}", e)))
        })
    }

    /// Create a new CUDA stream on this device.
    pub fn create_stream(&self) -> CudaResult<Arc<CudaStream>> {
        self.context
            .new_stream()
            .map_err(|e| CudaError::StreamCreation(format!("{:?}", e)))
    }

    /// Get device name.
    pub fn name(&self) -> String {
        self.context
            .name()
            .unwrap_or_else(|_| format!("CUDA Device {}", self.index))
    }

    /// Synchronize all operations on this device.
    pub fn synchronize(&self) {
        let _ = self.stream.synchronize();
    }
}

#[cfg(not(feature = "cuda"))]
impl CudaDevice {
    /// Create a new CUDA device (stub - always fails without CUDA feature).
    pub fn new(_index: usize) -> CudaResult<Self> {
        Err(CudaError::Initialization(
            "CUDA support not compiled. Enable the 'cuda' feature.".to_string(),
        ))
    }

    /// Get device name (stub).
    pub fn name(&self) -> String {
        format!("CUDA Device {} (unavailable)", self.index)
    }

    /// Synchronize all operations on this device (stub - no-op).
    pub fn synchronize(&self) {
        // No-op
    }
}

#[cfg(feature = "cuda")]
impl std::fmt::Debug for CudaDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaDevice")
            .field("index", &self.index)
            .field("name", &self.name())
            .field("memory_total", &self.memory_total)
            .field("compute_capability", &self.compute_capability)
            .finish()
    }
}

#[cfg(not(feature = "cuda"))]
impl std::fmt::Debug for CudaDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaDevice")
            .field("index", &self.index)
            .field("name", &"unavailable")
            .finish()
    }
}

impl Default for CudaDevice {
    fn default() -> Self {
        Self::new(0).expect("No CUDA device available")
    }
}

impl PartialEq for CudaDevice {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index
    }
}

impl Eq for CudaDevice {}

impl std::hash::Hash for CudaDevice {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.index.hash(state);
    }
}

#[cfg(feature = "cuda")]
impl DeviceOps<CudaBackend> for CudaDevice {
    fn id(&self) -> DeviceId {
        DeviceId::Cuda(self.index)
    }

    fn memory_capacity(&self) -> usize {
        self.memory_total
    }

    fn memory_available(&self) -> usize {
        // Query available memory from CUDA (simplified)
        self.memory_total / 2
    }

    fn compute_capability(&self) -> Option<(u32, u32)> {
        Some(self.compute_capability)
    }

    fn synchronize(&self) {
        let _ = self.stream.synchronize();
    }

    fn is_available(&self) -> bool {
        CudaContext::device_count()
            .map(|c| self.index < c as usize)
            .unwrap_or(false)
    }
}

// Note: DeviceOps impl for non-cuda builds is not provided because
// CudaBackend doesn't implement Backend without the cuda feature.

#[cfg(all(test, feature = "cuda"))]
mod tests {
    use super::*;

    #[test]
    fn test_device_creation() {
        // This test requires CUDA hardware
        if let Ok(device) = CudaDevice::new(0) {
            assert_eq!(device.index, 0);
            assert!(device.memory_total > 0);
            #[cfg(feature = "cuda")]
            assert!(device.is_available());
        }
    }
}