from __future__ import annotations

import shutil
from pathlib import Path

from adam.models import SystemSnapshot

try:
    import psutil
except ImportError:  # pragma: no cover - optional runtime dependency
    psutil = None


class SystemMonitor:
    def __init__(self, root: Path) -> None:
        self.root = root
        self._nvml = None
        self._gpu_handle = None
        try:
            import pynvml

            pynvml.nvmlInit()
            self._nvml = pynvml
            if pynvml.nvmlDeviceGetCount() > 0:
                self._gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
        except Exception:
            self._nvml = None
            self._gpu_handle = None

    def snapshot(self) -> SystemSnapshot:
        snapshot = SystemSnapshot()
        if psutil is not None:
            memory = psutil.virtual_memory()
            disk = psutil.disk_usage(str(self.root.anchor or self.root))
            snapshot.cpu_percent = float(psutil.cpu_percent(interval=None))
            snapshot.memory_percent = float(memory.percent)
            snapshot.memory_used_gb = memory.used / (1024**3)
            snapshot.memory_total_gb = memory.total / (1024**3)
            snapshot.disk_percent = float(disk.percent)
            snapshot.disk_used_gb = disk.used / (1024**3)
            snapshot.disk_total_gb = disk.total / (1024**3)
        else:
            total, used, _ = shutil.disk_usage(self.root)
            snapshot.disk_percent = used / total * 100 if total else 0.0
            snapshot.disk_used_gb = used / (1024**3)
            snapshot.disk_total_gb = total / (1024**3)

        if self._nvml is not None and self._gpu_handle is not None:
            try:
                name = self._nvml.nvmlDeviceGetName(self._gpu_handle)
                if isinstance(name, bytes):
                    name = name.decode("utf-8", errors="replace")
                utilization = self._nvml.nvmlDeviceGetUtilizationRates(self._gpu_handle)
                memory = self._nvml.nvmlDeviceGetMemoryInfo(self._gpu_handle)
                snapshot.gpu_name = str(name)
                snapshot.gpu_percent = float(utilization.gpu)
                snapshot.vram_used_gb = memory.used / (1024**3)
                snapshot.vram_total_gb = memory.total / (1024**3)
                snapshot.vram_percent = (
                    memory.used / memory.total * 100 if memory.total else 0.0
                )
                snapshot.gpu_temperature = float(
                    self._nvml.nvmlDeviceGetTemperature(
                        self._gpu_handle,
                        self._nvml.NVML_TEMPERATURE_GPU,
                    )
                )
            except Exception:
                pass
        return snapshot

    def close(self) -> None:
        if self._nvml is not None:
            try:
                self._nvml.nvmlShutdown()
            except Exception:
                pass
            self._nvml = None