use crate::data::{GpuMemory, GpuPowerUsage, GpuUsage};
use nvml_wrapper::enum_wrappers::device::TemperatureSensor;
use nvml_wrapper::{Device, Nvml};
use once_cell::sync::Lazy;
use std::borrow::Cow;

static NVIDIA: Lazy<Option<Nvml>> = Lazy::new(|| Nvml::init().ok());

fn devices() -> Option<impl Iterator<Item = Device<'static>>> {
    let count = NVIDIA.as_ref()?.device_count().unwrap_or_default();
    Some((0..count).flat_map(device))
}

fn device(index: u32) -> Option<Device<'static>> {
    NVIDIA.as_ref()?.device_by_index(index).ok()
}

pub fn temperature() -> Option<impl Iterator<Item = f32>> {
    Some(devices()?.flat_map(|device| {
        device
            .temperature(TemperatureSensor::Gpu)
            .ok()
            .map(|t| t as f32)
    }))
}

pub fn power() -> Option<impl Iterator<Item = GpuPowerUsage>> {
    Some(devices()?.flat_map(|device| {
        let power = device
            .total_energy_consumption()
            .ok()
            .map(|mj| mj * 1_000)?;
        Some(GpuPowerUsage {
            card: device.index().unwrap_or_default(),
            gpu_uj: power,
        })
    }))
}

pub fn memory() -> Option<impl Iterator<Item = GpuMemory>> {
    Some(devices()?.flat_map(|device| {
        let mem = device.memory_info().ok()?;
        Some(GpuMemory {
            card: device.index().unwrap_or_default(),
            total: mem.total,
            free: mem.free,
        })
    }))
}

pub fn utilization() -> Option<impl Iterator<Item = GpuUsage>> {
    let sources = devices()?.flat_map(|device| {
        let utilization = device.utilization_rates().ok();
        [
            ("compute", utilization.as_ref().map(|u| u.gpu)),
            ("memory", utilization.as_ref().map(|u| u.gpu)),
            (
                "encode",
                device.encoder_utilization().ok().map(|u| u.utilization),
            ),
            (
                "decode",
                device.decoder_utilization().ok().map(|u| u.utilization),
            ),
        ]
    });
    Some(sources.into_iter().flat_map(|(system, usage)| {
        Some(GpuUsage {
            card: 0,
            system: Cow::Borrowed(system),
            usage: usage?,
        })
    }))
}