// SPDX-License-Identifier: GPL-2.0
//
// Copyright (c) 2026 Galih Tama <galpt@v.recipes>
//
// This software may be used and distributed according to the terms of the GNU
// General Public License version 2.

//! NVML poller for active GPU tasks.
//!
//! Polls libnvidia-ml for compute and graphics pids and mirrors the
//! active set into the BPF hash map. Returns None when NVIDIA is
//! absent, so the scheduler continues without GPU awareness.

#![allow(dead_code)]
#![allow(unused_imports)]
#![allow(unused_variables)]
#![allow(clippy::all)]

use std::collections::HashMap;
use std::ffi::CString;
use std::os::fd::{AsFd, AsRawFd};
use std::os::raw::{c_char, c_int, c_uint, c_void};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use anyhow::Result;
use libbpf_rs::libbpf_sys;
use log::{info, warn};

use crate::bpf_intf::{mlfq_gpu_nvml_state, MLFQ_GPU_TRACE_NVML};

const RTLD_NOW: c_int = 2;
const RTLD_LOCAL: c_int = 0;
// RTLD_NODELETE keeps libnvidia-ml.so resident after dlclose for the
// lifetime of the poller thread; without it a dlclose in Drop could
// unload the library while the poller's NvmlPoller still holds function
// pointers, and a subsequent dlopen (e.g. topology banner probe) would
// re-resolve stale symbols. The handle is still dlclose'd in Drop to
// release the refcount; NODELETE only prevents the unload, not the
// close, so no leak (the OS keeps the mapping but not a duplicate).
const RTLD_NODELETE: c_int = 0x01000;
const RTLD_FLAGS: c_int = RTLD_NOW | RTLD_LOCAL | RTLD_NODELETE;

const POLL_MS: u64 = crate::bpf_intf::MLFQ_GPU_NVML_POLL_MS as u64;
const EVICT_MS: u64 = 500;
const NSEC_PER_MSEC: u64 = 1_000_000;

/// Pinned CUDA map path, mirrors MLFQ_CUDA_PIN_PATH from src/bpf/intf.h
/// ("/sys/fs/bpf/mlfq_cuda"). Kept as a Rust const so main.rs and this
/// module share one definition instead of hard-coding the string in two
/// sites (LOW-3). The BPF side defines the same via #define; the two
/// must stay in sync. Bindgen does not emit string macros, so we mirror
/// it here and assert at compile-time via the comment.
pub const MLFQ_CUDA_PIN_PATH: &str = "/sys/fs/bpf/mlfq_cuda";

#[repr(C)]
#[derive(Debug, Copy, Clone)]
struct NvmlProcessInfo {
    pid: c_uint,
    used_gpu_memory: u64,
    gpu_instance_id: c_uint,
    compute_instance_id: c_uint,
}

type NvmlReturn = c_int;
const NVML_SUCCESS: NvmlReturn = 0;

type NvmlInitFn = unsafe extern "C" fn() -> NvmlReturn;
type NvmlShutdownFn = unsafe extern "C" fn() -> NvmlReturn;
type NvmlDeviceGetCountFn = unsafe extern "C" fn(*mut c_uint) -> NvmlReturn;
type NvmlDeviceGetHandleByIndexFn = unsafe extern "C" fn(c_uint, *mut *mut c_void) -> NvmlReturn;
type NvmlDeviceGetComputeRunningProcessesFn =
    unsafe extern "C" fn(*mut c_void, *mut c_uint, *mut NvmlProcessInfo) -> NvmlReturn;
type NvmlDeviceGetGraphicsRunningProcessesFn =
    unsafe extern "C" fn(*mut c_void, *mut c_uint, *mut NvmlProcessInfo) -> NvmlReturn;

extern "C" {
    fn dlopen(filename: *const c_char, flag: c_int) -> *mut c_void;
    fn dlsym(handle: *mut c_void, symbol: *const c_char) -> *mut c_void;
    fn dlclose(handle: *mut c_void) -> c_int;
    fn dlerror() -> *mut c_char;
    fn clock_gettime(clk_id: c_int, tp: *mut Timespec) -> c_int;
}

