"""psutil-based health metrics — cross-platform."""
from __future__ import annotations

import logging
import time
from datetime import datetime

import psutil

from mesh.models import HealthSnapshot

logger = logging.getLogger(__name__)

_boot_time = psutil.boot_time()
_prev_net_bytes_sent: int | None = None
_prev_net_bytes_recv: int | None = None
_prev_net_timestamp: float | None = None
_SKIP_FSTYPES = {
    "", "autofs", "binfmt_misc", "bpf", "cgroup", "cgroup2", "configfs", "debugfs",
    "devpts", "devtmpfs", "efivarfs", "fusectl", "hugetlbfs", "mqueue", "overlay",
    "proc", "pstore", "rpc_pipefs", "securityfs", "selinuxfs", "squashfs", "sysfs",
    "tmpfs", "tracefs",
}
_SKIP_MOUNT_PREFIXES = (
    "/dev",
    "/proc",
    "/run",
    "/sys",
)


def collect_health() -> HealthSnapshot:
    """Collect a health snapshot of the current machine."""
    global _prev_net_bytes_sent, _prev_net_bytes_recv, _prev_net_timestamp

    vm = psutil.virtual_memory()
    disk = psutil.disk_usage("/")
    net = psutil.net_io_counters()
    load = psutil.getloadavg()
    freq = psutil.cpu_freq()
    now = time.time()

    net_bandwidth_mbps = 0.0
    if (
        _prev_net_bytes_sent is not None
        and _prev_net_bytes_recv is not None
        and _prev_net_timestamp is not None
    ):
        elapsed = now - _prev_net_timestamp
        if elapsed > 0:
            delta_sent = max(0, net.bytes_sent - _prev_net_bytes_sent)
            delta_recv = max(0, net.bytes_recv - _prev_net_bytes_recv)
            net_bandwidth_mbps = ((delta_sent + delta_recv) * 8.0) / (elapsed * 1_000_000.0)

    _prev_net_bytes_sent = net.bytes_sent
    _prev_net_bytes_recv = net.bytes_recv
    _prev_net_timestamp = now

    snap = HealthSnapshot(
        timestamp=datetime.utcnow(),
        cpu_pct=psutil.cpu_percent(interval=0.1),
        cpu_freq_mhz=freq.current if freq else 0.0,
        memory_used_pct=vm.percent,
        memory_available_mb=vm.available // (1024 * 1024),
        disk_used_pct=disk.percent,
        load_1m=load[0],
        load_5m=load[1],
        load_15m=load[2],
        net_bytes_sent=net.bytes_sent,
        net_bytes_recv=net.bytes_recv,
        uptime_s=time.time() - _boot_time,
        mount_free_mb=_collect_mount_free_mb(),
        net_bandwidth_mbps=net_bandwidth_mbps,
    )

    # GPU metrics
    snap.gpu_utilizations, snap.gpu_temperatures, snap.gpu_memory_used_mb = _gpu_metrics()

    return snap


def _collect_mount_free_mb() -> dict[str, int]:
    mount_free_mb: dict[str, int] = {}
    try:
        for partition in psutil.disk_partitions(all=False):
            if _skip_partition(partition.fstype, partition.mountpoint):
                continue
            try:
                usage = psutil.disk_usage(partition.mountpoint)
            except Exception:
                continue
            mount_free_mb[partition.mountpoint] = int(usage.free // (1024 * 1024))
    except Exception as e:
        logger.debug("Mount free space collection failed: %s", e)
    return mount_free_mb


def _skip_partition(fstype: str, mountpoint: str) -> bool:
    if fstype in _SKIP_FSTYPES:
        return True
    return any(mountpoint == prefix or mountpoint.startswith(prefix + "/")
               for prefix in _SKIP_MOUNT_PREFIXES)


def _gpu_metrics() -> tuple[list[float], list[float], list[int]]:
    """Collect GPU utilization, temperature, and memory from nvidia-smi or rocm-smi."""
    utils, temps, mems = [], [], []

    # NVIDIA
    try:
        import subprocess
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=utilization.gpu,temperature.gpu,memory.used",
             "--format=csv,noheader,nounits"],
            text=True, stderr=subprocess.DEVNULL, timeout=5
        )
        for line in out.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 3:
                utils.append(float(parts[0]) if parts[0] != "N/A" else 0.0)
                temps.append(float(parts[1]) if parts[1] != "N/A" else 0.0)
                mems.append(int(float(parts[2])) if parts[2] != "N/A" else 0)
    except Exception:
        pass

    # AMD
    try:
        import json, subprocess
        out = subprocess.check_output(
            ["rocm-smi", "--showuse", "--showtemp", "--showmeminfo", "vram", "--json"],
            text=True, stderr=subprocess.DEVNULL, timeout=5
        )
        data = json.loads(out)
        for key in sorted(data.keys()):
            if not key.startswith("card"):
                continue
            val = data[key]
            use = val.get("GPU use (%)", 0)
            utils.append(float(use) if use else 0.0)
            temp = val.get("Temperature (Sensor edge) (C)", 0)
            temps.append(float(temp) if temp else 0.0)
            vram_used = val.get("VRAM Total Used Memory (B)", 0)
            mems.append(int(vram_used) // (1024 * 1024) if vram_used else 0)
    except Exception:
        pass

    return utils, temps, mems
