//! Stream management and scheduling
//!
//! This module provides stream abstractions for asynchronous GPU execution
//! with dependency tracking and sub-microsecond scheduling overhead.

use crate::error::{Result, RuntimeError};
#[cfg(feature = "cuda")]
use cudarc::driver::{CudaContext, CudaStream};
use std::sync::Arc;

// Stub types when CUDA is not available
#[cfg(not(feature = "cuda"))]
pub struct CudaContext;
#[cfg(not(feature = "cuda"))]
pub struct CudaStream;

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

/// GPU stream for asynchronous execution
#[derive(Clone)]
pub struct Stream {
    /// Stream ID
    pub id: StreamId,
    /// The underlying CUDA stream (if using CUDA backend)
    cuda_stream: Option<Arc<CudaStream>>,
}

impl Stream {
    /// Create a new stream with just an ID (for non-CUDA backends)
    #[inline]
    pub fn new(id: StreamId) -> Self {
        Self {
            id,
            cuda_stream: None,
        }
    }

    /// Create a new CUDA stream from a CudaContext
    #[cfg(feature = "cuda")]
    pub fn new_cuda(context: Arc<CudaContext>) -> Result<Self> {
        // For simplicity, use the default stream
        let cuda_stream = context.default_stream();
        Ok(Self {
            id: StreamId(0), // Default stream ID
            cuda_stream: Some(cuda_stream),
        })
    }

    #[cfg(not(feature = "cuda"))]
    pub fn new_cuda(_context: Arc<CudaContext>) -> Result<Self> {
        Err(RuntimeError::InvalidOperation(
            "CUDA stream creation requires cuda feature".to_string(),
        ))
    }

    /// Create a Stream from an existing CudaStream
    #[cfg(feature = "cuda")]
    pub fn from_cuda_stream(stream: Arc<CudaStream>) -> Self {
        Self {
            id: StreamId(0), // Default ID for external streams
            cuda_stream: Some(stream),
        }
    }

    #[cfg(not(feature = "cuda"))]
    pub fn from_cuda_stream(_stream: Arc<CudaStream>) -> Self {
        Self {
            id: StreamId(0),
            cuda_stream: None,
        }
    }

    /// Get the raw CUDA stream handle for kernel launching
    #[cfg(feature = "cuda")]
    pub fn raw_stream(&self) -> Result<&Arc<CudaStream>> {
        self.cuda_stream
            .as_ref()
            .ok_or(RuntimeError::InvalidOperation(
                "raw_stream: Stream does not have CUDA backend".to_string(),
            ))
    }

    #[cfg(not(feature = "cuda"))]
    pub fn raw_stream(&self) -> Result<&Arc<CudaStream>> {
        Err(RuntimeError::InvalidOperation(
            "raw_stream requires cuda feature".to_string(),
        ))
    }

    /// Check if this is a CUDA stream
    #[inline]
    pub fn is_cuda(&self) -> bool {
        self.cuda_stream.is_some()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_stream_creation() {
        let stream = Stream::new(StreamId(0));
        assert_eq!(stream.id, StreamId(0));
        assert!(!stream.is_cuda());
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_cuda_stream_creation() {
        // This test requires CUDA to be available
        // CudaContext::new returns Arc<CudaContext> in cudarc
        if let Ok(context) = CudaContext::new(0) {
            let stream = Stream::new_cuda(context).expect("Failed to create CUDA stream");
            assert!(stream.is_cuda());
            assert!(stream.raw_stream().is_ok());
        }
    }
}