#[allow(unused_imports)]
use crate::error::ErrorCode;

use std::{ffi::CString, ptr, sync::Arc};

use singe_cuda_sys::driver;

use crate::{
    device::Device,
    error::{Error, Result},
    jit::JitOptions,
    library::Library,
    module::{Module, ModuleImage},
    nvrtc::{self, CompilationArtifact, OutputKind},
    try_cuda,
    types::Limit,
};

bitflags::bitflags! {
    /// Context creation flags.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct ContextFlags: u32 {
        const SCHEDULE_AUTO = driver::CUctx_flags::CU_CTX_SCHED_AUTO as _;
        const SCHEDULE_SPIN = driver::CUctx_flags::CU_CTX_SCHED_SPIN as _;
        const SCHEDULE_YIELD = driver::CUctx_flags::CU_CTX_SCHED_YIELD as _;
        const SCHEDULE_BLOCKING_SYNC = driver::CUctx_flags::CU_CTX_SCHED_BLOCKING_SYNC as _;
        const MAP_HOST = driver::CUctx_flags::CU_CTX_MAP_HOST as _;
        const LOCAL_MEMORY_RESIZE_TO_MAX = driver::CUctx_flags::CU_CTX_LMEM_RESIZE_TO_MAX as _;
        const COREDUMP_ENABLE = driver::CUctx_flags::CU_CTX_COREDUMP_ENABLE as _;
        const USER_COREDUMP_ENABLE = driver::CUctx_flags::CU_CTX_USER_COREDUMP_ENABLE as _;
        const SYNC_MEMORY_OPERATIONS = driver::CUctx_flags::CU_CTX_SYNC_MEMOPS as _;
    }
}

/// A shared CUDA driver context.
///
/// Unlike cuBLAS, cuDNN, cuFFT, and similar library handles, a CUDA context is
/// the underlying execution environment for a device. It is intended to be
/// shared by streams, modules, libraries, events, allocations, and higher-level
/// library wrappers.
///
/// This type is therefore reference-counted by returning `Arc&lt;Self&gt;` from the
/// constructors, and it remains `Send + Sync`. Shared references do not mutate
/// Rust-visible state on the [`Context`] object itself; methods such as `bind`
/// update the calling thread's current CUDA context in the driver.
///
/// Prefer one long-lived context per device and share it across dependent CUDA
/// objects instead of creating many short-lived contexts.
#[derive(Debug)]
pub struct Context {
    handle: driver::CUcontext,
    device: Device,
}

impl Context {
    pub fn create() -> Result<Arc<Self>> {
        Self::create_with_flags(ContextFlags::empty())
    }

    pub fn create_with_flags(flags: ContextFlags) -> Result<Arc<Self>> {
        let device = Device::current()?;
        Self::create_for_device_with_flags(device, flags)
    }

    pub fn create_for_device(device: Device) -> Result<Arc<Self>> {
        Self::create_for_device_with_flags(device, ContextFlags::empty())
    }

    pub fn create_for_device_with_flags(device: Device, flags: ContextFlags) -> Result<Arc<Self>> {
        unsafe {
            try_cuda!(driver::cuInit(0))?;

            let mut handle = ptr::null_mut();
            try_cuda!(driver::cuCtxCreate_v4(
                &raw mut handle,
                ptr::null_mut(), // CUctxCreateParams
                flags.bits(),
                device.id() as _,
            ))?;

            if handle.is_null() {
                return Err(Error::NullHandle);
            }

            Ok(Arc::new(Self { handle, device }))
        }
    }

