#!/usr/bin/env python3
"""CMP 170HX checker - shared library (preflight, detection, PCIe checks, BIOS, verify, unlock install, Vulkan)."""
import datetime
import hashlib
import os
import re
import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
HARNESS_DIR = SCRIPT_DIR / "harness"
CARDS_DIR = SCRIPT_DIR / "cards"
UNIDENTIFIED_DIR = SCRIPT_DIR / "unidentified"

REPOS = {
    "170th-Street": "https://github.com/amoghmunikote/170th-Street.git",
    "cmpunlocker": "https://github.com/amoghmunikote/cmpunlocker.git",
}
VERIFY_SH = HARNESS_DIR / "cmpunlocker" / "verify.sh"
INSTALL_SH = HARNESS_DIR / "cmpunlocker" / "install.sh"
GPU_IDS = ("20c2", "2082")
NVIDIA_MODULES = ("nvidia_uvm", "nvidia_drm", "nvidia_modeset", "nvidia")
VULKAN_TEST_SECONDS = 300

SPEED_TO_GEN = {"2.5GT/s": 1, "5GT/s": 2, "8GT/s": 3, "16GT/s": 4, "32GT/s": 5, "64GT/s": 6}

USE_COLOR = sys.stdout.isatty()


class C:
    RESET = "\033[0m" if USE_COLOR else ""
    BOLD = "\033[1m" if USE_COLOR else ""
    DIM = "\033[2m" if USE_COLOR else ""
    GREEN = "\033[32m" if USE_COLOR else ""
    YELLOW = "\033[33m" if USE_COLOR else ""
    RED = "\033[31m" if USE_COLOR else ""
    CYAN = "\033[36m" if USE_COLOR else ""


def info(msg):
    print(f"{C.CYAN}[*]{C.RESET} {msg}", flush=True)


def ok(msg):
    print(f"{C.GREEN}[+]{C.RESET} {msg}", flush=True)


def warn(msg):
    print(f"{C.YELLOW}[!]{C.RESET} {msg}", flush=True)


def err(msg):
    print(f"{C.RED}[-]{C.RESET} {msg}", flush=True)


def step(msg):
    print(f"\n{C.BOLD}== {msg} =={C.RESET}", flush=True)


def ask(prompt, default=""):
    suffix = f" [{default}]" if default else ""
    try:
        val = input(f"{C.BOLD}?{C.RESET} {prompt}{suffix}: ").strip()
    except EOFError:
        return default
    return val if val else default


def now_iso():
    return datetime.datetime.now().isoformat(timespec="seconds")


def sanitize(name):
    return re.sub(r"[^A-Za-z0-9._-]+", "_", name).strip("_") or "UNKNOWN"


# ---------------------------------------------------------------- subprocess

def run(cmd, timeout=60, env=None, cwd=None):
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env, cwd=cwd)
        return {"rc": p.returncode, "out": p.stdout, "err": p.stderr, "cmd": cmd}
    except FileNotFoundError:
        return {"rc": 127, "out": "", "err": f"not found: {cmd[0]}", "cmd": cmd}
    except subprocess.TimeoutExpired as e:
        out = e.stdout.decode(errors="replace") if isinstance(e.stdout, bytes) else (e.stdout or "")
        return {"rc": 124, "out": out, "err": "timeout", "cmd": cmd}


def run_live(cmd, cwd=None, timeout=600, log_path=None):
    """Run cmd while echoing every output line to the terminal as it arrives.

    The complete output is also written to log_path (operator keeps the full log,
    the terminal keeps the progress). stderr is merged into stdout, so out holds
    everything. Returns the same {rc, out, err, cmd} shape as run(): rc 127 =
    not found, rc 124 = timeout.
    """
    try:
        p = subprocess.Popen(cmd, cwd=cwd, stdout=subprocess.PIPE,
                             stderr=subprocess.STDOUT, text=True, bufsize=1)
    except FileNotFoundError:
        return {"rc": 127, "out": "", "err": f"not found: {cmd[0]}", "cmd": cmd}
    except OSError as e:
        return {"rc": 126, "out": "", "err": str(e), "cmd": cmd}

    lines, done = [], threading.Event()

    def _pump():
        try:
            for line in p.stdout or "":
                lines.append(line)
                print(line, end="", flush=True)
        except (ValueError, OSError):
            pass
        done.set()

    pump = threading.Thread(target=_pump, daemon=True)
    pump.start()

    timed_out = interrupted = False
    try:
        finished = done.wait(timeout)
    except KeyboardInterrupt:
        finished, interrupted = True, True
    if not finished:
        timed_out = True
    if timed_out or interrupted:
        try:
            p.kill()
        except OSError:
            pass
        pump.join(timeout=15)
    rc = p.wait()
    out = "".join(lines)
    err = ""
    if timed_out:
        rc, err = 124, "timeout"
    elif interrupted:
        err = "interrupted"

    if log_path:
        path = Path(log_path)
        try:
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_text(f"$ {' '.join(str(c) for c in cmd)}\n"
                            f"cwd: {cwd or os.getcwd()}\n"
                            f"{out}\nrc={rc}{f' ({err})' if err else ''}\n")
        except OSError as e:
            warn(f"could not write log {path}: {e}")
    if interrupted:
        raise KeyboardInterrupt
    return {"rc": rc, "out": out, "err": err, "cmd": cmd}