const CLOCK_MONOTONIC: c_int = 1;

#[repr(C)]
#[derive(Debug, Copy, Clone)]
struct Timespec {
    tv_sec: i64,
    tv_nsec: i64,
}

fn monotonic_ns() -> u64 {
    let mut ts = Timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    let ret = unsafe { clock_gettime(CLOCK_MONOTONIC, &mut ts as *mut _) };
    if ret != 0 {
        return 0;
    }
    (ts.tv_sec as u64)
        .wrapping_mul(1_000_000_000)
        .wrapping_add(ts.tv_nsec as u64)
}

fn dlerror_str() -> String {
    let p = unsafe { dlerror() };
    if p.is_null() {
        return "unknown dlerror".to_string();
    }
    unsafe { std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned() }
}

/// Encapsulated NVML poller, owns the dlopen handle and the BPF map
/// fd. The fd is borrowed from the skeleton; the caller guarantees
/// the BPF skeleton outlives the poller (joined on Drop/shutdown).
pub struct NvmlPoller {
    handle: *mut c_void,
    // NVML fn pointers, None when the symbol is missing.
    nvml_init: Option<NvmlInitFn>,
    nvml_shutdown: Option<NvmlShutdownFn>,
    nvml_device_get_count: Option<NvmlDeviceGetCountFn>,
    nvml_device_get_handle_by_index: Option<NvmlDeviceGetHandleByIndexFn>,
    nvml_device_get_compute_running_processes: Option<NvmlDeviceGetComputeRunningProcessesFn>,
    nvml_device_get_graphics_running_processes: Option<NvmlDeviceGetGraphicsRunningProcessesFn>,
    // BPF map fd for mlfq_gpu_nvml_map.
    map_fd: i32,
    // Last seen per pid, for eviction.
    seen: Arc<Mutex<HashMap<u32, u64>>>,
    backend: String,
}

unsafe impl Send for NvmlPoller {}
unsafe impl Sync for NvmlPoller {}

impl NvmlPoller {
    /// Try to create a poller. Returns None when NVIDIA is absent.
    pub fn try_new(skel: &crate::bpf_skel::BpfSkel) -> Option<Self> {
        let map_fd = skel.maps.mlfq_gpu_nvml_map.as_fd().as_raw_fd();

        let handle = Self::dlopen_nvml()?;
        let (init, shutdown, get_count, get_handle, get_compute, get_graphics) =
            Self::resolve_symbols(handle);

        // At least init and count must be present to be useful; otherwise
        // treat as not detected and close.
        if init.is_none() || get_count.is_none() {
            warn!("NVML library opened but required symbols missing, disabling NVML poller");
            unsafe { dlclose(handle) };
            return None;
        }

        // Try to init NVML now; failure is soft.
        let init_fn = init.unwrap();
        let ret = unsafe { init_fn() };
        if ret != NVML_SUCCESS {
            warn!("nvmlInit failed ({ret}), disabling NVML poller");
            unsafe { dlclose(handle) };
            return None;
        }

        info!(
            "NVML backend detected (libnvidia-ml), poll {}ms evict {}ms",
            POLL_MS, EVICT_MS
        );

        Some(Self {
            handle,
            nvml_init: init,
            nvml_shutdown: shutdown,
            nvml_device_get_count: get_count,
            nvml_device_get_handle_by_index: get_handle,
            nvml_device_get_compute_running_processes: get_compute,
            nvml_device_get_graphics_running_processes: get_graphics,
            map_fd,
            seen: Arc::new(Mutex::new(HashMap::new())),
            backend: "nvml".to_string(),
        })
    }

    fn dlopen_nvml() -> Option<*mut c_void> {
        for name in ["libnvidia-ml.so.1", "libnvidia-ml.so"] {
            let cname = CString::new(name).ok()?;
            let h = unsafe { dlopen(cname.as_ptr(), RTLD_FLAGS) };
            if !h.is_null() {
                return Some(h);
            }
            // Clear error for next try.
            unsafe { dlerror() };
        }
        None
    }

