mod driver;
pub mod nvapi;

use super::{CommonControllerInfo, FanControlHandle, GpuController};
use crate::{
    bindings::nvidia::NvPhysicalGpuHandle,
    server::gpu_controller::{
        NvApi,
        common::{fan_control::FanCurveExt, resolve_process_name},
        nvidia::nvapi::{CLOCK_CLIENT_CLK_VF_POINT_TYPE_PROG, ClockClientClkVfPointInfoV1},
    },
};
use amdgpu_sysfs::{
    gpu_handle::{PowerLevelId, fan_control::FanInfo, power_profile_mode::PowerProfileModesTable},
    hw_mon::Temperature,
};
use anyhow::{Context, anyhow, bail};
use driver::DriverHandle;
use futures::{FutureExt, future::LocalBoxFuture};
use indexmap::IndexMap;
use lact_schema::{
    ActivePowerStates, CacheInfo, ClocksInfo, ClocksTable, ClockspeedStats, DeviceApiInfo,
    DeviceFlag, DeviceInfo, DeviceStats, DeviceType, DrmInfo, DrmMemoryInfo, FanControlMode,
    FanStats, IntelDrmInfo, LinkInfo, NvidiaClockOffset, NvidiaClocksTable, NvidiaThermalInfo,
    NvidiaVfPoint, PmfwInfo, PowerState, PowerStates, PowerStats, ProcessInfo, ProcessList,
    ProcessType, ProcessUtilizationType, TemperatureEntry, VoltageStats, VramStats,
    config::{CurvePoint, FanControlSettings, FanCurve, GpuConfig},
};
use nvml_wrapper::{
    Device, Nvml,
    bitmasks::device::ThrottleReasons,
    enum_wrappers::device::{Clock, PerformanceState, TemperatureSensor, TemperatureThreshold},
    enums::device::{GpuLockedClocksSetting, UsedGpuMemory},
    error::NvmlError,
};
use std::{
    cell::{Cell, RefCell},
    cmp,
    collections::{BTreeMap, HashMap, btree_map::Entry},
    fmt::Write,
    rc::Rc,
    time::{Duration, Instant},
};
use tokio::{select, sync::Notify, time::sleep};
use tracing::{debug, error, trace, warn};

const SUPPORTED_UTIL_TYPES: &[ProcessUtilizationType] = &[
    ProcessUtilizationType::Graphics,
    ProcessUtilizationType::Memory,
    ProcessUtilizationType::Encode,
    ProcessUtilizationType::Decode,
];

pub struct NvidiaGpuController {
    nvml: &'static Nvml,
    common: CommonControllerInfo,
    fan_control_handle: RefCell<Option<FanControlHandle>>,
    initial_target_temp: Option<u32>,

    nvapi: Option<(&'static NvApi, NvPhysicalGpuHandle)>,
    driver_handle: Option<DriverHandle>,
    nvapi_thermals_mask: Option<i32>,

    last_util_timestamp: Cell<Option<u64>>,
    // Store last applied offsets as a workaround when the driver doesn't tell us the current offset
    last_applied_offsets: RefCell<HashMap<Clock, HashMap<PerformanceState, i32>>>,
    last_applied_gpu_locked_clocks: RefCell<Option<(u32, u32)>>,
    last_applied_vram_locked_clocks: RefCell<Option<(u32, u32)>>,
    // Check if reset is needed to avoid unnecessarily going to nvapi
    vf_curve_written: Cell<bool>,
    // Used as the initial value on cards which do not report base VF points themselves (Turing)
    base_vf_curve: RefCell<Option<Vec<NvidiaVfPoint>>>,
}

impl NvidiaGpuController {
    pub fn new(
        common: CommonControllerInfo,
        nvml: &'static Nvml,
        nvapi: Option<&'static NvApi>,
    ) -> anyhow::Result<Self> {
        let device = nvml
            .device_by_pci_bus_id(common.pci_slot_name.as_str())
            .with_context(|| {
                format!(
                    "Could not get PCI device '{}' from NVML",
                    common.pci_slot_name
                )
            })?;

        let (nvapi_handle, nvapi_thermals_mask) = match nvapi.as_ref() {
            Some(nvapi) => {
                let bus_id = common.get_slot_info()?.bus;
                let gpu_handle = nvapi
                    .find_matching_gpu(u32::from(bus_id))
                    .inspect_err(|err| error!("Could not get NvAPI GPU handle: {err}"))
                    .ok()
                    .flatten();

                let thermals_mask = gpu_handle.and_then(|handle| unsafe {
                    nvapi
                        .calculate_thermals_mask(handle)
                        .inspect(|mask| {
                            debug!("calculated NvAPI thermals mask {mask:x}");
                        })
                        .inspect_err(|err| {
                            error!("could not calculate NvAPI thermal mask: {err:#}");
                        })
                        .ok()
                });

                (gpu_handle, thermals_mask)
            }
            None => (None, None),
        };
        debug!("initialized NvAPI device handle {nvapi_handle:?}");

        let minor_number = device.minor_number()?;

        let driver_handle = match DriverHandle::open(minor_number) {
            Ok(handle) => {
                debug!("opened Nvidia driver handle");
                Some(handle)
            }
            Err(err) => {
                error!("could not get Nvidia driver handle: {err:#}");
                None
            }
        };

        let target_temp = device
            .temperature_threshold(TemperatureThreshold::AcousticCurr)
            .ok();

        Ok(Self {
            nvml,
            nvapi: nvapi.zip(nvapi_handle),
            common,
            driver_handle,
            nvapi_thermals_mask,
            initial_target_temp: target_temp,
            last_util_timestamp: Cell::new(None),
            fan_control_handle: RefCell::new(None),
            last_applied_offsets: RefCell::new(HashMap::new()),
            last_applied_gpu_locked_clocks: RefCell::new(None),
            last_applied_vram_locked_clocks: RefCell::new(None),
            vf_curve_written: Cell::new(false),
            base_vf_curve: RefCell::new(None),
        })
    }