def as_root(cmd, timeout=60, cwd=None):
    if os.geteuid() == 0:
        return run(cmd, timeout=timeout, cwd=cwd)
    if shutil.which("sudo"):
        r = run(["sudo", "-n"] + cmd, timeout=timeout, cwd=cwd)
        if r["rc"] != 0 and "password is required" in (r["err"] + r["out"]).lower():
            r["err"] = "sudo requires a password (sudo -n failed)"
        return r
    return {"rc": 1, "out": "", "err": "no root: neither uid 0 nor sudo available", "cmd": cmd}


def which(name):
    return shutil.which(name)


def find_nvflash():
    """nvflash binary: repo-bundled copy first (no system install needed), then PATH."""
    import platform
    arch_map = {"x86_64": "x64", "amd64": "x64", "aarch64": "aarch64",
                "ppc64le": "ppc64", "ppc64": "ppc64", "i686": "x86", "i386": "x86"}
    arch = arch_map.get(platform.machine(), "x64")
    candidates = [
        *sorted(SCRIPT_DIR.glob(f"nvflash_*_linux/{arch}/nvflash")),
        *sorted(SCRIPT_DIR.glob("nvflash_*_linux/x64/nvflash")),
        SCRIPT_DIR / "nvflash",
    ]
    for p in candidates:
        if p.is_file() and os.access(p, os.X_OK):
            return str(p)
    return which("nvflash")


def have_root():
    return os.geteuid() == 0 or bool(shutil.which("sudo"))


# ---------------------------------------------------------------- preflight

PM_CANDIDATES = [
    ("apt-get", ["apt-get", "install", "-y"]),
    ("dnf", ["dnf", "install", "-y"]),
    ("yum", ["yum", "install", "-y"]),
    ("pacman", ["pacman", "-S", "--noconfirm"]),
    ("zypper", ["zypper", "--non-interactive", "install"]),
    ("xbps-install", ["xbps-install", "-y"]),
]


def detect_pm():
    for name, cmd in PM_CANDIDATES:
        if which(name):
            return cmd
    return None


def install_packages(packages):
    pm = detect_pm()
    if not pm:
        return False, "no supported package manager found (apt/dnf/pacman/zypper/xbps)"
    installed, log = [], []
    for pkg in packages:
        r = as_root(pm + [pkg], timeout=300)
        log.append(f"$ {' '.join(pm + [pkg])}\nrc={r['rc']}\n{r['out'][-1500:]}{r['err'][-800:]}")
        if r["rc"] == 0:
            installed.append(pkg)
        else:
            warn(f"could not install '{pkg}' via {pm[0]} (rc={r['rc']})")
    return bool(installed), "\n".join(log)


def ensure_repos():
    for name, url in REPOS.items():
        dest = HARNESS_DIR / name
        if (dest / ".git").exists():
            continue
        info(f"cloning {name} ...")
        r = run(["git", "clone", "--depth", "1", url, str(dest)], timeout=300)
        if r["rc"] != 0:
            err(f"git clone failed for {name}: {r['err'].strip()}")
            return False
        ok(f"cloned {name}")
    return True


TOOL_CHECKS = {
    "git": ("critical", "clone harness repos"),
    "lspci": ("critical", "PCIe checks (package: pciutils)"),
    "setpci": ("critical", "link generation probe (package: pciutils)"),
    "nvidia-smi": ("critical-if-card-present", "GPU state / cmpunlocker verify"),
    "nvflash": ("critical", "VBIOS backup + serial (bundled in repo, or on PATH)"),
    "dmesg": ("optional", "kernel log scrape"),
    "journalctl": ("optional", "kernel log scrape fallback"),
    "vulkaninfo": ("optional", "vulkan presence (package: vulkan-tools)"),
    "vkmark": ("optional", "5-min VRAM stress (package: vkmark)"),
    "vkcube": ("optional", "VRAM stress fallback (package: vulkan-tools)"),
}
PKG_FOR = {"lspci": "pciutils", "setpci": "pciutils", "vulkaninfo": "vulkan-tools",
           "vkcube": "vulkan-tools", "vkmark": "vkmark", "git": "git"}