    fn resolve_symbols(
        handle: *mut c_void,
    ) -> (
        Option<NvmlInitFn>,
        Option<NvmlShutdownFn>,
        Option<NvmlDeviceGetCountFn>,
        Option<NvmlDeviceGetHandleByIndexFn>,
        Option<NvmlDeviceGetComputeRunningProcessesFn>,
        Option<NvmlDeviceGetGraphicsRunningProcessesFn>,
    ) {
        unsafe fn sym<T>(handle: *mut c_void, name: &str) -> Option<T> {
            let cname = CString::new(name).ok()?;
            let p = dlsym(handle, cname.as_ptr());
            if p.is_null() {
                None
            } else {
                Some(std::mem::transmute_copy::<*mut c_void, T>(&p))
            }
        }
        unsafe {
            (
                sym(handle, "nvmlInit"),
                sym(handle, "nvmlShutdown"),
                sym(handle, "nvmlDeviceGetCount_v2").or_else(|| sym(handle, "nvmlDeviceGetCount")),
                sym(handle, "nvmlDeviceGetHandleByIndex_v2")
                    .or_else(|| sym(handle, "nvmlDeviceGetHandleByIndex")),
                sym(handle, "nvmlDeviceGetComputeRunningProcesses"),
                sym(handle, "nvmlDeviceGetGraphicsRunningProcesses"),
            )
        }
    }

    /// Backend name for display (topology banner).
    pub fn backend(&self) -> &str {
        &self.backend
    }