    fn device(&self) -> Device<'_> {
        self.nvml
            .device_by_pci_bus_id(self.common.pci_slot_name.as_str())
            .expect("Can no longer get device")
    }

    fn get_nvidia_thermal_info(&self) -> NvidiaThermalInfo {
        NvidiaThermalInfo {
            target_temp: self.get_target_temp(),
            target_temp_default: self.initial_target_temp,
        }
    }

    fn get_target_temp(&self) -> Option<FanInfo> {
        let device = self.device();
        let current = device
            .temperature_threshold(TemperatureThreshold::AcousticCurr)
            .ok()?;
        let min = device
            .temperature_threshold(TemperatureThreshold::AcousticMin)
            .ok()?;
        let max = device
            .temperature_threshold(TemperatureThreshold::AcousticMax)
            .ok()?;

        Some(FanInfo {
            current,
            allowed_range: Some((min, max)),
        })
    }

    async fn start_curve_fan_control_task(
        &self,
        curve: FanCurve,
        settings: FanControlSettings,
    ) -> anyhow::Result<()> {
        // Stop existing task to re-apply new curve
        self.stop_fan_control().await?;

        let device = self.device();
        device
            .temperature(TemperatureSensor::Gpu)
            .context("Could not read temperature")?;

        let fan_count = device.num_fans().context("Could not read fan count")?;
        if fan_count == 0 {
            return Err(anyhow!("Device has no fans"));
        }

        let mut notify_guard = self
            .fan_control_handle
            .try_borrow_mut()
            .map_err(|err| anyhow!("Lock error: {err}"))?;

        let notify = Rc::new(Notify::new());
        let task_notify = notify.clone();

        let nvml = self.nvml;
        let pci_slot_id = self.common.pci_slot_name.clone();
        debug!("spawning new fan control task");

        let handle = tokio::task::spawn_local(async move {
            let mut device = nvml
                .device_by_pci_bus_id(pci_slot_id.as_str())
                .expect("Can no longer get device");

            let mut last_pwm = (None, Instant::now());
            let mut last_temp = 0;

            let interval = Duration::from_millis(settings.interval_ms);
            let spindown_delay = Duration::from_millis(settings.spindown_delay_ms.unwrap_or(0));
            #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
            let change_threshold = settings.change_threshold.unwrap_or(0) as i32;
            #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
            let auto_threshold = settings.auto_threshold.unwrap_or(0) as i32;

            let mut manual_mode = true;

            loop {
                select! {
                    () = sleep(interval) => (),
                    () = task_notify.notified() => break,
                }

                #[allow(clippy::cast_possible_wrap)]
                let current_temp = device
                    .temperature(TemperatureSensor::Gpu)
                    .expect("Could not read temperature") as i32;

                if (last_temp - current_temp).abs() < change_threshold {
                    trace!(
                        "temperature changed from {last_temp}°C to {current_temp}°C, which is less than the {change_threshold}°C threshold, skipping speed adjustment"
                    );
                    continue;
                }

                if current_temp < auto_threshold {
                    if manual_mode {
                        trace!("temperature below auto threshold, setting fan policy to auto");
                        for fan in 0..fan_count {
                            if let Err(err) = device.set_default_fan_speed(fan) {
                                error!(
                                    "could not set fan speed to auto: {err}, disabling fan control"
                                );
                                break;
                            }
                        }

                        manual_mode = false;
                    } else {
                        trace!("temperature below auto threshold, skipping control");
                    }
                    continue;
                }

                let target_pwm = curve.pwm_at_temp(Temperature {
                    #[allow(clippy::cast_precision_loss)]
                    current: Some(current_temp as f32),
                    crit: None,
                    crit_hyst: None,
                });
                let now = Instant::now();

                if let (Some(previous_pwm), previous_timestamp) = last_pwm {
                    let diff = now - previous_timestamp;
                    if target_pwm < previous_pwm && diff < spindown_delay {
                        trace!(
                            "delaying fan spindown ({}ms left)",
                            spindown_delay.checked_sub(diff).unwrap().as_millis()
                        );
                        continue;
                    }
                }

                last_pwm = (Some(target_pwm), now);
                last_temp = current_temp;

                trace!("fan control tick: setting pwm to {target_pwm}");

                for fan in 0..fan_count {
                    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
                    if let Err(err) =
                        device.set_fan_speed(fan, (f64::from(target_pwm) / 2.55) as u32)
                    {
                        error!("could not set fan speed: {err}, disabling fan control");
                        break;
                    }
                }
                manual_mode = true;
            }
            debug!("exited fan control task");
        });

        *notify_guard = Some((notify, handle));

        debug!(
            "started fan control with interval {}ms",
            settings.interval_ms
        );

        Ok(())
    }

    async fn stop_fan_control(&self) -> anyhow::Result<()> {
        let mut fail_on_error = false;

        let maybe_notify = self
            .fan_control_handle
            .try_borrow_mut()
            .map_err(|err| anyhow!("Lock error: {err}"))?
            .take();
        if let Some((notify, handle)) = maybe_notify {
            notify.notify_one();
            handle.await?;
            fail_on_error = true;
        }

        let mut device = self.device();
        let fan_count = device.num_fans().context("Could not get fan count")?;
        for i in 0..fan_count {
            if let Err(err) = device
                .set_default_fan_speed(i)
                .context("Could not reset fan control to default")
            {
                if fail_on_error {
                    return Err(err);
                }
                error!("{err:#?}");
            }
        }

        Ok(())
    }

