//! CUDA Graphs: capture a sequence of kernel launches and replay them with
//! near-zero host-side overhead (~3-5 us per launch eliminated).
//!
//! Usage:
//! ```ignore
//! let stream = CudaStream::new()?;
//! let mut graph = CudaGraph::new();
//!
//! // First call: capture
//! graph.begin_capture(&stream)?;
//! // ... launch kernels on `stream` ...
//! graph.end_capture(&stream)?;
//!
//! // Subsequent calls: replay
//! graph.launch(&stream)?;
//! ```
//!
//! Requirements for captured kernels:
//! - All tensor shapes must be identical between capture and replay.
//! - No host-side branching during the captured section.
//! - Memory addresses used during capture must remain valid during replay.

use crate::error::{self, Result};
use crate::ffi;
use crate::stream::CudaStream;

/// RAII wrapper around a captured CUDA graph and its executable instance.
pub struct CudaGraph {
    graph: ffi::CudaGraph,
    exec: ffi::CudaGraphExec,
}

impl CudaGraph {
    /// Create an empty graph handle (not yet captured).
    pub fn new() -> Self {
        Self {
            graph: std::ptr::null_mut(),
            exec: std::ptr::null_mut(),
        }
    }

    /// Returns true if a graph has been captured and instantiated.
    pub fn is_ready(&self) -> bool {
        !self.exec.is_null()
    }

    /// Begin capturing kernel launches on `stream`.
    /// All subsequent kernel launches on this stream are recorded into the
    /// graph instead of being executed.
    pub fn begin_capture(&mut self, stream: &CudaStream) -> Result<()> {
        // If we have an old graph, destroy it first
        self.destroy_inner();
        // THREAD_LOCAL: only "potentially unsafe" CUDA calls (cudaMalloc etc.)
        // made by THIS thread invalidate the capture. With GLOBAL mode, TP rank
        // threads capturing concurrently would poison each other's captures.
        error::check(unsafe {
            ffi::cudaStreamBeginCapture(stream.as_raw(), ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL)
        })
    }

    /// End capture and instantiate the executable graph.
    pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
        error::check(unsafe { ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph) })?;
        error::check(unsafe { ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0) })
    }

    /// Replay the captured graph on `stream`.
    /// Panics if no graph has been captured yet.
    pub fn launch(&self, stream: &CudaStream) -> Result<()> {
        assert!(self.is_ready(), "CudaGraph::launch called before capture");
        error::check(unsafe { ffi::cudaGraphLaunch(self.exec, stream.as_raw()) })
    }

    fn destroy_inner(&mut self) {
        if !self.exec.is_null() {
            unsafe { ffi::cudaGraphExecDestroy(self.exec) };
            self.exec = std::ptr::null_mut();
        }
        if !self.graph.is_null() {
            unsafe { ffi::cudaGraphDestroy(self.graph) };
            self.graph = std::ptr::null_mut();
        }
    }
}

impl Drop for CudaGraph {
    fn drop(&mut self) {
        self.destroy_inner();
    }
}

unsafe impl Send for CudaGraph {}