    /// Binds this CUDA context to the calling CPU thread.
    ///
    /// The "current context" is thread-local driver state. Calling this method
    /// does not mutate the Rust [`Context`] value itself; it makes this context
    /// current for subsequent CUDA driver and interoperating runtime calls on
    /// the current host thread.
    pub fn bind(&self) -> Result<()> {
        unsafe {
            let mut current_ctx = ptr::null_mut();
            try_cuda!(driver::cuCtxGetCurrent(&raw mut current_ctx))?;
            if current_ctx == self.as_raw() {
                return Ok(());
            }
            try_cuda!(driver::cuCtxSetCurrent(self.as_raw()))?;
        }
        Ok(())
    }

    /// Loads the corresponding module from the given image into the current context.
    /// The image may be a cubin or fatbin as output by **nvcc**, or a NUL-terminated PTX string, either as output by **nvcc** or hand-written, or Tile IR data.
    ///
    /// Note:
    ///
    /// Note that this function may also return error codes from previous, asynchronous launches.
    pub fn load_module(self: &Arc<Self>, image: &ModuleImage<'_>) -> Result<Module> {
        self.bind()?;

        unsafe {
            let mut module_handle = ptr::null_mut();
            try_cuda!(driver::cuModuleLoadData(
                &raw mut module_handle,
                image.as_ptr() as _,
            ))?;
            if module_handle.is_null() {
                return Err(Error::NullHandle);
            }
            Ok(Module::from_raw(module_handle, Arc::clone(self)))
        }
    }

    pub fn unload_module(self: &Arc<Self>, module: Module) -> Result<()> {
        drop(module);
        Ok(())
    }

    /// Loads the corresponding module from the given image into the current context.
    /// The image may be a cubin or fatbin as output by **nvcc**, or a NUL-terminated PTX string, either as output by **nvcc** or hand-written, or Tile IR data.
    ///
    /// Note:
    ///
    /// Note that this function may also return error codes from previous, asynchronous launches.
    pub fn load_module_with_options(
        self: &Arc<Self>,
        image: &ModuleImage<'_>,
        mut jit_options: JitOptions<'_>,
    ) -> Result<Module> {
        self.bind()?;

        let mut jit_options = jit_options.build();
        unsafe {
            let mut module_handle = ptr::null_mut();
            try_cuda!(driver::cuModuleLoadDataEx(
                &raw mut module_handle,
                image.as_ptr() as _,
                jit_options.names.len() as _,
                jit_options.names.as_mut_ptr() as _,
                jit_options.values.as_mut_ptr() as _,
            ))?;
            if module_handle.is_null() {
                return Err(Error::NullHandle);
            }
            Ok(Module::from_raw(module_handle, Arc::clone(self)))
        }
    }

    pub fn load_nvrtc_module(
        self: &Arc<Self>,
        program: &nvrtc::Program,
        output: OutputKind,
    ) -> Result<Module> {
        self.load_nvrtc_module_with_options(program, output, JitOptions::default())
    }

    pub fn load_nvrtc_module_with_options(
        self: &Arc<Self>,
        program: &nvrtc::Program,
        output: OutputKind,
        jit_options: JitOptions<'_>,
    ) -> Result<Module> {
        let image = module_loadable_image(program.artifact(output)?)?;
        self.load_module_with_options(&image, jit_options)
    }

    pub fn load_library(self: &Arc<Self>, image: &ModuleImage<'_>) -> Result<Library> {
        self.load_library_with_options(image, JitOptions::default())
    }

