#!/usr/bin/env python3
"""
Cross-platform host agent: push CPU / GPU / RAM / disk to CYD over USB serial.

Works on Windows 10/11, Linux, and macOS (CPU/RAM/disk). GPU utilization is
best-effort via optional vendor backends (NVIDIA NVML first, then CLI fallbacks).

Protocol (one line, ASCII-safe):
  SI|host=...|cpu=...|cpu%=..|gpu=...|gpu%=..|ram=used/total|disk=used/total

Deps (required):  pip install -r requirements.txt
  psutil, pyserial
Optional GPU:     pip install nvidia-ml-py
"""

from __future__ import annotations

import argparse
import glob
import os
import platform
import re
import socket
import subprocess
import sys
import time
from typing import List, Optional, Tuple

try:
    import psutil
except ImportError:
    print("Missing psutil. Run: pip install -r requirements.txt", file=sys.stderr)
    raise SystemExit(1)

try:
    import serial
    from serial.tools import list_ports
except ImportError:
    print("Missing pyserial. Run: pip install -r requirements.txt", file=sys.stderr)
    raise SystemExit(1)


def sanitize(s: str, max_len: int = 40) -> str:
    if not s:
        return ""
    t = re.sub(r"[|=\r\n]+", " ", str(s)).strip()
    t = re.sub(r"\s+", " ", t)
    return t[:max_len].strip()


def cpu_name() -> str:
    # Linux
    if sys.platform.startswith("linux"):
        try:
            with open("/proc/cpuinfo", "r", encoding="utf-8", errors="ignore") as f:
                for line in f:
                    if line.lower().startswith("model name") or line.lower().startswith(
                        "hardware"
                    ):
                        name = line.split(":", 1)[1].strip()
                        name = re.sub(r"\(R\)|\(TM\)|CPU @.*|with Radeon.*", "", name)
                        return sanitize(name, 44) or "CPU"
        except OSError:
            pass

    # macOS
    if sys.platform == "darwin":
        try:
            out = subprocess.check_output(
                ["sysctl", "-n", "machdep.cpu.brand_string"],
                stderr=subprocess.DEVNULL,
                text=True,
                timeout=2,
            ).strip()
            if out:
                return sanitize(out, 44)
        except (OSError, subprocess.SubprocessError):
            pass

    # Windows — registry (no admin)
    if sys.platform == "win32":
        try:
            import winreg

            key = winreg.OpenKey(
                winreg.HKEY_LOCAL_MACHINE,
                r"HARDWARE\DESCRIPTION\System\CentralProcessor\0",
            )
            name, _ = winreg.QueryValueEx(key, "ProcessorNameString")
            winreg.CloseKey(key)
            name = re.sub(r"\(R\)|\(TM\)|CPU @.*|with Radeon.*", "", str(name))
            return sanitize(name, 44) or "CPU"
        except OSError:
            pass

    # Fallback
    return sanitize(platform.processor() or "CPU", 44)


def cpu_pct() -> float:
    # Non-blocking after first call; warm up once at start.
    return round(float(psutil.cpu_percent(interval=None)), 1)


def ram_gb() -> Tuple[float, float]:
    vm = psutil.virtual_memory()
    total = round(vm.total / (1024**3), 1)
    used = round(vm.used / (1024**3), 1)
    return used, total


def disk_root() -> str:
    if sys.platform == "win32":
        return os.environ.get("SystemDrive", "C:") + "\\"
    return "/"


def disk_gb() -> Tuple[float, float]:
    root = disk_root()
    try:
        u = psutil.disk_usage(root)
    except OSError:
        # Last resort: first mount with size
        for p in psutil.disk_partitions(all=False):
            try:
                u = psutil.disk_usage(p.mountpoint)
                if u.total > 0:
                    break
            except OSError:
                continue
        else:
            return 0.0, 0.0
    total = round(u.total / (1024**3), 1)
    used = round(u.used / (1024**3), 1)
    return used, total


