import psutil
import subprocess

try:
    import GPUtil as _GPUtil
except ImportError:
    _GPUtil = None

try:
    import pynvml as _nvml
    _nvml.nvmlInit()
    _NVML_OK = True
except Exception:
    _nvml = None
    _NVML_OK = False


# ──────────────────────────────────────────────────────────────
# GPU helpers
# ──────────────────────────────────────────────────────────────
def _gpu_via_nvml() -> dict | None:
    """Read GPU metrics via pynvml (works on Colab GPU runtime)."""
    if not _NVML_OK:
        return None
    try:
        h = _nvml.nvmlDeviceGetHandleByIndex(0)
        util    = _nvml.nvmlDeviceGetUtilizationRates(h)
        temp    = _nvml.nvmlDeviceGetTemperature(h, _nvml.NVML_TEMPERATURE_GPU)
        mem     = _nvml.nvmlDeviceGetMemoryInfo(h)
        name    = _nvml.nvmlDeviceGetName(h)
        if isinstance(name, bytes):
            name = name.decode()
        vram_total = float(mem.total) / 1_048_576
        vram_used  = float(mem.used)  / 1_048_576
        vram_pct   = round(vram_used / vram_total * 100, 1) if vram_total > 0 else 0.0
        try:
            power_mw = _nvml.nvmlDeviceGetPowerUsage(h)
            power_w  = float(power_mw) / 1000.0
        except Exception:
            power_w = 0.0
        return {
            "gpu_util_avg":     float(util.gpu),
            "gpu_temp_c":       float(temp),
            "gpu_power_w":      power_w,
            "gpu_mem_used_mb":  vram_used,
            "gpu_mem_free_mb":  float(mem.free) / 1_048_576,
            "gpu_mem_total_mb": vram_total,
            "gpu_mem_used_pct": vram_pct,
            "hardware_type":    "GPU",
            "hardware_name":    name,
        }
    except Exception:
        return None


def _gpu_via_gputil() -> dict | None:
    """Read GPU metrics via GPUtil."""
    if _GPUtil is None:
        return None
    try:
        gpus = _GPUtil.getGPUs()
        if not gpus:
            return None
        gpu = gpus[0]
        vram_total = float(gpu.memoryTotal) if gpu.memoryTotal else 0.0
        vram_used  = float(gpu.memoryUsed)  if gpu.memoryUsed  else 0.0
        vram_pct   = round(vram_used / vram_total * 100, 1) if vram_total > 0 else 0.0
        return {
            "gpu_util_avg":     float(gpu.load * 100),
            "gpu_temp_c":       float(gpu.temperature),
            "gpu_power_w":      0.0,
            "gpu_mem_used_mb":  vram_used,
            "gpu_mem_free_mb":  float(gpu.memoryFree) if gpu.memoryFree else 0.0,
            "gpu_mem_total_mb": vram_total,
            "gpu_mem_used_pct": vram_pct,
            "hardware_type":    "GPU",
            "hardware_name":    gpu.name,
        }
    except Exception:
        return None


def _gpu_via_nvidiasmi() -> dict | None:
    """
    Last-resort: parse `nvidia-smi` output.
    Works on Colab + any Linux with nvidia-smi in PATH.
    """
    try:
        out = subprocess.check_output(
            ["nvidia-smi",
             "--query-gpu=name,utilization.gpu,temperature.gpu,"
             "memory.used,memory.free,memory.total,power.draw",
             "--format=csv,noheader,nounits"],
            stderr=subprocess.DEVNULL,
            timeout=4,
        ).decode().strip()
        parts = [p.strip() for p in out.split(",")]
        if len(parts) < 6:
            return None
        name      = parts[0]
        util      = float(parts[1])
        temp      = float(parts[2])
        mem_used  = float(parts[3])
        mem_free  = float(parts[4])
        mem_total = float(parts[5])
        power_w   = float(parts[6]) if len(parts) > 6 else 0.0
        vram_pct  = round(mem_used / mem_total * 100, 1) if mem_total > 0 else 0.0
        return {
            "gpu_util_avg":     util,
            "gpu_temp_c":       temp,
            "gpu_power_w":      power_w,
            "gpu_mem_used_mb":  mem_used,
            "gpu_mem_free_mb":  mem_free,
            "gpu_mem_total_mb": mem_total,
            "gpu_mem_used_pct": vram_pct,
            "hardware_type":    "GPU",
            "hardware_name":    name,
        }
    except Exception:
        return None