def os_release():
    data = read_file("/etc/os-release") or ""
    out = {}
    for line in data.splitlines():
        if "=" in line:
            k, v = line.split("=", 1)
            out[k] = v.strip().strip('"')
    return out


def sudo_precondition(interactive=True):
    """Precondition for package installation (--preflight): root or working sudo.

    Tries passwordless sudo first; falls back to an interactive `sudo -v` so the
    credentials get cached, after which plain sudo -n calls succeed.
    """
    if os.geteuid() == 0:
        return True, "running as root"
    if not shutil.which("sudo"):
        return False, "sudo is not installed"
    r = run(["sudo", "-n", "true"], timeout=15)
    if r["rc"] == 0:
        return True, "sudo available (credentials cached)"
    if interactive and sys.stdin.isatty():
        info("sudo authentication required for package installation - enter your password")
        try:
            p = subprocess.run(["sudo", "-v"])
        except KeyboardInterrupt:
            return False, "interrupted"
        if p.returncode == 0:
            return True, "sudo authenticated just now"
        return False, "sudo authentication failed"
    return False, "sudo requires a password (non-interactive session)"


def install_nvidia_driver():
    """Ubuntu: prefer 'ubuntu-drivers install', else the nvidia-driver metapackage."""
    if which("ubuntu-drivers"):
        info("installing NVIDIA driver via 'ubuntu-drivers install'")
        r = as_root(["ubuntu-drivers", "install"], timeout=900)
    else:
        pm = detect_pm()
        if not pm:
            return False, "no package manager available for driver install"
        info("installing NVIDIA driver metapackage 'nvidia-driver'")
        r = as_root(pm + ["nvidia-driver"], timeout=900)
    return r["rc"] == 0, (r["out"][-2000:] + r["err"][-1000:])


def preflight(auto_install=True, do_clone=True):
    step("Preflight")
    report = {"tools": {}, "root": have_root()}
    critical_missing, optional_missing = [], []

    rel = os_release()
    report["os"] = {"id": rel.get("ID", ""), "version_id": rel.get("VERSION_ID", ""),
                    "pretty": rel.get("PRETTY_NAME", "unknown")}
    info(f"OS: {report['os']['pretty']}")

    for tool, (level, why) in TOOL_CHECKS.items():
        path = find_nvflash() if tool == "nvflash" else which(tool)
        report["tools"][tool] = {"found": bool(path), "path": path, "level": level, "why": why}
        if not path:
            (critical_missing if level.startswith("critical") else optional_missing).append(tool)

    want_nvidia = False
    if auto_install:
        pkgs = sorted({PKG_FOR[t] for t in critical_missing + optional_missing if t in PKG_FOR})
        if "nvidia-smi" in critical_missing and sys.stdin.isatty():
            if ask("nvidia-smi missing - install the NVIDIA driver now? (y/N)", "n").lower().startswith("y"):
                want_nvidia = True
        if pkgs or want_nvidia:
            ok_sudo, why = sudo_precondition()
            report["sudo"] = {"ok": ok_sudo, "why": why}
            if not ok_sudo:
                err(f"sudo PRECONDITION FAILED: {why}")
                err("--preflight package installation requires sudo privileges. Fix one of:")
                err("  - run 'sudo -v' first, then re-run this command")
                err("  - run the script as root")
                err("  - re-run with --no-install to skip package installation")
                report["install_skipped"] = why
            else:
                info(f"sudo precondition OK: {why}")
                pm = detect_pm()
                if pkgs:
                    if pm and pm[0] == "apt-get":
                        info("apt-get update ...")
                        as_root(["apt-get", "update"], timeout=300)
                    if pm:
                        info(f"installing: {' '.join(pkgs)}")
                        _, log = install_packages(pkgs)
                        report["install_log"] = log
                    else:
                        report["install_log"] = "no package manager - manual install required"
                        warn(f"no package manager found - install manually: {' '.join(pkgs)}")
                if want_nvidia:
                    if pm or which("ubuntu-drivers"):
                        ok_d, log_d = install_nvidia_driver()
                        report["nvidia_driver"] = {"ok": ok_d, "log": log_d}
                        if ok_d:
                            ok("NVIDIA driver installed - reboot, then re-run the check")
                        else:
                            warn("NVIDIA driver install failed (see report['nvidia_driver']['log'])")
                    else:
                        warn("no package manager - cannot install NVIDIA driver automatically")
                for tool in list(critical_missing) + list(optional_missing):
                    if which(tool):
                        if tool in critical_missing:
                            critical_missing.remove(tool)
                        if tool in optional_missing:
                            optional_missing.remove(tool)
                        report["tools"][tool].update(found=True, path=which(tool))
        else:
            report["sudo"] = {"ok": True, "why": "nothing to install"}

    if do_clone:
        ensure_repos()
        report["repos_ok"] = all((HARNESS_DIR / n / ".git").exists() for n in REPOS)

    report["verify_sh_exists"] = VERIFY_SH.exists()

    for t, st in report["tools"].items():
        if st["found"]:
            mark = f"{C.GREEN}OK{C.RESET}"
        elif st["level"].startswith("critical"):
            mark = f"{C.RED}MISSING{C.RESET}"
        else:
            mark = f"{C.YELLOW}missing (optional){C.RESET}"
        print(f"   {t:<14} {mark}  {C.DIM}{st['why']}{C.RESET}")

    if not find_nvflash():
        warn("nvflash not found - VBIOS backup will fall back to sysfs ROM, serial to manual entry")
        print("     Bundle it in the repo: nvflash_<version>_linux/x64/nvflash (chmod +x),")
        print("     or pass --nvflash /path/to/nvflash, or place it on PATH.")
    else:
        info(f"nvflash: {find_nvflash()}")
    if not which("nvidia-smi"):
        warn("nvidia-smi missing - cmpunlocker verify will be skipped (install nvidia-open driver)")
    if not report["root"]:
        warn("no sudo and not root: privileged steps (rescan, retrain, rom, dmesg) will be skipped")

    blocking = [t for t in ("git", "lspci", "setpci") if not which(t)]
    report["blocking_missing"] = blocking
    report["preflight_ok"] = not blocking
    if blocking:
        err("blocking tools still missing: " + ", ".join(blocking))
    else:
        ok("preflight complete")
    return report


