import os
import platform
import warnings

import psutil

try:
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", FutureWarning)
        import pynvml
    NVML_AVAILABLE = True
except ImportError:
    NVML_AVAILABLE = False


class HardwareProbe:
    def __init__(self):
        self.gpu_enabled = False
        self.gpu_name = "N/A"
        self.vram_total = "N/A"

        if not NVML_AVAILABLE:
            return

        # On Windows, check for the NVML DLL before attempting init
        if platform.system() == "Windows":
            nvml_dll = os.path.join(
                os.environ.get("WINDIR", r"C:\Windows"),
                "System32", "nvml.dll",
            )
            if not os.path.exists(nvml_dll):
                return

        try:
            pynvml.nvmlInit()
            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
            self.gpu_name = pynvml.nvmlDeviceGetName(handle)
            if isinstance(self.gpu_name, bytes):
                self.gpu_name = self.gpu_name.decode("utf-8")
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            self.vram_total = f"{round(mem_info.total / (1024**3), 1)}GB"
            self.gpu_enabled = True
        except pynvml.NVMLError:
            self.gpu_enabled = False
        except Exception:
            self.gpu_enabled = False

    def get_stats(self) -> dict:
        """Return hardware telemetry with canonical keys.

        Always returns all keys — GPU fields fall back to "N/A" when
        no NVIDIA GPU is available.
        """
        stats = {
            "ram": f"{psutil.virtual_memory().percent}%",
            "cpu": f"{psutil.cpu_percent(interval=0.1)}%",
            "vram": "N/A",
            "gpu_temp": "N/A",
        }

        if not self.gpu_enabled:
            return stats

        # Each sensor read is independent — one failure shouldn't kill the rest
        try:
            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
        except (pynvml.NVMLError, Exception):
            return stats

        try:
            info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            stats["vram"] = f"{round((info.used / info.total) * 100, 1)}%"
        except (pynvml.NVMLError, Exception):
            pass

        try:
            temp = pynvml.nvmlDeviceGetTemperature(
                handle, pynvml.NVML_TEMPERATURE_GPU
            )
            stats["gpu_temp"] = f"{temp}°C"
        except (pynvml.NVMLError, Exception):
            pass

        return stats

    def get_system_summary(self) -> dict:
        """One-time summary for the hardware banner."""
        ram_total = round(psutil.virtual_memory().total / (1024**3), 1)
        return {
            "gpu_name": self.gpu_name,
            "vram_total": self.vram_total,
            "ram_total": f"{ram_total}GB",
            "gpu_enabled": self.gpu_enabled,
        }

    def shutdown(self):
        if self.gpu_enabled:
            try:
                pynvml.nvmlShutdown()
            except (pynvml.NVMLError, Exception):
                pass

    def __del__(self):
        self.shutdown()