    /// Loads the corresponding library from the given image based on the application defined library loading mode:
    ///
    /// * If module loading is set to EAGER, via the environment variables described in "Module loading", library is loaded eagerly into all contexts at the time of the call and future contexts at the time of creation until the library
    ///    is unloaded with [`sys::cuLibraryUnload`](singe_cuda_sys::driver::cuLibraryUnload).
    /// * If the environment variables are set to LAZY, library is not immediately loaded onto all existent contexts and will only be loaded when a function is needed for that context,
    ///    such as a kernel launch.
    ///
    /// These environment variables are described in the CUDA programming guide under the "CUDA environment variables" section.
    ///
    /// The code may be a cubin or fatbin as output by **nvcc**, or a NUL-terminated PTX string, either as output by **nvcc** or hand-written, or Tile IR data.
    /// A fatbin should also contain relocatable code when doing separate compilation.
    ///
    /// Note:
    ///
    /// If the library contains managed variables and no device in the system supports managed variables this call is expected to return [`ErrorCode::NotSupported`]
    pub fn load_library_with_options(
        self: &Arc<Self>,
        image: &ModuleImage<'_>,
        mut jit_options: JitOptions<'_>,
    ) -> Result<Library> {
        self.bind()?;

        let mut jit_options = jit_options.build();
        let mut handle = ptr::null_mut();
        unsafe {
            try_cuda!(driver::cuLibraryLoadData(
                &raw mut handle,
                image.as_ptr() as _,
                jit_options.names.as_mut_ptr() as _,
                jit_options.values.as_mut_ptr() as _,
                jit_options.names.len() as _,
                ptr::null_mut(),
                ptr::null_mut(),
                0,
            ))?;
        }
        if handle.is_null() {
            return Err(Error::NullHandle);
        }
        Ok(unsafe { Library::from_raw(handle, Arc::clone(self)) })
    }

    pub fn load_nvrtc_library(
        self: &Arc<Self>,
        program: &nvrtc::Program,
        output: OutputKind,
    ) -> Result<Library> {
        self.load_nvrtc_library_with_options(program, output, JitOptions::default())
    }

    pub fn load_nvrtc_library_with_options(
        self: &Arc<Self>,
        program: &nvrtc::Program,
        output: OutputKind,
        jit_options: JitOptions<'_>,
    ) -> Result<Library> {
        let image = library_loadable_image(program.artifact(output)?)?;
        self.load_library_with_options(&image, jit_options)
    }

    /// Loads the corresponding library from the given file based on the application defined library loading mode:
    ///
    /// * If module loading is set to EAGER, via the environment variables described in "Module loading", library is loaded eagerly into all contexts at the time of the call and future contexts at the time of creation until the library
    ///    is unloaded with [`sys::cuLibraryUnload`](singe_cuda_sys::driver::cuLibraryUnload).
    /// * If the environment variables are set to LAZY, library is not immediately loaded onto all existent contexts and will only be loaded when a function is needed for that context,
    ///    such as a kernel launch.
    ///
    /// These environment variables are described in the CUDA programming guide under the "CUDA environment variables" section.
    ///
    /// The file should be a cubin file as output by **nvcc**, or a PTX file either as output by **nvcc** or handwritten, or a fatbin file as output by **nvcc** or hand-written, or Tile IR file.
    /// A fatbin should also contain relocatable code when doing separate compilation.
    ///
    /// Note:
    ///
    /// If the library contains managed variables and no device in the system supports managed variables this call is expected to return [`ErrorCode::NotSupported`]
    pub fn load_library_from_file(self: &Arc<Self>, path: &str) -> Result<Library> {
        self.bind()?;
        let path = CString::new(path)?;
        let mut handle = ptr::null_mut();
        unsafe {
            try_cuda!(driver::cuLibraryLoadFromFile(
                &raw mut handle,
                path.as_ptr(),
                ptr::null_mut(),
                ptr::null_mut(),
                0,
                ptr::null_mut(),
                ptr::null_mut(),
                0,
            ))?;
        }
        if handle.is_null() {
            return Err(Error::NullHandle);
        }
        Ok(unsafe { Library::from_raw(handle, Arc::clone(self)) })
    }

    /// Blocks until the current context has completed all preceding requested tasks.
    /// If the current context is the primary context, child contexts that have been created will also be synchronized.
    /// [`Context::synchronize`] returns an error if one of the preceding tasks failed.
    /// If the context was created with [`ContextFlags::SCHEDULE_BLOCKING_SYNC`], the CPU thread will block until the GPU context has finished its work.
    ///
    /// Note:
    ///
    /// Note that this function may also return error codes from previous, asynchronous launches.
    pub fn synchronize(&self) -> Result<()> {
        self.bind()?;
        unsafe {
            try_cuda!(driver::cuCtxSynchronize())?;
        }
        Ok(())
    }