    fn try_get_power_states(&self) -> anyhow::Result<PowerStates> {
        let device = self.device();

        let supported_states = device
            .supported_performance_states()
            .context("Could not get supported pstates")?;

        let mut power_states = PowerStates::default();

        for pstate in supported_states {
            let (gpu_min, gpu_max) = device
                .min_max_clock_of_pstate(Clock::Graphics, pstate)
                .context("Could not read GPU pstates")?;

            power_states.core.push(PowerState {
                enabled: true,
                min_value: Some(u64::from(gpu_min)),
                value: u64::from(gpu_max),
                id: Some(PowerLevelId::Index(
                    pstate
                        .as_c()
                        .try_into()
                        .expect("Power state always fits in u8"),
                )),
            });

            let (mem_min, mem_max) = device
                .min_max_clock_of_pstate(Clock::Memory, pstate)
                .context("Could not read memory pstates")?;

            power_states.vram.push(PowerState {
                enabled: true,
                min_value: Some(u64::from(mem_min)),
                value: u64::from(mem_max),
                id: Some(PowerLevelId::Index(
                    pstate
                        .as_c()
                        .try_into()
                        .expect("Power state always fits in u8"),
                )),
            });
        }

        Ok(power_states)
    }

    fn get_vf_curve(&self) -> anyhow::Result<Vec<NvidiaVfPoint>> {
        if let Some((nvapi, handle)) = self.nvapi.as_ref() {
            let info;
            let status;

            unsafe {
                info = nvapi.clock_client_clk_vf_points_get_info(*handle)?;
                status =
                    nvapi.clock_client_clk_vf_points_get_status(*handle, info.vf_points_mask)?;
            }

            let base_vf_curve = self.base_vf_curve.borrow();
            let point_count = point_count_from_mask(info.vf_points_mask);
            let mut curve = Vec::with_capacity(point_count);

            for i in 0..point_count {
                let point = status.vf_points[i];
                let point_info = info.vf_points[i];

                // Only report configurable and voltage-based points
                if !vf_curve_point_is_editable(point_info) {
                    continue;
                }

                let (base_freq, base_voltage) = if status.b_vf_tuple_base_supported == 0 {
                    // If vf_tuple_base is not supported, base_vf_curve must be populated before any write
                    // Otherwise, current values are base
                    if let Some(base_curve) = &*base_vf_curve {
                        let base_point = base_curve
                            .iter()
                            .find(|point| point.index as usize == i)
                            .expect("Mismatched base point");
                        (base_point.base_freq, base_point.base_voltage)
                    } else {
                        (point.freq_khz / 1000, point.voltage_uv / 1000)
                    }
                } else {
                    (
                        point.vf_tuple_base.freq_khz / 1000,
                        point.vf_tuple_base.voltage_uv / 1000,
                    )
                };

                curve.push(NvidiaVfPoint {
                    index: u8::try_from(i).expect("max 255 points"),
                    freq: point.freq_khz / 1000,
                    voltage: point.voltage_uv / 1000,
                    base_freq,
                    base_voltage,
                });
            }

            Ok(curve)
        } else {
            Err(anyhow!("NvAPI not available"))
        }
    }

    #[expect(clippy::similar_names)]
    fn apply_vf_curve(&self, curve: &IndexMap<u8, CurvePoint>) -> anyhow::Result<()> {
        let (nvapi, handle) = self.nvapi.as_ref().context("NvAPI not available")?;

        debug!("applying curve with {} points", curve.len());

        let info = unsafe { nvapi.clock_client_clk_vf_points_get_info(*handle)? };
        let current_vf_curve =
            unsafe { nvapi.clock_client_clk_vf_points_get_status(*handle, info.vf_points_mask)? };
        let mut curve_control =
            unsafe { nvapi.clock_client_clk_vf_get_control(*handle, info.vf_points_mask)? };

        let offset_info = self
            .device()
            .clock_offset(Clock::Graphics, PerformanceState::Zero)
            .context("Could not get offset info")?;

        // If a base curve is needed, we should get a clean unmodified version
        if current_vf_curve.b_vf_tuple_base_supported == 0 && self.base_vf_curve.borrow().is_none()
        {
            self.reset_vf_curve().context("Could not reset curve")?;
            self.base_vf_curve.replace(Some(self.get_vf_curve()?));

            return self.apply_vf_curve(curve);
        }

        let base_vf_curve = self.base_vf_curve.borrow();

        for (i, configured_point) in curve {
            let i = *i as usize;

            let point_info = info.vf_points[i];
            let current_point = current_vf_curve.vf_points[i];
            let point_control = &mut curve_control.vf_points[i];

            if !vf_curve_point_is_editable(point_info) {
                bail!("Trying to edit non-configurable point {i}");
            }

            if let Some(configured_mv) = configured_point.voltage {
                let current_mv = (current_point.voltage_uv / 1000).cast_signed();
                if configured_mv != current_mv {
                    bail!(
                        "Voltage is immutable - point {i} was attempted to be changed from {current_mv}mV to {configured_mv}mV"
                    );
                }
            }

            if let Some(configured_mhz) = configured_point.clockspeed {
                let base_freq: i32 = if let Some(base_points) = &*base_vf_curve {
                    let base_point = base_points
                        .iter()
                        .find(|point| point.index as usize == i)
                        .expect("Mismatched base point");
                    base_point.base_freq.cast_signed() * 1000
                } else {
                    current_point.vf_tuple_base.freq_khz.try_into()?
                };
                let offset_khz = configured_mhz * 1000 - base_freq;
                let offset_mhz = offset_khz / 1000;

                let min_offset = cmp::min(offset_info.min_clock_offset_mhz, -(base_freq / 1000));

                if !(min_offset..=offset_info.max_clock_offset_mhz).contains(&offset_mhz) {
                    bail!("Configured offset {offset_mhz}MHz is outside of the allowed range",);
                }

                debug!("writing offset {offset_khz}KHz to point {i}");

                point_control.data.prog.freq_offset_khz = offset_khz;
            }
        }

        unsafe {
            nvapi.clock_client_clk_vf_set_control(*handle, curve_control)?;
        }

        self.vf_curve_written.set(true);

        Ok(())
    }