# ---------------------------------------------------------------- GPU detect

def detect_gpus():
    bdfs = []
    for dev_id in GPU_IDS:
        r = run(["lspci", "-D", "-d", f"10de:{dev_id}"])
        for line in r["out"].splitlines():
            parts = line.split()
            if not parts:
                continue
            bdf = parts[0]
            if re.match(r"^[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9a-f]$", bdf):
                bdfs.append(bdf.lower())
    return sorted(set(bdfs))


def rescan_bus():
    r = as_root(["sh", "-c", "echo 1 > /sys/bus/pci/rescan"], timeout=15)
    if r["rc"] == 0:
        time.sleep(2)
    return r["rc"] == 0


def sysfs(bdf):
    return Path("/sys/bus/pci/devices") / bdf


def read_file(path, binary=False):
    try:
        with open(path, "rb" if binary else "r", errors="replace") as f:
            return f.read()
    except (PermissionError, FileNotFoundError, OSError):
        return None


# ---------------------------------------------------------------- PCIe checks

def parse_link(text):
    out = {"lnksta_speed": None, "lnksta_width": None, "lnksta_gen": None,
           "lnkcap_speed": None, "lnkcap_width": None, "lnksta_raw": None}
    for line in text.splitlines():
        s = line.strip()
        if s.startswith("LnkSta:") and not s.startswith("LnkSta2"):
            out["lnksta_raw"] = s
            m = re.search(r"Speed\s+([^,\s]+)", s)
            if m:
                out["lnksta_speed"] = m.group(1)
                out["lnksta_gen"] = SPEED_TO_GEN.get(m.group(1))
            m = re.search(r"Width\s+(x\d+)", s)
            if m:
                out["lnksta_width"] = m.group(1)
        elif s.startswith("LnkCap:"):
            m = re.search(r"Speed\s+([^,\s]+)", s)
            if m:
                out["lnkcap_speed"] = m.group(1)
            m = re.search(r"Width\s+(x\d+)", s)
            if m:
                out["lnkcap_width"] = m.group(1)
    return out


def parse_aer_flags(text):
    res = {"CESta": [], "UESta": []}
    for line in text.splitlines():
        s = line.strip()
        for key in ("CESta", "UESta"):
            if s.startswith(key + ":"):
                body = s.split(":", 1)[1]
                res[key].extend(re.findall(r"([A-Za-z0-9_]+)\+", body))
    return res


def parent_bdf(bdf):
    path = Path(os.path.realpath(sysfs(bdf)))
    while True:
        path = path.parent
        name = path.name
        if re.match(r"^[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9a-f]$", name):
            return name.lower()
        if name.startswith("pci") or path.parent == path:
            return None