_VIRTUAL_GPU_RE = re.compile(
    r"Microsoft Basic|Remote Desktop|Oray|IddDriver|Virtual|Parsec|"
    r"TeamViewer|VNC|Indirect|QXL|VMware|VirtualBox|Hyper-V|llvmpipe|SVGA",
    re.I,
)


def _is_real_gpu_name(name: str) -> bool:
    return bool(name) and not _VIRTUAL_GPU_RE.search(name)


# --- GPU: optional layered backends ---

def _gpu_nvidia_nvml() -> Optional[Tuple[str, float]]:
    try:
        import pynvml  # provided by nvidia-ml-py
    except ImportError:
        return None
    try:
        pynvml.nvmlInit()
        try:
            count = pynvml.nvmlDeviceGetCount()
            if count < 1:
                return None
            # Prefer highest utilization discrete card
            best = None
            for i in range(count):
                h = pynvml.nvmlDeviceGetHandleByIndex(i)
                raw = pynvml.nvmlDeviceGetName(h)
                name = raw.decode("utf-8", "ignore") if isinstance(raw, bytes) else str(raw)
                name = name.replace("NVIDIA GeForce ", "").replace("NVIDIA ", "")
                util = pynvml.nvmlDeviceGetUtilizationRates(h).gpu
                cand = (sanitize(name, 44), float(util))
                if best is None or cand[1] > best[1]:
                    best = cand
            return best
        finally:
            try:
                pynvml.nvmlShutdown()
            except Exception:
                pass
    except Exception:
        return None


def _gpu_nvidia_smi() -> Optional[Tuple[str, float]]:
    try:
        out = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=name,utilization.gpu",
                "--format=csv,noheader,nounits",
            ],
            stderr=subprocess.DEVNULL,
            text=True,
            timeout=3,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    best = None
    for line in out.splitlines():
        parts = [p.strip() for p in line.split(",")]
        if len(parts) < 2:
            continue
        name = parts[0].replace("NVIDIA GeForce ", "").replace("NVIDIA ", "")
        try:
            util = float(parts[1])
        except ValueError:
            continue
        cand = (sanitize(name, 44), util)
        if best is None or cand[1] > best[1]:
            best = cand
    return best


def _gpu_amd_cli() -> Optional[Tuple[str, float]]:
    # Linux ROCm / Windows amd-smi when present
    for cmd in (
        ["amd-smi", "metric", "--json"],
        ["rocm-smi", "--showuse", "--showproductname", "--csv"],
        ["rocm-smi", "-u", "--showproductname"],
    ):
        try:
            out = subprocess.check_output(
                cmd, stderr=subprocess.DEVNULL, text=True, timeout=4
            )
        except (OSError, subprocess.SubprocessError):
            continue
        # Heuristic parse — keep resilient across CLI churn
        name = "AMD GPU"
        util = -1.0
        m = re.search(r"Card series:\s*(.+)", out, re.I)
        if m:
            name = m.group(1).strip()
        m = re.search(r"GPU\[0\].*?=\s*(\d+)", out)
        if not m:
            m = re.search(r"GPU use \(%\):\s*(\d+)", out, re.I)
        if not m:
            m = re.search(r"\"gpu_activity\"\s*:\s*\{\s*\"value\"\s*:\s*(\d+)", out)
        if m:
            util = float(m.group(1))
        if util >= 0:
            return sanitize(name, 44), util
    return None


