# coding: utf-8
"""GPU 信息采集模块，基于 pynvml 获取 NVIDIA GPU 占用率等"""

from typing import List

from .schema import GPUInfo

_pynvml_available = False
try:
    import pynvml
    _pynvml_available = True
except ImportError:
    pass


def get_gpu_info() -> List[GPUInfo]:
    """
    获取所有 NVIDIA GPU 的占用率等信息。

    Returns:
        GPUInfo 列表，若无 GPU 或 pynvml 不可用则返回空列表。
    """
    if not _pynvml_available:
        return []

    result: List[GPUInfo] = []
    try:
        pynvml.nvmlInit()
        gpu_count = pynvml.nvmlDeviceGetCount()

        for i in range(gpu_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)

            gpu_name = pynvml.nvmlDeviceGetName(handle)
            if isinstance(gpu_name, bytes):
                gpu_name = gpu_name.decode("utf-8")

            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            mem_used_mb = mem_info.used / 1024**2
            mem_total_mb = mem_info.total / 1024**2
            mem_usage = (mem_info.used / mem_info.total) * 100 if mem_info.total > 0 else 0

            util_rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
            gpu_util = util_rates.gpu
            mem_util = util_rates.memory

            result.append(GPUInfo(
                index=i,
                name=gpu_name,
                gpu_util=gpu_util,
                mem_usage=round(mem_usage, 1),
                mem_util=mem_util,
                mem_used_mb=round(mem_used_mb, 1),
                mem_total_mb=round(mem_total_mb, 1),
            ))
    except Exception:
        result = []
    finally:
        try:
            pynvml.nvmlShutdown()
        except Exception:
            pass

    return result