    /// Returns the flags of the current context.
    /// See `cuCtxCreate_v4` for flag values.
    ///
    /// Note:
    ///
    /// Note that this function may also return error codes from previous, asynchronous launches.
    pub fn flags(&self) -> Result<ContextFlags> {
        self.bind()?;
        unsafe {
            let mut flags = 0;
            try_cuda!(driver::cuCtxGetFlags(&raw mut flags))?;
            Ok(ContextFlags::from_bits_truncate(flags))
        }
    }

    /// Returns the current size of limit.
    /// The supported `CUlimit` values are:
    ///
    /// * [`Limit::StackSize`]: stack size in bytes of each GPU thread.
    /// * [`Limit::PrintfFifoSize`]: size in bytes of the FIFO used by the printf() device system call.
    /// * [`Limit::MallocHeapSize`]: size in bytes of the heap used by the malloc() and free() device system calls.
    /// * [`Limit::DevRuntimeSyncDepth`]: maximum grid depth at which a thread can issue the device runtime call [`Device::synchronize`] to wait on child grid launches to complete.
    /// * [`Limit::DevRuntimePendingLaunchCount`]: maximum number of outstanding device runtime launches that can be made from this context.
    /// * [`Limit::MaxL2FetchGranularity`]: L2 cache fetch granularity.
    /// * [`Limit::PersistingL2CacheSize`]: Persisting L2 cache size in bytes
    ///
    /// Note:
    ///
    /// Note that this function may also return error codes from previous, asynchronous launches.
    pub fn limit(&self, limit: Limit) -> Result<usize> {
        self.bind()?;
        unsafe {
            let mut value = 0;
            try_cuda!(driver::cuCtxGetLimit(&raw mut value, limit.into()))?;
            Ok(value as usize)
        }
    }

