//! NVTX profiling markers for Nsight Systems integration.
//!
//! Enable: add `nvtx` feature to ml-core AND set `FOXHUNT_NVTX=1` at runtime.
//! Without the feature or env var, all calls are zero-cost no-ops.
//!
//! # Architecture
//!
//! When the `nvtx` feature is **disabled** (default), every function in this
//! module is an `#[inline(always)]` no-op that compiles to zero instructions.
//!
//! When the `nvtx` feature is **enabled**, the functions call into
//! `libnvToolsExt.so` via FFI. A runtime check of the `FOXHUNT_NVTX` env var
//! gates the calls so that even a feature-enabled build incurs zero overhead
//! unless explicitly activated. The library is loaded lazily via `dlopen`
//! at first use; if `libnvToolsExt.so` is not present the calls silently
//! degrade to no-ops.
//!
//! # Usage
//!
//! ```rust,ignore
//! use ml_core::nvtx::NvtxRange;
//!
//! fn hot_path() {
//!     let _nvtx = NvtxRange::new("my_kernel_launch");
//!     // ... kernel launch code ...
//! } // range popped automatically on drop
//! ```

// ── Feature-gated FFI implementation ──────────────────────────────────────

#[cfg(feature = "nvtx")]
#[allow(unsafe_code)]
mod ffi {
    use std::ffi::CString;
    use std::sync::OnceLock;

    /// FFI function pointers loaded from libnvToolsExt.so.
    struct NvtxFfi {
        range_push: unsafe extern "C" fn(*const i8) -> i32,
        range_pop: unsafe extern "C" fn() -> i32,
        // Note: we intentionally leak the dlopen handle (process-lifetime library).
    }

    // Safety: the FFI function pointers are thread-safe (NVTX is designed for multi-threaded use).
    unsafe impl Send for NvtxFfi {}
    unsafe impl Sync for NvtxFfi {}

    /// `None` = NVTX disabled or unavailable. `Some(ffi)` = ready to call.
    static NVTX: OnceLock<Option<NvtxFfi>> = OnceLock::new();

    fn try_init() -> Option<NvtxFfi> {
        let enabled = std::env::var("FOXHUNT_NVTX")
            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
            .unwrap_or(false);

        if !enabled {
            return None;
        }

        let lib_name = b"libnvToolsExt.so\0";
        let handle = unsafe {
            libc::dlopen(
                lib_name.as_ptr().cast::<i8>(),
                libc::RTLD_NOW | libc::RTLD_LOCAL,
            )
        };

        if handle.is_null() {
            return None;
        }

        let push_name = b"nvtxRangePushA\0";
        let pop_name = b"nvtxRangePop\0";

        let push_sym = unsafe { libc::dlsym(handle, push_name.as_ptr().cast::<i8>()) };
        let pop_sym = unsafe { libc::dlsym(handle, pop_name.as_ptr().cast::<i8>()) };

        if push_sym.is_null() || pop_sym.is_null() {
            unsafe { libc::dlclose(handle) };
            return None;
        }

        // Safety: function signatures match the NVTX C API exactly.
        // We intentionally leak the handle -- the library stays loaded for the process lifetime.
        Some(NvtxFfi {
            range_push: unsafe { std::mem::transmute(push_sym) },
            range_pop: unsafe { std::mem::transmute(pop_sym) },
        })
    }

    #[inline]
    fn get() -> Option<&'static NvtxFfi> {
        NVTX.get_or_init(try_init).as_ref()
    }

    /// Returns `true` if NVTX profiling is active (feature enabled + env var + library found).
    #[inline]
    pub(super) fn is_enabled() -> bool {
        get().is_some()
    }

    /// Push a named NVTX range (pairs with `range_pop`).
    #[inline]
    pub(super) fn range_push(name: &str) {
        if let Some(ffi) = get() {
            if let Ok(c_name) = CString::new(name) {
                unsafe { (ffi.range_push)(c_name.as_ptr()) };
            }
        }
    }

    /// Pop the most recently pushed NVTX range.
    #[inline]
    pub(super) fn range_pop() {
        if let Some(ffi) = get() {
            unsafe { (ffi.range_pop)() };
        }
    }
}

// ── Public API ────────────────────────────────────────────────────────────

/// RAII guard -- pushes an NVTX range on creation, pops on drop.
///
/// When the `nvtx` feature is disabled (the default), construction and
/// destruction are zero-cost no-ops that optimise away entirely.
#[cfg(feature = "nvtx")]
pub struct NvtxRange {
    _active: bool,
}

#[cfg(not(feature = "nvtx"))]
pub struct NvtxRange;

impl NvtxRange {
    /// Create a new NVTX range with the given label.
    ///
    /// The range is pushed immediately and will be popped when this guard
    /// is dropped. Without the `nvtx` feature or the `FOXHUNT_NVTX=1`
    /// env var, this is a zero-cost no-op.
    #[inline(always)]
    pub fn new(_name: &str) -> Self {
        #[cfg(feature = "nvtx")]
        {
            if ffi::is_enabled() {
                ffi::range_push(_name);
                return Self { _active: true };
            }
            return Self { _active: false };
        }
        #[cfg(not(feature = "nvtx"))]
        {
            Self
        }
    }
}

#[cfg(feature = "nvtx")]
impl Drop for NvtxRange {
    #[inline(always)]
    fn drop(&mut self) {
        if self._active {
            ffi::range_pop();
        }
    }
}

// Manual Debug impl -- simple, no sensitive data.
impl std::fmt::Debug for NvtxRange {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NvtxRange").finish()
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────

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

    #[test]
    fn test_nvtx_range_creation_is_noop() {
        // Without the nvtx feature, this should be a complete no-op.
        // With the feature but without FOXHUNT_NVTX=1, also a no-op.
        let _range = NvtxRange::new("test_range");
        // Just ensure it compiles and doesn't panic.
    }

    #[test]
    fn test_nvtx_range_drop_is_noop() {
        {
            let _range = NvtxRange::new("scoped_range");
        }
        // Dropped without issue.
    }

    #[test]
    fn test_nvtx_range_debug() {
        let range = NvtxRange::new("debug_test");
        let debug_str = format!("{range:?}");
        assert!(debug_str.contains("NvtxRange"));
    }

    #[test]
    fn test_multiple_nested_ranges() {
        let _outer = NvtxRange::new("outer");
        let _inner = NvtxRange::new("inner");
        // Both dropped in reverse order -- correct NVTX push/pop nesting.
    }
}