def _gpu_name_windows() -> str:
    try:
        import winreg

        base = r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}"
        root = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, base)
        best = ""
        i = 0
        while True:
            try:
                sub = winreg.EnumKey(root, i)
            except OSError:
                break
            i += 1
            if not re.fullmatch(r"\d{4}", sub):
                continue
            try:
                k = winreg.OpenKey(root, sub)
                name, _ = winreg.QueryValueEx(k, "DriverDesc")
                winreg.CloseKey(k)
            except OSError:
                continue
            name = str(name)
            if not _is_real_gpu_name(name):
                continue
            if re.search(r"NVIDIA|GeForce|RTX|GTX|Radeon|AMD|Intel Arc", name, re.I):
                best = name
                break
            if not best:
                best = name
        winreg.CloseKey(root)
        best = best.replace("NVIDIA GeForce ", "").replace("NVIDIA ", "").replace("AMD ", "")
        return sanitize(best, 44)
    except OSError:
        return ""


def _gpu_name_linux() -> str:
    # lspci is widely available
    try:
        out = subprocess.check_output(
            ["lspci"], stderr=subprocess.DEVNULL, text=True, timeout=3
        )
    except (OSError, subprocess.SubprocessError):
        # DRM card names
        for path in sorted(glob.glob("/sys/class/drm/card*/device/vendor")):
            try:
                card = path.rsplit("/device/", 1)[0]
                # uevent MODALIAS / device product via lspci preferred
                _ = card
            except OSError:
                pass
        return ""
    for line in out.splitlines():
        if "VGA" in line or "3D" in line or "Display" in line:
            name = line.split(":", 2)[-1].strip() if ":" in line else line
            name = re.sub(r"^.*?(VGA|3D|Display) compatible controller:\s*", "", line)
            if _is_real_gpu_name(name):
                name = name.replace("NVIDIA Corporation ", "").replace(
                    "Advanced Micro Devices, Inc. [AMD/ATI] ", ""
                )
                return sanitize(name, 44)
    return ""


def _gpu_name_macos() -> str:
    try:
        out = subprocess.check_output(
            ["system_profiler", "SPDisplaysDataType"],
            stderr=subprocess.DEVNULL,
            text=True,
            timeout=5,
        )
    except (OSError, subprocess.SubprocessError):
        return "Apple GPU"
    m = re.search(r"Chipset Model:\s*(.+)", out)
    if m:
        return sanitize(m.group(1).strip(), 44)
    return "Apple GPU"


def gpu_info() -> Tuple[str, float]:
    """Return (name, utilization%) — util -1 if unknown."""
    for fn in (_gpu_nvidia_nvml, _gpu_nvidia_smi, _gpu_amd_cli):
        got = fn()
        if got:
            return got

    # Name-only fallbacks
    name = ""
    if sys.platform == "win32":
        name = _gpu_name_windows()
    elif sys.platform.startswith("linux"):
        name = _gpu_name_linux()
    elif sys.platform == "darwin":
        name = _gpu_name_macos()
    return (name or "GPU", -1.0)


def build_line() -> str:
    host = sanitize(socket.gethostname(), 36)
    cpu = cpu_name()
    c_pct = cpu_pct()
    gpu, g_pct = gpu_info()
    ram_u, ram_t = ram_gb()
    disk_u, disk_t = disk_gb()
    gpu_field = f"gpu%={g_pct}" if g_pct >= 0 else "gpu%=-1"
    return (
        f"SI|host={host}|cpu={cpu}|cpu%={c_pct}|gpu={gpu}|"
        f"{gpu_field}|ram={ram_u}/{ram_t}|disk={disk_u}/{disk_t}"
    )


_USB_HINT = re.compile(
    r"CH340|CP210|CP2102|CP210x|SILABS|Silicon Labs|USB[- ]?SERIAL|"
    r"USB[- ]?UART|FT232|FTDI|wchusb|QinHeng|ESP32|Espressif|1A86|10C4",
    re.I,
)


def list_port_infos():
    infos = list(list_ports.comports())
    if infos:
        return infos
    # Extra Linux globs when pyserial enumeration is empty
    out = []
    if sys.platform.startswith("linux"):
        for path in sorted(
            glob.glob("/dev/ttyUSB*")
            + glob.glob("/dev/ttyACM*")
            + glob.glob("/dev/tty.usb*")
        ):
            out.append(type("P", (), {"device": path, "description": path, "hwid": ""})())
    return out