    /// Poll once, fixed 100 ms cadence is enforced by the caller.
    /// Poll NVML and update the BPF map. Evicts entries older than
    /// 500ms. The BPF side considers entries inactive after 200ms.
    pub fn poll_once(&self) -> Result<()> {
        let now = monotonic_ns();
        let mut active_pids: Vec<u32> = Vec::new();

        // Query NVML for compute + graphics pids.
        if let (Some(get_count), Some(get_handle)) = (
            self.nvml_device_get_count,
            self.nvml_device_get_handle_by_index,
        ) {
            let mut count: c_uint = 0;
            let ret = unsafe { get_count(&mut count as *mut _) };
            if ret == NVML_SUCCESS {
                for idx in 0..count {
                    let mut dev: *mut c_void = std::ptr::null_mut();
                    let ret = unsafe { get_handle(idx, &mut dev as *mut _) };
                    if ret != NVML_SUCCESS || dev.is_null() {
                        continue;
                    }
                    // Query compute pids. Probe with null buffer first
                    // to get the count, then fetch the list.
                    if let Some(f) = self.nvml_device_get_compute_running_processes {
                        let mut n: c_uint = 0;
                        let _ = unsafe { f(dev, &mut n as *mut _, std::ptr::null_mut()) };
                        // Check n after the call.
                        if n > 0 {
                            let mut buf = vec![
                                NvmlProcessInfo {
                                    pid: 0,
                                    used_gpu_memory: 0,
                                    gpu_instance_id: 0,
                                    compute_instance_id: 0,
                                };
                                n as usize
                            ];
                            let mut n2 = n;
                            let ret2 = unsafe { f(dev, &mut n2 as *mut _, buf.as_mut_ptr()) };
                            if ret2 == NVML_SUCCESS {
                                // n2 may be less than n if races; respect n2.
                                for info in buf.iter().take(n2 as usize) {
                                    if info.pid != 0 {
                                        active_pids.push(info.pid);
                                    }
                                }
                            }
                        }
                    }
                    // Query graphics pids.
                    if let Some(f) = self.nvml_device_get_graphics_running_processes {
                        let mut n: c_uint = 0;
                        let _ = unsafe { f(dev, &mut n as *mut _, std::ptr::null_mut()) };
                        if n > 0 {
                            let mut buf = vec![
                                NvmlProcessInfo {
                                    pid: 0,
                                    used_gpu_memory: 0,
                                    gpu_instance_id: 0,
                                    compute_instance_id: 0,
                                };
                                n as usize
                            ];
                            let mut n2 = n;
                            let ret2 = unsafe { f(dev, &mut n2 as *mut _, buf.as_mut_ptr()) };
                            if ret2 == NVML_SUCCESS {
                                for info in buf.iter().take(n2 as usize) {
                                    if info.pid != 0 {
                                        active_pids.push(info.pid);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        active_pids.sort_unstable();
        active_pids.dedup();

        // Collect updates under lock, then update BPF maps
        // without holding the lock.
        let to_evict: Vec<u32>;
        {
            let mut seen = self.seen.lock().unwrap();
            for pid in &active_pids {
                seen.insert(*pid, now);
            }
            // Evict >500 ms stale: collect under lock, remove from seen,
            // then drop lock before BPF deletes.
            let mut evict = Vec::new();
            for (pid, last) in seen.iter() {
                // last_seen 0 guard already in BPF (mlfq_gpu_nvml_active);
                // userspace also uses wrapping_sub for drift safety.
                if now.wrapping_sub(*last) > EVICT_MS * NSEC_PER_MSEC {
                    evict.push(*pid);
                }
            }
            for pid in &evict {
                seen.remove(pid);
            }
            to_evict = evict;
        }
        // Outside lock: BPF map updates for active pids.
        for pid in &active_pids {
            let state = mlfq_gpu_nvml_state {
                active: 1,
                pad: 0,
                last_seen: now,
            };
            let ret = unsafe {
                libbpf_sys::bpf_map_update_elem(
                    self.map_fd,
                    pid as *const _ as *const c_void,
                    &state as *const _ as *const c_void,
                    0,
                )
            };
            if ret != 0 {
                warn!(
                    "nvml map update pid {pid} failed: {}",
                    std::io::Error::last_os_error()
                );
            }
        }
        // Outside lock: BPF map deletes for evicted.
        for pid in to_evict {
            let ret = unsafe {
                libbpf_sys::bpf_map_delete_elem(self.map_fd, &pid as *const _ as *const c_void)
            };
            if ret != 0 {
                // Best effort; stale entries are harmless due to 200 ms gate.
                let _ = ret;
            }
        }

        // Also update trace mask bit for web UI gauge via BPF BSS?
        // The BSS mask is seeded in main.rs; we keep it alive here by
        // ensuring at least one poll sets it if backend is present.
        // The actual BSS write is done in main.rs after try_new.

        // Ensure the map's entries beyond our seen are also evicted
        // if they were written by a previous run. We rely on the seen
        // eviction; the BPF side's mlfq_gpu_nvml_active already gates
        // on 200 ms, so stale entries are harmless until evicted.

        Ok(())
    }

    /// Fixed interval for the poll thread.
    pub fn poll_interval() -> Duration {
        Duration::from_millis(POLL_MS)
    }

    /// Whether the poller detected an NVIDIA backend.
    pub fn is_detected(&self) -> bool {
        !self.handle.is_null()
    }
}

impl Drop for NvmlPoller {
    fn drop(&mut self) {
        if let Some(f) = self.nvml_shutdown {
            let ret = unsafe { f() };
            if ret != NVML_SUCCESS {
                warn!("nvmlShutdown failed ({ret})");
            }
        }
        if !self.handle.is_null() {
            unsafe { dlclose(self.handle) };
            self.handle = std::ptr::null_mut();
        }
    }
}

/// Helper to check if the trace mask should include NVML without
/// constructing a poller (for probe fallback).
pub fn nvml_library_present() -> bool {
    for name in ["libnvidia-ml.so.1", "libnvidia-ml.so"] {
        if let Ok(cname) = CString::new(name) {
            let h = unsafe { dlopen(cname.as_ptr(), RTLD_FLAGS) };
            if !h.is_null() {
                unsafe { dlclose(h) };
                return true;
            }
            unsafe { dlerror() };
        }
    }
    false
}

/// Detect backend string for topology banner.
pub fn detect_backend() -> Option<String> {
    if nvml_library_present() {
        Some("nvml".to_string())
    } else {
        None
    }
}

// Ensure the trace bit is referenced so the bindgen constant is not
// dead-code warned when the poller is compiled but not used.
#[allow(dead_code)]
fn _assert_trace_bit() {
    let _ = MLFQ_GPU_TRACE_NVML;
}
