//! CUDA backend: cudarc driver API + the sm_120a cubins produced by build.rs (embedded in the binary).
use crate::backend::{Arg, Backend};
use anyhow::{anyhow, Result};
use cudarc::driver::{CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

include!(concat!(env!("OUT_DIR"), "/kernel_index.rs"));

pub struct CudaBackend {
    pub ctx: Arc<CudaContext>,
    pub stream: Arc<CudaStream>,
    modules: HashMap<&'static str, Arc<CudaModule>>,
    funcs: Mutex<HashMap<String, CudaFunction>>,
    allocs: Mutex<HashMap<u64, CudaSlice<u8>>>,
    /// second stream for an independent branch (dense-MLP || router+MoE): fork/use/join through events (also valid inside graph capture)
    stream2: Arc<CudaStream>,
    ev_fork: cudarc::driver::CudaEvent,
    ev_join: cudarc::driver::CudaEvent,
    on_branch: std::sync::atomic::AtomicBool,
}

impl CudaBackend {
    pub fn new(ordinal: usize) -> Result<Arc<Self>> {
        if KERNEL_FILES.is_empty() {
            return Err(anyhow!("dllm-kernels was built without nvcc: no device code available (set NVCC=/path/to/nvcc and rebuild)"));
        }
        let ctx = CudaContext::new(ordinal)?;
        unsafe { ctx.disable_event_tracking() }; // explicit fork/join events only (branch_fork / branch_join)
        // a dedicated (non-default) stream: the legacy default stream cannot be captured into CUDA graphs
        let stream = ctx.new_stream()?;
        let mut modules = HashMap::new();
        for (name, cubin) in KERNEL_FILES {
            modules.insert(*name, ctx.load_module(Ptx::from_binary(cubin.to_vec()))?);
        }
        let stream2 = ctx.new_stream()?;
        let flags = Some(cudarc::driver::sys::CUevent_flags::CU_EVENT_DISABLE_TIMING);
        let (ev_fork, ev_join) = (ctx.new_event(flags)?, ctx.new_event(flags)?);
        Ok(Arc::new(CudaBackend { stream2, ev_fork, ev_join, on_branch: std::sync::atomic::AtomicBool::new(false), ctx, stream, modules, funcs: Mutex::new(HashMap::new()), allocs: Mutex::new(HashMap::new()) }))
    }
    #[inline]
    fn cur(&self) -> &Arc<CudaStream> {
        if self.on_branch.load(std::sync::atomic::Ordering::Relaxed) { &self.stream2 } else { &self.stream }
    }
    fn func(&self, file: &str, sym: &str) -> Result<CudaFunction> {
        let key = format!("{file}::{sym}");
        let mut f = self.funcs.lock().unwrap();
        if let Some(x) = f.get(&key) {
            return Ok(x.clone());
        }
        let m = self.modules.get(file).ok_or_else(|| anyhow!("module {file} not built"))?;
        let func = m.load_function(sym).map_err(|e| anyhow!("{file}::{sym}: {e}"))?;
        f.insert(key, func.clone());
        Ok(func)
    }
}

impl Backend for CudaBackend {
    fn name(&self) -> &'static str {
        "cuda"
    }
    fn alloc(&self, bytes: usize) -> Result<u64> {
        let s = self.stream.alloc_zeros::<u8>(bytes)?;
        let p = s.device_ptr(&self.stream).0;
        self.allocs.lock().unwrap().insert(p, s);
        Ok(p)
    }
    fn free(&self, ptr: u64) {
        self.allocs.lock().unwrap().remove(&ptr);
    }
    fn upload(&self, dst: u64, src: &[u8]) -> Result<()> {
        let mut a = self.allocs.lock().unwrap();
        // dst may point inside an allocation: find the owner
        let (base, s) = a.iter_mut().find(|(b, s)| **b <= dst && dst < **b + s.len() as u64).ok_or_else(|| anyhow!("upload: unknown device pointer"))?;
        let off = (dst - *base) as usize;
        let mut view = s.slice_mut(off..off + src.len());
        self.stream.memcpy_htod(src, &mut view)?;
        Ok(())
    }
    fn download(&self, src: u64, dst: &mut [u8]) -> Result<()> {
        let a = self.allocs.lock().unwrap();
        let (base, s) = a.iter().find(|(b, s)| **b <= src && src < **b + s.len() as u64).ok_or_else(|| anyhow!("download: unknown device pointer"))?;
        let off = (src - *base) as usize;
        let view = s.slice(off..off + dst.len());
        self.stream.memcpy_dtoh(&view, dst)?;
        self.stream.synchronize()?;
        Ok(())
    }
    fn copy_d2d(&self, dst: u64, src: u64, bytes: usize) -> Result<()> {
        unsafe { cudarc::driver::result::memcpy_dtod_async(dst, src, bytes, self.stream.cu_stream())? };
        Ok(())
    }
    fn memset0(&self, dst: u64, bytes: usize) -> Result<()> {
        unsafe { cudarc::driver::result::memset_d8_async(dst, 0, bytes, self.stream.cu_stream())? };
        Ok(())
    }
    fn launch(&self, file: &str, sym: &str, grid: [u32; 3], block: [u32; 3], args: &[Arg]) -> Result<()> {
        let f = self.func(file, sym)?;
        let cfg = LaunchConfig { grid_dim: (grid[0], grid[1], grid[2]), block_dim: (block[0], block[1], block[2]), shared_mem_bytes: 0 };
        let mut b = self.cur().launch_builder(&f);
        for a in args {
            match a {
                Arg::P(v) => b.arg(v),
                Arg::I(v) => b.arg(v),
                Arg::F(v) => b.arg(v),
                Arg::L(v) => b.arg(v),
            };
        }
        unsafe { b.launch(cfg) }.map_err(|e| anyhow!("launch {file}::{sym}: {e}"))?;
        Ok(())
    }
    fn branch_fork(&self) -> Result<()> {
        self.ev_fork.record(&self.stream)?;
        self.stream2.wait(&self.ev_fork)?;
        Ok(())
    }
    fn branch_use(&self, on: bool) {
        self.on_branch.store(on, std::sync::atomic::Ordering::Relaxed);
    }
    fn branch_join(&self) -> Result<()> {
        self.ev_join.record(&self.stream2)?;
        self.stream.wait(&self.ev_join)?;
        Ok(())
    }
    fn sync(&self) -> Result<()> {
        self.stream.synchronize()?;
        Ok(())
    }
    fn graph_begin(&self) -> Result<()> {
        self.stream.begin_capture(cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)?;
        Ok(())
    }
    fn graph_end(&self) -> Result<Box<dyn crate::backend::GraphExec>> {
        let g = self.stream.end_capture(cudarc::driver::sys::CUgraphInstantiate_flags(0))?.ok_or_else(|| anyhow!("empty CUDA graph"))?;
        Ok(Box::new(CudaGraphExec(g)))
    }
}

struct CudaGraphExec(cudarc::driver::CudaGraph);
// the graph is only ever launched from the engine thread that captured it; the handle itself is a plain driver object
unsafe impl Send for CudaGraphExec {}
unsafe impl Sync for CudaGraphExec {}
impl crate::backend::GraphExec for CudaGraphExec {
    fn launch(&self) -> Result<()> {
        self.0.launch()?;
        Ok(())
    }
}