def aer_snapshot(bdf):
    snap = {"aer_sysfs_correctable": None, "aer_sysfs_fatal": None,
            "aer_correctable_sum": 0, "ce_active": [], "ue_active": [],
            "parent_bdf": None, "parent_ce_active": []}
    base = sysfs(bdf)
    data = read_file(base / "aer_dev_correctable")
    if data is not None:
        snap["aer_sysfs_correctable"] = data.strip()
        snap["aer_correctable_sum"] = sum(int(x) for x in re.findall(r"\d+", data))
    fatal = read_file(base / "aer_dev_fatal")
    if fatal is not None:
        snap["aer_sysfs_fatal"] = fatal.strip()
    r = run(["lspci", "-vvv", "-D", "-s", bdf], timeout=30)
    if r["rc"] == 0:
        aer = parse_aer_flags(r["out"])
        snap["ce_active"] = aer["CESta"]
        snap["ue_active"] = aer["UESta"]
    parent = parent_bdf(bdf)
    if parent:
        snap["parent_bdf"] = parent
        rp = run(["lspci", "-vvv", "-D", "-s", parent], timeout=30)
        if rp["rc"] == 0:
            snap["parent_ce_active"] = parse_aer_flags(rp["out"])["CESta"]
    return snap


def link_gen_setpci(bdf):
    r = as_root(["setpci", "-s", bdf, "CAP_EXP+12.w"], timeout=15)
    val = r["out"].strip()
    if re.match(r"^[0-9a-fA-F]{4}$", val):
        return int(val, 16) & 0x0F
    return None


def check_link(bdf):
    r = run(["lspci", "-vvv", "-D", "-s", bdf], timeout=30)
    link = parse_link(r["out"]) if r["rc"] == 0 else {}
    link["setpci_current_gen"] = link_gen_setpci(bdf)
    link["lspci_rc"] = r["rc"]
    return link, r["out"]


def dump_config(bdf):
    r = run(["lspci", "-nnvvxxxx", "-D", "-s", bdf], timeout=30)
    return r["out"] if r["rc"] == 0 else f"lspci failed rc={r['rc']}: {r['err']}"


def retrain_test(bdf):
    result = {"attempted": False, "ok": None, "note": "", "width_before": None, "width_after": None}
    link_before, _ = check_link(bdf)
    result["width_before"] = link_before.get("lnksta_width")
    node = sysfs(bdf) / "link" / "retrain"
    if not node.exists():
        result["note"] = "sysfs link/retrain not available on this kernel"
        return result
    result["attempted"] = True
    r = as_root(["sh", "-c", f"echo 1 > {node}"], timeout=15)
    if r["rc"] != 0:
        result["note"] = f"retrain write failed: {r['err'].strip()}"
        return result
    time.sleep(3)
    if bdf not in detect_gpus():
        result.update(ok=False, note="device vanished after retrain")
        return result
    link_after, _ = check_link(bdf)
    result["width_after"] = link_after.get("lnksta_width")
    result["ok"] = result["width_after"] == result["width_before"] and result["width_after"] is not None
    result["note"] = f"{result['width_before']} -> {result['width_after']}"
    return result


def rescan_test(bdf):
    result = {"attempted": True, "ok": None, "note": ""}
    if not rescan_bus():
        result["note"] = "bus rescan not permitted (no root)"
        return result
    time.sleep(1)
    present = bdf in detect_gpus()
    result["ok"] = present
    result["note"] = "re-enumerated at same BDF" if present else "device missing after rescan"
    return result


def dmesg_capture(bdf):
    text, source = "", None
    r = as_root(["dmesg", "--ctime"], timeout=20)
    if r["rc"] == 0 and r["out"]:
        text, source = r["out"], "dmesg"
    else:
        r2 = run(["journalctl", "-k", "-b", "--no-pager", "-o", "short-precise"], timeout=25)
        if r2["rc"] == 0 and r2["out"]:
            text, source = r2["out"], "journalctl"
    if not text:
        return {"source": None, "related": [], "errors": [], "note": "kernel log unavailable (no root?)"}
    short = bdf[5:]
    related, errors = [], []
    for line in text.splitlines():
        low = line.lower()
        if (bdf in line or short in line or "10de" in low or "nvrm" in low
                or "xid" in low or "aer:" in low or "pcieport" in low):
            related.append(line)
            if any(k in low for k in ("xid", "error", "aer:", "fail", "corrupted", "fallen off")):
                errors.append(line)
    return {"source": source, "related": related[-200:], "errors": errors[-50:],
            "note": f"{len(related)} related lines"}


