//! Device abstraction. Pointers are plain u64 device addresses; kernels receive them as 8-byte pointer arguments.
use anyhow::Result;
use std::sync::Arc;

#[derive(Debug, Clone, Copy)]
pub enum Arg {
    P(u64),
    I(i32),
    F(f32),
    L(i64),
}

pub trait Backend: Send + Sync {
    fn name(&self) -> &'static str;
    /// zero-initialised, >= 256-byte aligned
    fn alloc(&self, bytes: usize) -> Result<u64>;
    fn free(&self, ptr: u64);
    fn upload(&self, dst: u64, src: &[u8]) -> Result<()>;
    fn download(&self, src: u64, dst: &mut [u8]) -> Result<()>;
    fn copy_d2d(&self, dst: u64, src: u64, bytes: usize) -> Result<()>;
    fn memset0(&self, dst: u64, bytes: usize) -> Result<()>;
    /// `file` = cuda/<file>.cu compilation unit (cubin module), `sym` = extern "C" kernel name
    fn launch(&self, file: &str, sym: &str, grid: [u32; 3], block: [u32; 3], args: &[Arg]) -> Result<()>;
    fn sync(&self) -> Result<()>;
    /// two-branch execution: `branch_fork` (branch stream waits for everything enqueued so far), `branch_use(true)` routes the following launches to the branch stream, `branch_use(false)` back to the main one,
    /// `branch_join` (main waits for the branch). Default: no-ops (everything runs on the one stream, still correct).
    fn branch_fork(&self) -> Result<()> {
        Ok(())
    }
    fn branch_use(&self, _on: bool) {}
    fn branch_join(&self) -> Result<()> {
        Ok(())
    }
    /// CUDA graphs: kernels launched between `graph_begin` and `graph_end` are recorded, not run; `graph_end` returns the replayable executable.
    /// Device pointers / launch parameters are frozen into the graph (metadata buffers must have fixed addresses; contents may change between replays).
    fn graph_begin(&self) -> Result<()> {
        anyhow::bail!("{}: CUDA graph capture is not supported", self.name())
    }
    fn graph_end(&self) -> Result<Box<dyn GraphExec>> {
        anyhow::bail!("{}: CUDA graph capture is not supported", self.name())
    }
}

/// an instantiated CUDA graph
pub trait GraphExec: Send + Sync {
    fn launch(&self) -> Result<()>;
}

/// owned device allocation
pub struct DBuf {
    pub ptr: u64,
    pub bytes: usize,
    be: Arc<dyn Backend>,
}

impl DBuf {
    pub fn zeros(be: &Arc<dyn Backend>, bytes: usize) -> Result<DBuf> {
        Ok(DBuf { ptr: be.alloc(bytes.max(16))?, bytes, be: be.clone() })
    }
    pub fn from_bytes(be: &Arc<dyn Backend>, src: &[u8]) -> Result<DBuf> {
        let b = DBuf::zeros(be, src.len())?;
        be.upload(b.ptr, src)?;
        Ok(b)
    }
    pub fn from_slice<T: Copy>(be: &Arc<dyn Backend>, src: &[T]) -> Result<DBuf> {
        let bytes = unsafe { std::slice::from_raw_parts(src.as_ptr() as *const u8, std::mem::size_of_val(src)) };
        DBuf::from_bytes(be, bytes)
    }
    pub fn upload_slice<T: Copy>(&self, src: &[T]) -> Result<()> {
        let bytes = unsafe { std::slice::from_raw_parts(src.as_ptr() as *const u8, std::mem::size_of_val(src)) };
        anyhow::ensure!(bytes.len() <= self.bytes, "upload {} > buffer {}", bytes.len(), self.bytes);
        self.be.upload(self.ptr, bytes)
    }
    pub fn to_vec<T: Copy + Default>(&self, n: usize) -> Result<Vec<T>> {
        let mut v = vec![T::default(); n];
        let bytes = unsafe { std::slice::from_raw_parts_mut(v.as_mut_ptr() as *mut u8, n * std::mem::size_of::<T>()) };
        anyhow::ensure!(bytes.len() <= self.bytes, "download {} > buffer {}", bytes.len(), self.bytes);
        self.be.download(self.ptr, bytes)?;
        Ok(v)
    }
    pub fn zero(&self) -> Result<()> {
        self.be.memset0(self.ptr, self.bytes)
    }
    /// address `off` bytes into the buffer
    pub fn at(&self, off: usize) -> u64 {
        self.ptr + off as u64
    }
}

impl Drop for DBuf {
    fn drop(&mut self) {
        self.be.free(self.ptr);
    }
}