    fn reset_vf_curve(&self) -> anyhow::Result<()> {
        let (nvapi, handle) = self.nvapi.as_ref().context("NvAPI not available")?;

        let info = unsafe { nvapi.clock_client_clk_vf_points_get_info(*handle)? };
        let mut curve_control =
            unsafe { nvapi.clock_client_clk_vf_get_control(*handle, info.vf_points_mask)? };

        for i in 0..point_count_from_mask(info.vf_points_mask) {
            let point_info = info.vf_points[i];
            if vf_curve_point_is_editable(point_info) {
                curve_control.vf_points[i].data.prog.freq_offset_khz = 0;
            }
        }

        unsafe {
            nvapi.clock_client_clk_vf_set_control(*handle, curve_control)?;
        }

        Ok(())
    }

    fn reset_target_temp(&self) -> anyhow::Result<()> {
        if let Some(initial) = self.initial_target_temp {
            let device = self.device();

            let current = device.temperature_threshold(TemperatureThreshold::AcousticCurr)?;

            if current != initial {
                debug!("resetting target temperature to {initial}");
                device.set_temperature_threshold(
                    TemperatureThreshold::AcousticCurr,
                    initial.cast_signed(),
                )?;
            }
        } else {
            debug!("no initial target temperature was read, skipping reset");
        }

        Ok(())
    }
}

fn vf_curve_point_is_editable(point: ClockClientClkVfPointInfoV1) -> bool {
    point.b_voltage_based == 1 && point.type_ == CLOCK_CLIENT_CLK_VF_POINT_TYPE_PROG
}

fn point_count_from_mask(mask: [u32; 8]) -> usize {
    let count: usize = mask.iter().map(|i| i.count_ones() as usize).sum();
    assert!(u8::try_from(count).is_ok());
    count
}

impl GpuController for NvidiaGpuController {
    fn controller_info(&self) -> &CommonControllerInfo {
        &self.common
    }

    fn device_type(&self) -> DeviceType {
        // No clue what happens on Tegra chips
        DeviceType::Dedicated
    }

    fn friendly_name(&self) -> Option<String> {
        self.device()
            .name()
            .ok()
            .or_else(|| self.common.pci_info.device_pci_info.model.clone())
    }