# ---------------------------------------------------------------- BIOS / serial

def nvflash_bin(args):
    if getattr(args, "nvflash", None):
        return args.nvflash
    return find_nvflash()


def nvflash_list(nf):
    for flags in (["--list"], ["-l"]):
        r = as_root([nf] + flags, timeout=40)
        if r["rc"] == 0 and r["out"].strip():
            return r["out"]
    return None


def get_serial(nf, args):
    if getattr(args, "serial", None):
        return args.serial, "cli-override"
    if nf:
        listing = nvflash_list(nf)
        if listing:
            m = re.search(r"UUID:\s*(GPU-[0-9a-fA-F\-]+)", listing, re.I)
            if m:
                info(f"serial auto-read from nvflash UUID: {m.group(1)}")
                return m.group(1), "nvflash-uuid"
            m = re.search(r"Serial(?:Number)?\s*[:=]\s*(\S+)", listing, re.I)
            if m:
                info(f"serial auto-read from nvflash: {m.group(1)}")
                return m.group(1), "nvflash-list"
        q = as_root([nf, "--query"], timeout=40)
        m = re.search(r"Serial(?:Number)?\s*[:=]\s*(\S+)", q["out"], re.I)
        if m:
            info(f"serial auto-read from nvflash --query: {m.group(1)}")
            return m.group(1), "nvflash-query"
        warn("nvflash present but exposes no serial/UUID for this card")
    else:
        warn("nvflash unavailable - cannot auto-read serial")
    return "N/A", "unavailable"


def save_bios(nf, out_path, args):
    result = {"method": None, "path": None, "sha256": None, "size": 0, "ok": False, "note": ""}
    if getattr(args, "skip_bios", False):
        result["note"] = "bios backup skipped (--skip-bios)"
        return result
    notes = []
    if nf:
        ok_save, note = try_nvflash_save(nf, out_path)
        if ok_save:
            data = out_path.read_bytes()
            result.update(method="nvflash", path=str(out_path), size=len(data),
                          sha256=hashlib.sha256(data).hexdigest(), ok=True)
            return result
        notes.append(note)
    else:
        notes.append("nvflash unavailable")

    for bdf in detect_gpus():
        rom = sysfs(bdf) / "rom"
        if not rom.exists():
            notes.append(f"no sysfs rom node for {bdf} (ROM BAR not exposed - fallback unavailable)")
            continue
        if as_root(["sh", "-c", f"echo 1 > {rom}"], timeout=15)["rc"] != 0:
            notes.append(f"sysfs rom enable denied for {bdf}")
            continue
        data = read_file(rom, binary=True)
        as_root(["sh", "-c", f"echo 0 > {rom}"], timeout=15)
        if data and len(data) > 512 and data.strip(b"\x00") != b"":
            out_path.write_bytes(data)
            result.update(method="sysfs-rom", path=str(out_path), size=len(data),
                          sha256=hashlib.sha256(data).hexdigest(), ok=True)
            return result
        notes.append(f"sysfs rom empty/all-zero for {bdf}")
    if any("no sysfs rom node" in n for n in notes):
        notes.append(f"host lockdown={lockdown_state()} secureboot={secure_boot_state()}")
    result["note"] = " | ".join(notes)

    # Opt-in last resort: unload the driver, retry nvflash --save, always reload.
    if nf and any("driver loaded" in n for n in notes):
        if getattr(args, "bios_unload", False) or (
                sys.stdin.isatty()
                and ask("Unload NVIDIA driver and retry nvflash --save now? (y/N)",
                        "n").lower().startswith("y")):
            result["note"] += " | " + retry_save_unloaded(nf, out_path, result)
        else:
            result["note"] += (" | skipped driver-unload retry (use --bios-unload; "
                               "local console may freeze while the driver is out)")
    return result


def try_nvflash_save(nf, out_path):
    """One nvflash --save attempt. Returns (ok, note)."""
    try:
        if out_path.exists():
            out_path.unlink()
    except OSError:
        pass
    r = as_root([nf, "--save", str(out_path)], timeout=90)
    text = r["out"] + r["err"]
    if r["rc"] == 0 and out_path.exists() and out_path.stat().st_size > 0:
        data = out_path.read_bytes()
        if data.strip(b"\x00") == b"":
            return False, "nvflash wrote an all-zero image"
        return True, ""
    if "unload" in text.lower() and "kernel driver" in text.lower():
        return False, (f"nvflash --save refused: NVIDIA driver loaded (rc={r['rc']}) - "
                       "unload it first or use --bios-unload")
    return False, f"nvflash --save failed rc={r['rc']}: {text[-200:].strip()}"