    /// Setting limit to value is a request by the application to update the current limit maintained by the context.
    /// The driver is free to modify the requested value to meet h/w requirements (this could be clamping to minimum or maximum values, rounding up to nearest element size, etc).
    /// The application can use [`Context::limit`] to find out exactly what the limit has been set to.
    ///
    /// Setting each `CUlimit` has its own specific restrictions, so each is discussed here.
    ///
    /// * [`Limit::StackSize`] controls the stack size in bytes of each GPU thread.
    ///    The driver automatically increases the per-thread stack size for each
    ///    kernel launch as needed.
    ///    This size isn't reset back to the original value after each launch.
    ///    Setting this value will take
    ///    effect immediately, and if necessary, the device will block until all preceding requested tasks are complete.
    ///
    /// * [`Limit::PrintfFifoSize`] controls the size in bytes of the FIFO used by the printf() device system call.
    ///    Setting [`Limit::PrintfFifoSize`] must be performed before launching any kernel that uses the printf() device system call, otherwise [`ErrorCode::InvalidValue`] will be returned.
    ///
    /// * [`Limit::MallocHeapSize`] controls the size in bytes of the heap used by the malloc() and free() device system calls.
    ///    Setting [`Limit::MallocHeapSize`] must be performed before launching any kernel that uses the malloc() or free() device system calls, otherwise [`ErrorCode::InvalidValue`] will be returned.
    ///
    /// * [`Limit::DevRuntimeSyncDepth`] controls the maximum nesting depth of a grid at which a thread can safely call [`Device::synchronize`].
    ///    Setting this limit must be performed before any launch of a kernel that uses the device runtime and calls [`Device::synchronize`] above the default sync depth, two levels of grids.
    ///    Calls to [`Device::synchronize`] will fail if the limitation is violated.
    ///    This limit can be set smaller than the default or up the maximum launch depth of 24.
    ///    When
    ///    setting this limit, keep in mind that additional levels of sync depth require the driver to reserve large amounts of device
    ///    memory which can no longer be used for user allocations.
    ///    If these reservations of device memory fail, [`Context::set_limit`] will return [`ErrorCode::OutOfMemory`], and the limit can be reset to a lower value.
    ///    This limit is only applicable to devices of compute capability &lt; 9.0.
    ///    Attempting
    ///    to set this limit on devices of other compute capability versions will result in the error [`ErrorCode::UnsupportedLimit`] being returned.
    ///
    /// * [`Limit::DevRuntimePendingLaunchCount`] controls the maximum number of outstanding device runtime launches that can be made from the current context.
    ///    A grid is outstanding
    ///    from the point of launch up until the grid is known to have been completed.
    ///    Device runtime launches which violate this limitation will fail.
    ///    If more pending launches than the default (2048 launches) are needed for a module using the device
    ///    runtime, this limit can be increased.
    ///    Keep in mind that being able to sustain additional pending launches will require the
    ///    driver to reserve larger amounts of device memory upfront which can no longer be used for allocations.
    ///    If these reservations
    ///    fail, [`Context::set_limit`] will return [`ErrorCode::OutOfMemory`], and the limit can be reset to a lower value.
    ///    This limit is only applicable to devices of compute capability 3.5 and higher.
    ///    Attempting to set this limit on devices of compute capability less than 3.5 will result in the error [`ErrorCode::UnsupportedLimit`] being returned.
    ///
    /// * [`Limit::MaxL2FetchGranularity`] controls the L2 cache fetch granularity.
    ///    Values can range from 0B to 128B.
    ///    This is purely a performance hint and it can be
    ///    ignored or clamped depending on the platform.
    ///
    /// * [`Limit::PersistingL2CacheSize`] controls size in bytes available for persisting L2 cache.
    ///    This is purely a performance hint and it can be ignored or clamped
    ///    depending on the platform.
    ///
    /// Note:
    ///
    /// Note that this function may also return error codes from previous, asynchronous launches.
    pub fn set_limit(&self, limit: Limit, value: usize) -> Result<()> {
        self.bind()?;
        unsafe {
            try_cuda!(driver::cuCtxSetLimit(limit.into(), value as _))?;
        }
        Ok(())
    }

    pub const fn device(&self) -> Device {
        self.device
    }

    pub const unsafe fn as_raw(&self) -> driver::CUcontext {
        self.handle
    }
}

unsafe impl Send for Context {}

// CUDA driver contexts are shared execution environments, not per-thread
// library handles. The Rust wrapper only stores the raw context pointer and the
// owning device, while current-context selection is maintained by CUDA as
// thread-local driver state.
unsafe impl Sync for Context {}

impl Drop for Context {
    fn drop(&mut self) {
        unsafe {
            if let Err(err) = try_cuda!(driver::cuCtxDestroy_v2(self.handle)) {
                #[cfg(debug_assertions)]
                eprintln!("failed to destroy CUDA context: {err}");
            }
        }
    }
}

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

impl Eq for Context {}

fn module_loadable_image(artifact: CompilationArtifact) -> Result<ModuleImage<'static>> {
    match artifact {
        CompilationArtifact::Ptx(image) | CompilationArtifact::Cubin(image) => Ok(image),
        CompilationArtifact::LtoIr(_) | CompilationArtifact::OptixIr(_) => Err(Error::InvalidValue),
    }
}

fn library_loadable_image(artifact: CompilationArtifact) -> Result<ModuleImage<'static>> {
    match artifact {
        CompilationArtifact::Ptx(image) | CompilationArtifact::Cubin(image) => Ok(image),
        CompilationArtifact::LtoIr(_) | CompilationArtifact::OptixIr(_) => Err(Error::InvalidValue),
    }
}
