# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Sequence
from pprint import pformat
from typing import NamedTuple, Optional

from oumi.utils.logging import logger

try:
    # The library is only useful for NVIDIA GPUs, and
    # may not be installed for other vendors e.g., AMD
    import pynvml  # pyright: ignore[reportMissingImports]
except ModuleNotFoundError:
    pynvml = None

# TODO: OPE-562 - Add support for `amdsmi.amdsmi_init()`` for AMD GPUs


def _initialize_pynvml() -> bool:
    """Attempts to initialize pynvml library. Returns True on success."""
    global pynvml
    if pynvml is None:
        return False

    try:
        pynvml.nvmlInit()
    except Exception:
        logger.error(
            "Failed to initialize pynvml library. All pynvml calls will be disabled."
        )
        pynvml = None

    return pynvml is not None


def _initialize_pynvml_and_get_pynvml_device_count() -> Optional[int]:
    """Attempts to initialize pynvml library.

    Returns device count on success, or None otherwise.
    """
    global pynvml
    # The call to `pynvml is None` is technically redundant but exists here
    # to make pyright happy.
    if pynvml is None or not _initialize_pynvml():
        return None
    return int(pynvml.nvmlDeviceGetCount())



[docs]
class NVidiaGpuRuntimeInfo(NamedTuple):
    """Contains misc NVIDIA GPU measurements and stats retrieved by `pynvml`.

    The majority of fields are optional. You can control whether they are
    populated by setting boolean query parameters of
    `_get_nvidia_gpu_runtime_info_impl(, ...)` such as `memory`, `temperature`,
    `fan_speed`, etc.
    """

    device_index: int
    """Zero-based device index."""

    device_count: int
    """Total number of GPU devices on this node."""

    used_memory_mb: Optional[float] = None
    """Used GPU memory in MB."""

    temperature: Optional[int] = None
    """GPU temperature in Celcius."""

    fan_speed: Optional[int] = None
    """GPU fan speed in [0,100] range."""

    fan_speeds: Optional[Sequence[int]] = None
    """An array of GPU fan speeds.

    The array's length is equal to the number of fans per GPU (can be multiple).
    Speed values are in [0, 100] range.
    """

    power_usage_watts: Optional[float] = None
    """GPU power usage in Watts."""

    power_limit_watts: Optional[float] = None
    """GPU power limit in Watts."""

    gpu_utilization: Optional[int] = None
    """GPU compute utilization. Range: [0,100]."""

    memory_utilization: Optional[int] = None
    """GPU memory utilization. Range: [0,100]."""

    performance_state: Optional[int] = None
    """See `nvmlPstates_t`. Valid values are in [0,15] range, or 32 if unknown.

    0 corresponds to Maximum Performance.
    15 corresponds to Minimum Performance.
    """

    clock_speed_graphics: Optional[int] = None
    """Graphics clock speed (`NVML_CLOCK_GRAPHICS`) in MHz."""

    clock_speed_sm: Optional[int] = None
    """SM clock speed (`NVML_CLOCK_SM`) in MHz."""

    clock_speed_memory: Optional[int] = None
    """Memory clock speed (`NVML_CLOCK_MEM`) in MHz."""