    fn get_info(
        &self,
        unique_vendor: bool,
        include_api_info: bool,
    ) -> LocalBoxFuture<'_, DeviceInfo> {
        Box::pin(async move {
            let api_info = if include_api_info {
                self.get_api_info(unique_vendor).await
            } else {
                DeviceApiInfo::default()
            };

            let device = self.device();
            let driver_handle = self.driver_handle.as_ref();

            DeviceInfo {
                pci_info: Some(self.common.pci_info.clone()),
                api_info,
                driver: format!(
                    "nvidia {}",
                    self.nvml.sys_driver_version().unwrap_or_default()
                ), // NVML should always be "nvidia"
                vbios_version: device
                    .vbios_version()
                    .map_err(|err| error!("could not get VBIOS version: {err}"))
                    .ok(),
                link_info: LinkInfo {
                    current_width: device.current_pcie_link_width().map(|v| v.to_string()).ok(),
                    current_speed: device
                        .pcie_link_speed()
                        .map(|v| {
                            let mut output = format!("{} GT/s", v / 1000);
                            if let Ok(link_gen) = device.current_pcie_link_gen() {
                                let _ = write!(output, " Gen {link_gen}");
                            }
                            output
                        })
                        .ok(),
                    max_width: device.max_pcie_link_width().map(|v| v.to_string()).ok(),
                    max_speed: device
                        .max_pcie_link_speed()
                        .ok()
                        .and_then(|v| v.as_integer())
                        .map(|v| {
                            let mut output = format!("{} GT/s", v / 1000);
                            if let Ok(link_gen) = device.max_pcie_link_gen() {
                                let _ = write!(output, " Gen {link_gen}");
                            }
                            output
                        }),
                },
                drm_info: Some(DrmInfo {
                    device_name: device.name().ok(),
                    pci_revision_id: None,
                    family_name: device.architecture().map(|arch| arch.to_string()).ok(),
                    family_id: None,
                    asic_name: None,
                    chip_class: device.architecture().map(|arch| arch.to_string()).ok(),
                    compute_units: None,
                    streaming_multiprocessors: driver_handle
                        .and_then(|handle| handle.get_sm_count().ok()),
                    cuda_cores: device.num_cores().ok(),
                    vram_type: driver_handle
                        .and_then(|handle| handle.get_ram_type().ok())
                        .map(str::to_owned),
                    vram_clock_ratio: 1.0,
                    isa: None,
                    vram_vendor: driver_handle
                        .and_then(|handle| handle.get_ram_vendor().ok())
                        .map(str::to_owned),
                    vram_bit_width: driver_handle.and_then(|handle| handle.get_bus_width().ok()),
                    vram_max_bw: None,
                    cache_info: driver_handle
                        .and_then(|handle| handle.get_l2_cache_size().ok())
                        .map(|size| CacheInfo::Nvidia { l2: size }),
                    rop_info: driver_handle
                        .as_ref()
                        .and_then(|handle| handle.get_rop_info().ok()),
                    memory_info: device
                        .bar1_memory_info()
                        .map(|bar_info| DrmMemoryInfo {
                            cpu_accessible_used: bar_info.used,
                            cpu_accessible_total: bar_info.total,
                            resizeable_bar: device
                                .memory_info()
                                .ok()
                                .map(|memory_info| bar_info.total >= memory_info.total),
                        })
                        .ok(),
                    amd_ip_info: vec![],
                    intel: IntelDrmInfo::default(),
                }),
                flags: vec![
                    DeviceFlag::ConfigurableFanControl,
                    DeviceFlag::AutoFanThreshold,
                ],
            }
        })
    }

    #[allow(
        clippy::cast_precision_loss,
        clippy::cast_possible_truncation,
        clippy::cast_sign_loss,
        clippy::too_many_lines
    )]
    fn get_stats(&self, gpu_config: Option<&GpuConfig>) -> DeviceStats {
        let device = self.device();

        let mut temps = HashMap::new();

        if let Ok(temp) = device.temperature(TemperatureSensor::Gpu) {
            let crit = device
                .temperature_threshold(TemperatureThreshold::Shutdown)
                .map(|value| value as f32)
                .ok();

            let value = Temperature {
                current: Some(temp as f32),
                crit,
                crit_hyst: None,
            };

            temps.insert(
                "GPU".to_owned(),
                TemperatureEntry {
                    value,
                    primary: true,
                    display_only: false,
                },
            );
        }

        let mut voltage = None;

        if let Some((nvapi, handle)) = self.nvapi.as_ref() {
            unsafe {
                if let Some(mask) = self.nvapi_thermals_mask
                    && let Ok(thermals) = nvapi.get_thermals(*handle, mask)
                {
                    if let Some(hotspot) = thermals.hotspot() {
                        temps.insert(
                            "GPU Hotspot".to_owned(),
                            TemperatureEntry {
                                value: Temperature {
                                    current: Some(hotspot as f32),
                                    crit: None,
                                    crit_hyst: None,
                                },
                                primary: true,
                                display_only: true,
                            },
                        );
                    }

                    if let Some(vram) = thermals.vram() {
                        temps.insert(
                            "VRAM".to_owned(),
                            TemperatureEntry {
                                value: Temperature {
                                    current: Some(vram as f32),
                                    crit: None,
                                    crit_hyst: None,
                                },
                                primary: true,
                                display_only: true,
                            },
                        );
                    }
                }

                if let Ok(value) = nvapi.get_voltage(*handle) {
                    voltage = Some(u64::from(value) / 1000);
                }
            }
        }

        let fan_settings = gpu_config.and_then(|config| config.fan_control_settings.as_ref());

        let num_fans = device.num_fans().unwrap_or(0);

        let (pwm_current, speed_current) = if num_fans == 0 {
            (None, None)
        } else {
            let fan_speeds = (0..num_fans)
                .flat_map(|idx| device.fan_speed(idx))
                .collect::<Vec<_>>();

            let pwm_current = if fan_speeds.is_empty() {
                None
            } else {
                let avg_speed: u32 = fan_speeds.iter().sum::<u32>() / fan_speeds.len() as u32;
                Some((f64::from(avg_speed) * 2.55) as u8)
            };

            let fan_speeds_rpm = (0..num_fans)
                .flat_map(|idx| device.fan_speed_rpm(idx))
                .collect::<Vec<_>>();

            let speed_current = if fan_speeds_rpm.is_empty() {
                None
            } else {
                Some(fan_speeds_rpm.iter().sum::<u32>() / fan_speeds_rpm.len() as u32)
            };

            (pwm_current, speed_current)
        };

        let vram = device
            .memory_info()
            .map(|info| VramStats {
                total: Some(info.total),
                used: Some(info.used),
            })
            .unwrap_or_default();

        let active_pstate = device
            .performance_state()
            .map(|pstate| {
                PowerLevelId::Index(
                    pstate
                        .as_c()
                        .try_into()
                        .expect("Power state always fits in u8"),
                )
            })
            .ok();

        let fan_range = device.min_max_fan_speed().ok();

        DeviceStats {
            temps,
            fan: FanStats {
                control_enabled: gpu_config.is_some_and(|config| config.fan_control_enabled),
                control_mode: fan_settings.map(|settings| settings.mode),
                static_speed: fan_settings.map(|settings| settings.static_speed),
                curve: fan_settings.map(|settings| settings.curve.0.clone()),
                spindown_delay_ms: fan_settings.and_then(|settings| settings.spindown_delay_ms),
                change_threshold: fan_settings.and_then(|settings| settings.change_threshold),
                auto_threshold: fan_settings.and_then(|settings| settings.auto_threshold),
                temperature_key: None,
                speed_current,
                speed_max: None,
                speed_min: None,
                pwm_current,
                pwm_max: fan_range.map(|(_, max)| (f64::from(max) * 2.55).round() as u32),
                pwm_min: fan_range.map(|(min, _)| (f64::from(min) * 2.55).round() as u32),
                temperature_range: None,
                pmfw_info: PmfwInfo::default(),
            },
            nvidia_thermal_info: self.get_nvidia_thermal_info(),
            power: PowerStats {
                average: None,
                current: device.power_usage().map(|mw| f64::from(mw) / 1000.0).ok(),
                cap_current: device
                    .power_management_limit()
                    .map(|mw| f64::from(mw) / 1000.0)
                    .ok(),
                cap_max: device
                    .power_management_limit_constraints()
                    .map(|constraints| f64::from(constraints.max_limit) / 1000.0)
                    .ok(),
                cap_min: device
                    .power_management_limit_constraints()
                    .map(|constraints| f64::from(constraints.min_limit) / 1000.0)
                    .ok(),
                cap_default: device
                    .power_management_limit_default()
                    .map(|mw| f64::from(mw) / 1000.0)
                    .ok(),
                sensors: HashMap::new(),
            },
            busy_percent: device
                .utilization_rates()
                .map(|utilization| u8::try_from(utilization.gpu).expect("Invalid percentage"))
                .ok(),
            vram,
            clockspeed: ClockspeedStats {
                gpu_clockspeed: device.clock_info(Clock::Graphics).map(Into::into).ok(),
                vram_clockspeed: device.clock_info(Clock::Memory).map(Into::into).ok(),
                target_gpu_clockspeed: None,
                sensors: [
                    ("SM", device.clock_info(Clock::SM)),
                    ("Video", device.clock_info(Clock::Video)),
                ]
                .into_iter()
                .filter_map(|(label, result)| Some((label.to_owned(), result.ok()?.into())))
                .collect(),
            },
            throttle_info: device.current_throttle_reasons().ok().map(|reasons| {
                reasons
                    .iter()
                    .filter(|reason| *reason != ThrottleReasons::GPU_IDLE)
                    .map(|reason| {
                        let mut name = String::new();
                        bitflags::parser::to_writer(&reason, &mut name).unwrap();
                        (name, vec![])
                    })
                    .collect()
            }),
            voltage: VoltageStats {
                gpu: voltage,
                ..Default::default()
            },
            performance_level: None,
            active_power_states: active_pstate.map(|active_pstate| ActivePowerStates {
                core: Some(active_pstate),
                memory: Some(active_pstate),
                pcie: None,
            }),
        }
    }

    #[allow(clippy::cast_possible_wrap)]
    fn get_clocks_info(&self, _gpu_config: Option<&GpuConfig>) -> anyhow::Result<ClocksInfo> {
        let device = self.device();

        let mut gpu_offsets = IndexMap::new();
        let mut mem_offsets = IndexMap::new();
        let mut gpu_clock_range = None;
        let mut vram_clock_range = None;

        let supported_pstates = device.supported_performance_states()?;

        let mut clock_types = [
            (Clock::Graphics, &mut gpu_offsets, &mut gpu_clock_range),
            (Clock::Memory, &mut mem_offsets, &mut vram_clock_range),
        ];

        for pstate in supported_pstates.iter().rev() {
            for (clock_type, offsets, clock_range) in &mut clock_types {
                if let Ok(offset) = device.clock_offset(*clock_type, *pstate) {
                    let mut offset = NvidiaClockOffset {
                        current: offset.clock_offset_mhz,
                        min: offset.min_clock_offset_mhz,
                        max: offset.max_clock_offset_mhz,
                    };

                    // On some driver versions, the applied offset values are not reported.
                    // In these scenarios we must store them manually for reporting.
                    if offset.current == 0
                        && let Some(applied_offsets) =
                            self.last_applied_offsets.borrow().get(clock_type)
                        && let Some(applied_offset) = applied_offsets.get(pstate)
                    {
                        offset.current = *applied_offset;
                    }

                    offsets.insert(pstate.as_c(), offset);
                }

                // Find min and max clockspeeds from any pstate
                if let Ok((pstate_min, pstate_max)) =
                    device.min_max_clock_of_pstate(*clock_type, *pstate)
                {
                    match clock_range {
                        Some((current_min, current_max)) => {
                            *current_min = cmp::min(*current_min, pstate_min);
                            *current_max = cmp::max(*current_max, pstate_max);
                        }
                        None => {
                            **clock_range = Some((pstate_min, pstate_max));
                        }
                    }
                }
            }
        }

        let gpu_vf_curve = self
            .get_vf_curve()
            .inspect_err(|err| warn!("could not get VF curve: {err:#}"))
            .unwrap_or_default();

        let table = NvidiaClocksTable {
            gpu_offsets,
            mem_offsets,
            gpu_locked_clocks: *self.last_applied_gpu_locked_clocks.borrow(),
            vram_locked_clocks: *self.last_applied_vram_locked_clocks.borrow(),
            gpu_clock_range,
            vram_clock_range,
            gpu_vf_curve,
        };

        Ok(ClocksInfo {
            table: Some(ClocksTable::Nvidia(table)),
            ..Default::default()
        })
    }

    fn get_power_states(&self, _gpu_config: Option<&GpuConfig>) -> PowerStates {
        self.try_get_power_states().unwrap_or_else(|err| {
            warn!("could not get pstates info: {err:#}");
            PowerStates::default()
        })
    }

    fn get_power_profile_modes(&self) -> anyhow::Result<PowerProfileModesTable> {
        Err(anyhow!("Not supported on Nvidia"))
    }

    fn vbios_dump(&self) -> anyhow::Result<Vec<u8>> {
        Err(anyhow!("Not supported on Nvidia"))
    }

    #[allow(clippy::cast_possible_wrap, clippy::cast_sign_loss)]
    fn apply_config<'a>(&'a self, config: &'a GpuConfig) -> LocalBoxFuture<'a, anyhow::Result<()>> {
        Box::pin(async {
            let mut device = self.device();

            if let Some(cap) = config.power_cap {
                #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
                let cap = (cap * 1000.0) as u32;

                let current_cap = device
                    .power_management_limit()
                    .context("Could not get current cap")?;

                if current_cap != cap {
                    debug!("setting power cap to {cap}");
                    device
                        .set_power_management_limit(cap)
                        .context("Could not set power cap")?;
                }
            } else {
                let current_cap = device.power_management_limit();
                let default_cap = device.power_management_limit_default();

                if let (Ok(current_cap), Ok(default_cap)) = (current_cap, default_cap)
                    && current_cap != default_cap
                {
                    debug!("resetting power cap to {default_cap}");
                    device
                        .set_power_management_limit(default_cap)
                        .context("Could not reset power cap")?;
                }
            }

            self.reset_clocks()?;

            let clocks = &config.clocks_configuration;

            match (clocks.min_core_clock, clocks.max_core_clock) {
                (Some(min), Some(max)) => {
                    debug!("applying GPU locked clocks: {min}..{max}");
                    device
                        .set_gpu_locked_clocks(GpuLockedClocksSetting::Numeric {
                            min_clock_mhz: min as u32,
                            max_clock_mhz: max as u32,
                        })
                        .context("Could not apply GPU locked clocks")?;
                    self.last_applied_gpu_locked_clocks
                        .replace(Some((min as u32, max as u32)));
                }
                (None, None) => (),
                _ => bail!("Min and max GPU clock must be set together"),
            }

            match (clocks.min_memory_clock, clocks.max_memory_clock) {
                (Some(min), Some(max)) => {
                    debug!("applying VRAM locked clocks: {min}..{max}");
                    device
                        .set_mem_locked_clocks(min as u32, max as u32)
                        .context("Could not apply VRAM locked clocks")?;
                    self.last_applied_vram_locked_clocks
                        .replace(Some((min as u32, max as u32)));
                }
                (None, None) => (),
                _ => bail!("Min and max VRAM clock must be set together"),
            }

            for (pstate, offset) in &clocks.gpu_clock_offsets {
                let pstate = PerformanceState::try_from(*pstate)
                    .map_err(|_| anyhow!("Invalid pstate '{pstate}'"))?;
                debug!("applying offset {offset} for GPU pstate {pstate:?}");
                device
                    .set_clock_offset(Clock::Graphics, pstate, *offset)
                    .with_context(|| {
                        format!("Could not set clock offset {offset} for GPU pstate {pstate:?}")
                    })?;

                self.last_applied_offsets
                    .borrow_mut()
                    .entry(Clock::Graphics)
                    .or_default()
                    .insert(pstate, *offset);
            }

            for (pstate, offset) in &clocks.mem_clock_offsets {
                let pstate = PerformanceState::try_from(*pstate)
                    .map_err(|_| anyhow!("Invalid pstate '{pstate}'"))?;
                debug!("applying offset {offset} for VRAM pstate {pstate:?}");
                device
                    .set_clock_offset(Clock::Memory, pstate, *offset)
                    .with_context(|| {
                        format!("Could not set clock offset {offset} for VRAM pstate {pstate:?}")
                    })?;

                self.last_applied_offsets
                    .borrow_mut()
                    .entry(Clock::Memory)
                    .or_default()
                    .insert(pstate, *offset);
            }

            if !clocks.gpu_vf_curve.is_empty() {
                self.apply_vf_curve(&clocks.gpu_vf_curve)
                    .context("Could not apply VF curve")?;
            }

            if config.fan_control_enabled {
                let settings = config
                    .fan_control_settings
                    .as_ref()
                    .context("Fan control enabled with no settings")?;
                match settings.mode {
                    FanControlMode::Static => {
                        self.stop_fan_control()
                            .await
                            .context("Could not reset fan control")?;

                        #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
                        let speed = (settings.static_speed * 100.0) as u32;

                        let fan_count = device.num_fans().context("Could not get fan count")?;
                        for fan in 0..fan_count {
                            device
                                .set_fan_speed(fan, speed)
                                .context("Could not reset fan speed to default")?;
                        }
                    }

                    FanControlMode::Curve => {
                        let (min_speed, max_speed) = device
                            .min_max_fan_speed()
                            .context("Could not get fan speed range")?;

                        for point in settings.curve.0.values() {
                            #[allow(clippy::cast_possible_truncation)]
                            if !(min_speed..=max_speed).contains(&((*point * 100.0) as u32)) {
                                bail!(
                                    "Fan speed {}% outside of the allowed range {min_speed}% to {max_speed}%",
                                    point * 100.0
                                );
                            }
                        }

                        self.start_curve_fan_control_task(settings.curve.clone(), settings.clone())
                            .await?;
                    }
                }
            } else {
                self.stop_fan_control()
                    .await
                    .context("Could not reset fan control")?;
            }

            if let Some(target_temp_info) = self.get_target_temp()
                && let Some((min, max)) = target_temp_info.allowed_range
            {
                if let Some(target_temp) = config.nvidia_thermal_options.target_temperature {
                    let target_temp = target_temp.clamp(min, max);

                    if target_temp_info.current != target_temp {
                        debug!("setting target temperature to {target_temp}");
                        if let Err(err) = device.set_temperature_threshold(
                            TemperatureThreshold::AcousticCurr,
                            target_temp.cast_signed(),
                        ) {
                            warn!("Could not set target temperature: {err:#}");
                        }
                    }
                } else if let Err(err) = self.reset_target_temp() {
                    warn!("could not reset target temperature: {err:#}");
                }
            }

            Ok(())
        })
    }

    fn reset_clocks(&self) -> anyhow::Result<()> {
        let mut device = self.device();

        if let Ok(supported_pstates) = device.supported_performance_states() {
            for pstate in supported_pstates {
                for clock_type in [Clock::Graphics, Clock::Memory] {
                    if let Ok(current_offset) = device.clock_offset(clock_type, pstate)
                        && (current_offset.clock_offset_mhz != 0
                            || self
                                .last_applied_offsets
                                .borrow()
                                .get(&clock_type)
                                .and_then(|applied_offsets| applied_offsets.get(&pstate))
                                .is_some_and(|offset| *offset != 0))
                    {
                        debug!("resetting clock offset for {clock_type:?} pstate {pstate:?}");
                        if let Err(err) = device.set_clock_offset(clock_type, pstate, 0) {
                            warn!("could not reset {clock_type:?} pstate {pstate:?}: {err:#}");
                        }
                    }

                    if let Some(applied_offsets) =
                        self.last_applied_offsets.borrow_mut().get_mut(&clock_type)
                    {
                        applied_offsets.remove(&pstate);
                    }
                }
            }
        }

        if self.last_applied_gpu_locked_clocks.borrow().is_some() {
            device
                .reset_gpu_locked_clocks()
                .context("Could not reset locked GPU clocks")?;
            self.last_applied_gpu_locked_clocks.take();
        }

        if self.last_applied_vram_locked_clocks.borrow().is_some() {
            device
                .reset_mem_locked_clocks()
                .context("Could not reset locked GPU clocks")?;
            self.last_applied_vram_locked_clocks.take();
        }

        if self.vf_curve_written.get() {
            self.reset_vf_curve().context("Could not reset VF curve")?;
        }

        Ok(())
    }

    fn cleanup(&self) -> LocalBoxFuture<'_, ()> {
        async {
            if let Some((fan_notify, fan_handle)) = self.fan_control_handle.take() {
                debug!("sending stop notification to old fan control task");
                fan_notify.notify_one();
                fan_handle.await.unwrap();
                debug!("finished controller cleanup");
            }
        }
        .boxed_local()
    }

    fn process_list(&self) -> anyhow::Result<ProcessList> {
        fn map_process(
            process: &nvml_wrapper::struct_wrappers::device::ProcessInfo,
            process_type: ProcessType,
        ) -> ProcessInfo {
            #[allow(clippy::cast_possible_wrap)]
            let (name, args) = resolve_process_name((process.pid as i32).into())
                .unwrap_or_else(|_| ("<Unknown>".to_owned(), String::new()));

            ProcessInfo {
                name,
                args,
                memory_used: match process.used_gpu_memory {
                    UsedGpuMemory::Used(size) => size,
                    UsedGpuMemory::Unavailable => 0,
                },
                types: vec![process_type],
                util: SUPPORTED_UTIL_TYPES.iter().map(|util| (*util, 0)).collect(),
            }
        }

        let device = self.device();

        let mut processes = BTreeMap::new();

        for process in device
            .running_graphics_processes()
            .context("Could not get graphics processes")?
        {
            processes.insert(process.pid, map_process(&process, ProcessType::Graphics));
        }

        for process in device
            .running_compute_processes()
            .context("Could not get compute processes")?
        {
            match processes.entry(process.pid) {
                Entry::Vacant(entry) => {
                    entry.insert(map_process(&process, ProcessType::Compute));
                }
                Entry::Occupied(mut entry) => {
                    entry.get_mut().types.push(ProcessType::Compute);
                }
            }
        }

        match device.process_utilization_stats(self.last_util_timestamp.get()) {
            Ok(stats) => {
                if let Some(stat) = stats.first() {
                    self.last_util_timestamp.set(Some(stat.timestamp));
                }

                for stat in stats {
                    if let Some(info) = processes.get_mut(&stat.pid) {
                        info.util
                            .insert(ProcessUtilizationType::Graphics, stat.sm_util);
                        info.util
                            .insert(ProcessUtilizationType::Memory, stat.mem_util);
                        info.util
                            .insert(ProcessUtilizationType::Encode, stat.enc_util);
                        info.util
                            .insert(ProcessUtilizationType::Decode, stat.dec_util);
                    }
                }
            }
            Err(NvmlError::NotFound) => (),
            Err(err) => {
                error!("could not get process util stats: {err}");
            }
        }
        Ok(ProcessList {
            processes,
            supported_util_types: SUPPORTED_UTIL_TYPES.iter().copied().collect(),
        })
    }

    #[cfg(feature = "display-info")]
    fn populate_displays_info(&self, info: &mut lact_schema::DisplaysInfo) -> anyhow::Result<()> {
        use lact_schema::DisplayConnector;
        use std::fs;
        use std::os::fd::AsRawFd as _;

        if let Some(handle) = &self.driver_handle {
            let drm_path = self.common.get_drm_render()?;
            let drm_file = fs::OpenOptions::new()
                .read(true)
                .write(true)
                .open(drm_path)
                .context("Could not open DRM file")?;

            for (key, display_info) in &mut info.displays {
                match driver::connector_id_to_display_id(
                    display_info.connector_id,
                    drm_file.as_raw_fd(),
                ) {
                    Ok(display_id) => {
                        if let DisplayConnector::DisplayPort {
                            lanes, bandwidth, ..
                        } = &mut display_info.connector_type
                        {
                            match handle.get_dp_link_config(display_id) {
                                Ok(params) => {
                                    *lanes = Some(params.laneCount.try_into()?);
                                    *bandwidth = Some(if params.linkBW != 0 {
                                        crate::server::display::dp1_rate_to_bandwidth(params.linkBW)
                                    } else {
                                        crate::server::display::dp2_rate_to_bandwidth(
                                            params.dp2LinkBW,
                                        )
                                    });
                                }
                                Err(err) => {
                                    warn!("could not fetch DP info for display {key}: {err:#}");
                                }
                            }
                        }
                    }
                    Err(err) => {
                        warn!(
                            "could not resolve display '{key}' into the driver display id: {err:#}"
                        );
                    }
                }
            }
        }

        Ok(())
    }
}