# ──────────────────────────────────────────────────────────────
# CPU TEMP helpers
# ──────────────────────────────────────────────────────────────
# Sensor names seen across platforms
_CPU_SENSOR_NAMES = [
    "coretemp",   # Intel Linux
    "k10temp",    # AMD Linux
    "acpitz",     # ACPI thermal (common on VMs / Colab)
    "cpu-thermal",# ARM / Raspberry Pi
    "cpu_thermal",
    "cpu0-thermal",
]


def _cpu_temp_psutil() -> float:
    """Read CPU temp via psutil (works on Linux including Colab if exposed)."""
    try:
        temps = psutil.sensors_temperatures()
        if not temps:
            return 0.0
        # try known names first
        for name in _CPU_SENSOR_NAMES:
            if name in temps and temps[name]:
                return float(temps[name][0].current)
        # fallback: first available sensor
        for entries in temps.values():
            if entries:
                t = entries[0].current
                if t and t > 0:
                    return float(t)
    except Exception:
        pass
    return 0.0


def _cpu_temp_thermal_zone() -> float:
    """
    Read /sys/class/thermal/thermal_zone*/temp (Linux / Colab).
    Returns °C or 0 on failure.
    """
    try:
        import glob
        zones = sorted(glob.glob("/sys/class/thermal/thermal_zone*/temp"))
        for zone in zones:
            with open(zone) as f:
                val = int(f.read().strip())
                if val > 0:
                    return val / 1000.0
    except Exception:
        pass
    return 0.0


def _read_gpu_temp_wmi() -> float | None:
    """Windows WMI GPU temp (requires Open Hardware Monitor running)."""
    try:
        import wmi
        w = wmi.WMI(namespace="root\\OpenHardwareMonitor")
        for sensor in w.Sensor():
            if sensor.SensorType == "Temperature" and "gpu" in sensor.Name.lower():
                return float(sensor.Value)
    except Exception:
        pass
    return None


# ──────────────────────────────────────────────────────────────
def get_metrics() -> dict:
    """
    Safe system metric snapshot. Always returns a complete dict.
    Tries multiple backends for GPU and CPU temperature so values
    are populated on Colab, Windows, macOS, and bare Linux.
    """
    m: dict = {}

    # ══════════════════════════════════════════════════════════
    # CPU UTILISATION
    # ══════════════════════════════════════════════════════════
    try:
        m["cpu_util_avg"] = float(psutil.cpu_percent(interval=None))
    except Exception:
        m["cpu_util_avg"] = 0.0

    # ══════════════════════════════════════════════════════════
    # CPU TEMPERATURE  (multi-backend with fallbacks)
    # ══════════════════════════════════════════════════════════
    cpu_temp = _cpu_temp_psutil()
    if cpu_temp == 0.0:
        cpu_temp = _cpu_temp_thermal_zone()
    m["cpu_temp_c"] = cpu_temp

    # ══════════════════════════════════════════════════════════
    # RAM (system memory)
    # ══════════════════════════════════════════════════════════
    try:
        vm = psutil.virtual_memory()
        m["ram_used_pct"]  = float(vm.percent)
        m["ram_used_mb"]   = float(vm.used  / 1_048_576)
        m["ram_total_mb"]  = float(vm.total / 1_048_576)
    except Exception:
        m["ram_used_pct"] = m["ram_used_mb"] = m["ram_total_mb"] = 0.0

    # ══════════════════════════════════════════════════════════
    # GPU  — try three backends in order of reliability
    # ══════════════════════════════════════════════════════════
    gpu_data = _gpu_via_nvml() or _gpu_via_gputil() or _gpu_via_nvidiasmi()
    if gpu_data:
        m.update(gpu_data)
        # WMI GPU temp override (Windows only)
        if m.get("gpu_temp_c", 0) == 0.0:
            wmi_temp = _read_gpu_temp_wmi()
            if wmi_temp:
                m["gpu_temp_c"] = wmi_temp
    else:
        # CPU-only fallback
        m.update({
            "gpu_util_avg":     0.0,
            "gpu_temp_c":       0.0,
            "gpu_power_w":      0.0,
            "gpu_mem_used_mb":  0.0,
            "gpu_mem_free_mb":  0.0,
            "gpu_mem_total_mb": 0.0,
            "gpu_mem_used_pct": 0.0,
            "hardware_type":    "CPU",
            "hardware_name":    "CPU Only",
        })

    # ══════════════════════════════════════════════════════════
    # CPU POWER ESTIMATE  (TDP-scaled heuristic)
    # ══════════════════════════════════════════════════════════
    try:
        m["cpu_power_w"] = max(2.0, m["cpu_util_avg"] * 0.35)
    except Exception:
        m["cpu_power_w"] = 3.0

    return m