def retry_save_unloaded(nf, out_path, result):
    """Unload NVIDIA modules, retry nvflash --save, always reload. Returns note."""
    warn("unloading NVIDIA driver - local display may freeze until reload")
    _, unload_log = unload_nvidia_modules()
    try:
        ok_save, note = try_nvflash_save(nf, out_path)
        if ok_save:
            data = out_path.read_bytes()
            result.update(method="nvflash-unload", path=str(out_path), size=len(data),
                          sha256=hashlib.sha256(data).hexdigest(), ok=True)
            return "saved via nvflash after driver unload"
        return f"nvflash --save still failed after unload: {note}"
    finally:
        ok_reload, reload_log = load_nvidia_modules()
        for line in unload_log + reload_log:
            info(line)
        if ok_reload and which("nvidia-smi"):
            ok("NVIDIA driver reloaded")
        else:
            err("NVIDIA driver reload incomplete - reboot the host")


def unload_nvidia_modules():
    """Best-effort rmmod of the NVIDIA stack. Returns (all_ok, log_lines)."""
    log, all_ok = [], True
    for mod in NVIDIA_MODULES:
        r = as_root(["rmmod", mod], timeout=60)
        log.append(f"rmmod {mod}: rc={r['rc']} {(r['err'] or r['out']).strip()[:160]}")
        if r["rc"] != 0:
            all_ok = False
    return all_ok, log


def load_nvidia_modules():
    """Reload the NVIDIA stack in reverse order. Returns (all_ok, log_lines)."""
    log, all_ok = [], True
    for mod in reversed(NVIDIA_MODULES):
        r = as_root(["modprobe", mod], timeout=120)
        log.append(f"modprobe {mod}: rc={r['rc']} {(r['err'] or r['out']).strip()[:160]}")
        if r["rc"] != 0:
            all_ok = False
    return all_ok, log


def lockdown_state():
    data = read_file("/sys/kernel/security/lockdown")
    return data.strip() if data else "lockdown file unreadable"


def secure_boot_state():
    if not which("mokutil"):
        return "mokutil not installed"
    r = run(["mokutil", "--sb-state"], timeout=15)
    lines = (r["out"] or r["err"]).strip().splitlines()
    return lines[0].strip() if lines else "mokutil gave no output"


# ---------------------------------------------------------------- nvidia / verify

def nvidia_smi_info():
    if not which("nvidia-smi"):
        return {"available": False}
    q = run(["nvidia-smi", "--query-gpu=name,uuid,memory.total,pcie.link.gen.current,"
                          "pcie.link.gen.max,pcie.link.width.current,temperature.gpu,driver_version",
             "--format=csv,noheader"], timeout=30)
    if q["rc"] != 0:
        return {"available": False, "error": (q["err"] or q["out"])[-300:]}
    rows = [[f.strip() for f in line.split(",")] for line in q["out"].strip().splitlines()]
    return {"available": True, "gpus": rows}


def get_gpu_uuid():
    """First GPU UUID from nvidia-smi (e.g. GPU-xxxx...), or None."""
    if not which("nvidia-smi"):
        return None
    q = run(["nvidia-smi", "--query-gpu=uuid", "--format=csv,noheader"], timeout=30)
    if q["rc"] != 0:
        return None
    for line in q["out"].splitlines():
        line = line.strip()
        if line.startswith("GPU-"):
            return line
    return None


def run_verify_unlocker():
    if not VERIFY_SH.exists():
        return {"ran": False, "note": "harness/cmpunlocker not cloned"}
    if not which("nvidia-smi"):
        return {"ran": False, "note": "nvidia-smi missing - unlock verify skipped"}
    r = as_root(["bash", str(VERIFY_SH)], timeout=180)
    lines = [l for l in (r["out"] + r["err"]).splitlines() if l.strip()]
    status = "unknown"
    if r["rc"] == 0 and any("unlocked memory" in l.lower() for l in lines):
        status = "unlocked-ok"
    elif "STOCK" in r["out"]:
        status = "stock-memory"
    elif "MISSING" in r["out"]:
        status = "gpu-missing-in-smi"
    elif r["rc"] != 0:
        status = f"verify-failed-rc{r['rc']}"
    return {"ran": True, "rc": r["rc"], "status": status, "tail": lines[-30:]}


# ---------------------------------------------------------------- unlock (opt-in install)