def _get_nvidia_gpu_runtime_info_impl(
    device_index: int = 0,
    *,
    memory: bool = False,
    temperature: bool = False,
    fan_speed: bool = False,
    power_usage: bool = False,
    utilization: bool = False,
    performance_state: bool = False,
    clock_speed: bool = False,
) -> Optional[NVidiaGpuRuntimeInfo]:
    global pynvml
    if pynvml is None:
        return None

    device_count = _initialize_pynvml_and_get_pynvml_device_count()
    if device_count is None or device_count <= 0:
        return None
    elif device_index < 0 or device_index >= device_count:
        raise ValueError(
            f"Device index ({device_index}) must be "
            f"within the [0, {device_count}) range."
        )

    try:
        gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
    except Exception:
        logger.exception(f"Failed to get GPU handle for device: {device_index}")
        return None

    used_memory_mb_value: Optional[float] = None
    if memory:
        try:
            info = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle)
            used_memory_mb_value = float(info.used) // 1024**2
        except Exception:
            logger.exception(
                f"Failed to get GPU memory info for device: {device_index}"
            )
            return None

    temperature_value: Optional[int] = None
    if temperature:
        try:
            temperature_value = pynvml.nvmlDeviceGetTemperature(
                gpu_handle, pynvml.NVML_TEMPERATURE_GPU
            )
        except Exception:
            logger.exception(
                f"Failed to get GPU temperature for device: {device_index}"
            )
            return None

    fan_speed_value: Optional[int] = None
    fan_speeds_value: Optional[Sequence[int]] = None
    if fan_speed:
        try:
            fan_speed_value = pynvml.nvmlDeviceGetFanSpeed(gpu_handle)
        except Exception:
            # The `GetFanSpeed` function fails on many systems
            # Only do DEBUG-level logging to reduce noise.
            logger.debug(
                f"Failed to get GPU fan speed for device: {device_index}", exc_info=True
            )

        if fan_speed_value is not None:
            fan_speeds_value = tuple([fan_speed_value])
            if hasattr(pynvml, "nvmlDeviceGetNumFans"):
                try:
                    fan_count = pynvml.nvmlDeviceGetNumFans(gpu_handle)
                    value = [0] * fan_count
                    for i in range(fan_count):
                        speed = pynvml.nvmlDeviceGetFanSpeed_v2(gpu_handle, i)
                        value[i] = speed
                    # Make it immutable.
                    fan_speeds_value = tuple(value)
                except Exception:
                    fan_speeds_value = tuple([fan_speed_value])

    power_usage_watts_value: Optional[float] = None
    power_limit_watts_value: Optional[float] = None
    if power_usage:
        try:
            milliwatts = pynvml.nvmlDeviceGetPowerUsage(gpu_handle)
            power_usage_watts_value = float(milliwatts) * 1e-3

            milliwatts = pynvml.nvmlDeviceGetPowerManagementLimit(gpu_handle)
            power_limit_watts_value = float(milliwatts) * 1e-3
        except Exception:
            logger.exception(
                f"Failed to get GPU power usage for device: {device_index}"
            )
            return None

    gpu_utilization_value: Optional[float] = None
    memory_utilization_value: Optional[float] = None
    if utilization:
        try:
            result = pynvml.nvmlDeviceGetUtilizationRates(gpu_handle)
            gpu_utilization_value = int(result.gpu)
            memory_utilization_value = int(result.memory)
        except Exception:
            logger.exception(
                f"Failed to get GPU utilization for device: {device_index}"
            )
            return None

    performance_state_value: Optional[int] = None
    if performance_state:
        try:
            performance_state_value = int(
                pynvml.nvmlDeviceGetPerformanceState(gpu_handle)
            )
        except Exception:
            logger.exception(
                f"Failed to get GPU performance state for device: {device_index}"
            )
            return None

    clock_speed_graphics_value: Optional[int] = None
    clock_speed_sm_value: Optional[int] = None
    clock_speed_memory_value: Optional[int] = None
    if clock_speed:
        try:
            clock_speed_graphics_value = int(
                pynvml.nvmlDeviceGetClockInfo(gpu_handle, pynvml.NVML_CLOCK_GRAPHICS)
            )
            clock_speed_sm_value = int(
                pynvml.nvmlDeviceGetClockInfo(gpu_handle, pynvml.NVML_CLOCK_SM)
            )
            clock_speed_memory_value = int(
                pynvml.nvmlDeviceGetClockInfo(gpu_handle, pynvml.NVML_CLOCK_MEM)
            )
        except Exception:
            logger.exception(
                f"Failed to get GPU clock speed for device: {device_index}"
            )
            return None

    return NVidiaGpuRuntimeInfo(
        device_index=device_index,
        device_count=device_count,
        used_memory_mb=used_memory_mb_value,
        temperature=temperature_value,
        fan_speed=fan_speed_value,
        fan_speeds=fan_speeds_value,
        power_usage_watts=power_usage_watts_value,
        power_limit_watts=power_limit_watts_value,
        gpu_utilization=gpu_utilization_value,
        memory_utilization=memory_utilization_value,
        performance_state=performance_state_value,
        clock_speed_graphics=clock_speed_graphics_value,
        clock_speed_sm=clock_speed_sm_value,
        clock_speed_memory=clock_speed_memory_value,
    )