def find_ports() -> List[str]:
    return sorted({p.device for p in list_port_infos()})


def score_port(p) -> int:
    blob = f"{getattr(p, 'description', '')} {getattr(p, 'hwid', '')} {p.device}"
    score = 0
    if _USB_HINT.search(blob):
        score += 100
    if sys.platform.startswith("linux") and (
        "ttyUSB" in p.device or "ttyACM" in p.device
    ):
        score += 50
    if sys.platform == "darwin" and "usb" in p.device.lower():
        score += 50
    return score


def pick_port(explicit: str, wait: bool = True) -> str:
    if explicit:
        while True:
            if any(p.device == explicit for p in list_port_infos()) or os.path.exists(
                explicit
            ):
                return explicit
            if not wait:
                raise SystemExit(f"Port not found: {explicit}")
            print(f"Waiting for {explicit} ...", flush=True)
            time.sleep(1.0)

    while True:
        infos = list_port_infos()
        if infos:
            ranked = sorted(infos, key=lambda p: (score_port(p), p.device))
            best = ranked[-1]
            if len(infos) > 1:
                print(
                    "Available ports:",
                    ", ".join(f"{p.device}({getattr(p, 'description', '')})" for p in infos),
                    flush=True,
                )
            print(f"Using port: {best.device}", flush=True)
            return best.device
        if not wait:
            raise SystemExit("No serial ports found. Plug in the CYD USB cable.")
        print("Waiting for CYD USB serial ...", flush=True)
        time.sleep(1.0)


def open_serial(port: str, baud: int):
    ser = serial.Serial()
    ser.port = port
    ser.baudrate = baud
    ser.timeout = 0.2
    ser.write_timeout = 2
    # Avoid toggling DTR/RTS reset on open when possible
    ser.dtr = False
    ser.rts = False
    ser.open()
    return ser


def run_session(port: str, baud: int, interval_ms: int, once: bool) -> None:
    ser = open_serial(port, baud)
    time.sleep(0.35)
    try:
        while True:
            line = build_line()
            ser.write((line + "\n").encode("utf-8", errors="ignore"))
            ser.flush()
            print(line, flush=True)
            if once:
                return
            time.sleep(max(0.2, interval_ms / 1000.0))
    finally:
        try:
            ser.close()
        except Exception:
            pass


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Push host sysinfo to CYD over USB serial (no Python needed if using .exe)"
    )
    ap.add_argument("--port", "-p", default="", help="Serial port (COMx / /dev/ttyUSB0)")
    ap.add_argument("--baud", "-b", type=int, default=460800)
    ap.add_argument("--interval-ms", "-i", type=int, default=1000)
    ap.add_argument("--once", action="store_true", help="Send one line and exit")
    ap.add_argument(
        "--no-wait",
        action="store_true",
        help="Exit immediately if port missing (default: wait / reconnect)",
    )
    args = ap.parse_args()

    # Warm CPU percent baseline
    psutil.cpu_percent(interval=None)

    print(
        f"CYD host agent ({platform.system()} {platform.release()})  "
        f"baud={args.baud} interval={args.interval_ms}ms  Ctrl+C to stop",
        flush=True,
    )

    try:
        while True:
            port = pick_port(args.port, wait=not args.no_wait)
            print(
                f"Sending host info -> {port}",
                flush=True,
            )
            try:
                run_session(port, args.baud, args.interval_ms, args.once)
                if args.once:
                    return 0
            except serial.SerialException as e:
                print(f"Serial lost ({e}); will retry ...", flush=True)
                if args.once or args.no_wait:
                    return 2
                time.sleep(1.0)
            except OSError as e:
                print(f"Port error ({e}); will retry ...", flush=True)
                if args.once or args.no_wait:
                    return 2
                time.sleep(1.0)
    except KeyboardInterrupt:
        print("\nStopped.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())