def update_cmpunlocker():
    """Bring harness/cmpunlocker up to date with upstream main before installing.

    A failed `git pull --ff-only` (diverged checkout, offline host, ...) is not
    fatal: we warn, keep the existing checkout, and let the install proceed.
    """
    dest = HARNESS_DIR / "cmpunlocker"
    if not (dest / ".git").exists():
        info("harness/cmpunlocker missing - cloning ...")
        if not ensure_repos():
            return {"ok": False, "note": "harness/cmpunlocker not cloned and git clone failed"}
        return {"ok": True, "note": "harness/cmpunlocker cloned"}
    r = run(["git", "-C", str(dest), "pull", "--ff-only"], timeout=180)
    if r["rc"] != 0:
        msg = (r["err"] or r["out"]).strip().splitlines()
        msg = msg[-1] if msg else f"rc={r['rc']}"
        warn(f"git pull --ff-only failed: {msg} - continuing with the existing checkout")
        return {"ok": False, "note": f"git pull --ff-only failed ({msg}) - using existing checkout"}
    last = (r["out"].strip().splitlines() or ["already up to date"])[-1]
    ok(f"cmpunlocker harness updated: {last}")
    return {"ok": True, "note": f"updated: {last}"}


def run_installer(log_path):
    """Run install.sh from the harness clone (opt-in, root required).

    Output streams live to the terminal and is fully saved to log_path.
    Returns {"attempted", "ok", "rc", "note", "log"}; attempted=False when a
    precondition (clone present, sudo usable) fails.
    """
    log_path = Path(log_path)
    result = {"attempted": False, "ok": False, "rc": None, "note": "", "log": str(log_path)}
    if not INSTALL_SH.exists():
        result["note"] = "harness/cmpunlocker not cloned"
        warn(f"{result['note']} - run --preflight first")
        return result
    ok_sudo, why = sudo_precondition()
    if not ok_sudo:
        result["note"] = f"sudo precondition failed: {why}"
        warn(result["note"])
        return result

    cmd = ["bash", str(INSTALL_SH)]
    if os.geteuid() != 0:
        cmd = ["sudo", "-n"] + cmd   # credentials cached by sudo_precondition()
    info(f"running install.sh in {INSTALL_SH.parent} - kernel module build takes 5-15 min, "
         f"live output -> {log_path}")
    r = run_live(cmd, cwd=str(INSTALL_SH.parent), timeout=1800, log_path=log_path)
    result.update(attempted=True, rc=r["rc"])
    if r["rc"] == 0:
        result.update(ok=True, note="install.sh completed")
        ok("cmpunlocker installed - cold reboot required: sudo shutdown -h now (then power on), "
           "then re-run this script for the final PCIe check")
    else:
        tail = (r["err"] or r["out"]).strip().splitlines()
        tail = tail[-1] if tail else ""
        result["note"] = f"install.sh failed rc={r['rc']}" + (f": {tail}" if tail else "")
        warn(f"{result['note']} (see {log_path})")
    return result


# ---------------------------------------------------------------- Vulkan test

def vulkan_tool():
    for t in ("vkmark", "vkcube", "vulkaninfo"):
        if which(t):
            return t
    return None


def run_vulkan_test(seconds=VULKAN_TEST_SECONDS):
    result = {"tool": None, "ran": False, "requested": True, "duration_s": seconds,
              "rc": None, "note": "", "log_tail": []}
    tool = vulkan_tool()
    if not tool:
        result["note"] = "no vulkan tool available (install vkmark or vulkan-tools)"
        warn(result["note"])
        return result
    result["tool"] = tool
    info(f"running {tool} for {seconds}s VRAM stress - press Ctrl+C to abort early")
    output = []
    deadline = time.time() + seconds
    try:
        if tool == "vkmark":
            while time.time() < deadline:
                r = run(["vkmark", "--benchmark"], timeout=60)
                output.append(r["out"])
                if r["rc"] not in (0, 124):
                    break
        elif tool == "vkcube":
            p = subprocess.Popen(["vkcube"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
            try:
                p.wait(timeout=seconds)
            except subprocess.TimeoutExpired:
                p.terminate()
                try:
                    p.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    p.kill()
            output.append(p.communicate()[0] or "")
            result["rc"] = p.returncode
        else:
            r = run(["vulkaninfo"], timeout=seconds)
            output.append(r["out"] + r["err"])
            result["rc"] = r["rc"]
    except KeyboardInterrupt:
        result["note"] = "aborted by user"
        warn("vulkan test aborted")
    text = "\n".join(output)
    result["ran"] = True
    result["log_tail"] = [l for l in text.splitlines() if l.strip()][-40:]
    bad = [l for l in text.splitlines()
           if re.search(r"(error|fail|out of memory|oom|corrupt)", l, re.I)]
    result["tool_errors"] = bad[-20:]
    ok(f"{tool} stress finished ({seconds}s)")
    return result