[docs]
def get_nvidia_gpu_runtime_info(
    device_index: int = 0,
) -> Optional[NVidiaGpuRuntimeInfo]:
    """Returns runtime stats for Nvidia GPU."""
    return _get_nvidia_gpu_runtime_info_impl(
        device_index=device_index,
        memory=True,
        temperature=True,
        fan_speed=True,
        power_usage=True,
        utilization=True,
        performance_state=True,
        clock_speed=True,
    )




[docs]
def log_nvidia_gpu_runtime_info(device_index: int = 0, log_prefix: str = "") -> None:
    """Prints the current NVIDIA GPU runtime info."""
    info = get_nvidia_gpu_runtime_info(device_index)
    logger.info(f"{log_prefix.rstrip()} GPU runtime info: {pformat(info)}.")




[docs]
def get_nvidia_gpu_memory_utilization(device_index: int = 0) -> float:
    """Returns amount of memory being used on an Nvidia GPU in MiB."""
    info: Optional[NVidiaGpuRuntimeInfo] = _get_nvidia_gpu_runtime_info_impl(
        device_index=device_index, memory=True
    )
    return (
        info.used_memory_mb
        if (info is not None and info.used_memory_mb is not None)
        else 0.0
    )




[docs]
def log_nvidia_gpu_memory_utilization(
    device_index: int = 0, log_prefix: str = ""
) -> None:
    """Prints amount of memory being used on an Nvidia GPU."""
    memory_mib = get_nvidia_gpu_memory_utilization(device_index)
    logger.info(f"{log_prefix.rstrip()} GPU memory occupied: {memory_mib} MiB.")




[docs]
def get_nvidia_gpu_temperature(device_index: int = 0) -> int:
    """Returns the current temperature readings for the device, in degrees C."""
    info: Optional[NVidiaGpuRuntimeInfo] = _get_nvidia_gpu_runtime_info_impl(
        device_index=device_index,
        temperature=True,
    )
    return (
        info.temperature if (info is not None and info.temperature is not None) else 0
    )




[docs]
def log_nvidia_gpu_temperature(device_index: int = 0, log_prefix: str = "") -> None:
    """Prints the current temperature readings for the device, in degrees C."""
    temperature = get_nvidia_gpu_temperature(device_index)
    logger.info(f"{log_prefix.rstrip()} GPU temperature: {temperature} C.")




[docs]
def get_nvidia_gpu_fan_speeds(device_index: int = 0) -> Sequence[int]:
    """Returns the current fan speeds for NVIDIA GPU device."""
    info: Optional[NVidiaGpuRuntimeInfo] = _get_nvidia_gpu_runtime_info_impl(
        device_index=device_index, fan_speed=True
    )
    return (
        info.fan_speeds
        if (info is not None and info.fan_speeds is not None)
        else tuple()
    )




[docs]
def log_nvidia_gpu_fan_speeds(device_index: int = 0, log_prefix: str = "") -> None:
    """Prints the current NVIDIA GPU fan speeds."""
    fan_speeds = get_nvidia_gpu_fan_speeds(device_index)
    logger.info(f"{log_prefix.rstrip()} GPU fan speeds: {fan_speeds}.")




[docs]
def get_nvidia_gpu_power_usage(device_index: int = 0) -> float:
    """Returns the current power usage for NVIDIA GPU device."""
    info: Optional[NVidiaGpuRuntimeInfo] = _get_nvidia_gpu_runtime_info_impl(
        device_index=device_index, power_usage=True
    )
    return (
        info.power_usage_watts
        if (info is not None and info.power_usage_watts is not None)
        else 0.0
    )




[docs]
def log_nvidia_gpu_power_usage(device_index: int = 0, log_prefix: str = "") -> None:
    """Prints the current NVIDIA GPU power usage."""
    power_usage = get_nvidia_gpu_power_usage(device_index)
    logger.info(f"{log_prefix.rstrip()} GPU power usage: {power_usage:.2f}W.")