#!/usr/bin/env python3
"""CMP 170HX setup, unlock, PCIe checks, and optional diagnostics."""
import datetime
import hashlib
import io
import json
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
import zipfile
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"

UNLOCKER_COMMIT = "f51da030ef4260656288edf07e95ae45f339aea8"  # upstream v0.3
UNLOCKER_ZIP_SHA256 = "a10b226ed4667fb2ae4c9c5ad1575938a85ceb707f8c40ea146887778930ce21"
REPOS = {"cmpunlocker": ("commits", UNLOCKER_COMMIT)}
VERIFY_SH = HARNESS_DIR / "cmpunlocker" / "verify.sh"
INSTALL_SH = HARNESS_DIR / "cmpunlocker" / "install.sh"
PENDING_FILE = HARNESS_DIR / "unlock_pending.json"
GPU_IDS = ("20c2", "2082")

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}[ERROR] {msg}{C.RESET}", 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, errors="replace",
                           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 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 run_logged_as_root(cmd, log_path, timeout, cwd=None):
    """Show a long-running command live while keeping its full output on disk."""
    actual = cmd if os.geteuid() == 0 else ["sudo", "-n"] + cmd
    log_path.parent.mkdir(parents=True, exist_ok=True)
    tail = bytearray()
    try:
        with log_path.open("wb") as log:
            process = subprocess.Popen(actual, stdout=subprocess.PIPE,
                                       stderr=subprocess.STDOUT, cwd=cwd,
                                       start_new_session=True)

            def copy_output():
                show_console = True
                while True:
                    try:
                        chunk = os.read(process.stdout.fileno(), 4096)
                    except OSError:
                        break
                    if not chunk:
                        break
                    log.write(chunk)
                    log.flush()
                    if show_console:
                        try:
                            if hasattr(sys.stdout, "buffer"):
                                sys.stdout.buffer.write(chunk)
                            else:
                                sys.stdout.write(chunk.decode(errors="replace"))
                            sys.stdout.flush()
                        except BrokenPipeError:
                            show_console = False
                    tail.extend(chunk)
                    if len(tail) > 4096:
                        del tail[:-4096]

            def stop_group(sig):
                try:
                    os.killpg(process.pid, sig)
                except ProcessLookupError:
                    pass

            reader = threading.Thread(target=copy_output, daemon=True)
            reader.start()
            try:
                rc = process.wait(timeout=timeout)
            except subprocess.TimeoutExpired:
                stop_group(signal.SIGTERM)
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    stop_group(signal.SIGKILL)
                    process.wait()
                rc = 124
            except KeyboardInterrupt:
                stop_group(signal.SIGTERM)
                try:
                    process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    stop_group(signal.SIGKILL)
                    process.wait()
                raise
            finally:
                reader.join(timeout=10)
                if reader.is_alive():
                    stop_group(signal.SIGTERM)
                    reader.join(timeout=5)
                if reader.is_alive():
                    stop_group(signal.SIGKILL)
                    reader.join(timeout=5)
                process.stdout.close()
                reader.join(timeout=2)
    except OSError as exc:
        return {"rc": 127, "out": "", "err": str(exc), "cmd": actual,
                "log": str(log_path)}
    return {"rc": rc, "out": tail.decode(errors="replace"), "err": "",
            "cmd": actual, "log": str(log_path)}


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

def ensure_repos():
    """Fetch upstream source archives without changing any Git remotes."""
    HARNESS_DIR.mkdir(exist_ok=True)
    complete = True
    for name, (kind, ref) in REPOS.items():
        dest = HARNESS_DIR / name
        if (dest / ("install.sh" if name == "cmpunlocker" else "README.md")).exists():
            continue
        suffix = f"archive/{ref}.zip" if kind == "commits" else f"archive/refs/{kind}/{ref}.zip"
        url = f"https://github.com/amoghmunikote/{name}/{suffix}"
        info(f"downloading {name} ({ref}) ...")
        try:
            with urllib.request.urlopen(url, timeout=60) as response:
                archive = response.read()
            if name == "cmpunlocker" and hashlib.sha256(archive).hexdigest() != UNLOCKER_ZIP_SHA256:
                raise ValueError("cmpunlocker archive checksum differs from pinned v0.3 source")
            with tempfile.TemporaryDirectory(dir=HARNESS_DIR) as temp:
                staging = Path(temp) / name
                staging.mkdir()
                archive_root = f"{name}-{ref[1:] if kind == 'tags' and ref.startswith('v') else ref}"
                with zipfile.ZipFile(io.BytesIO(archive)) as zf:
                    for member in zf.infolist():
                        parts = Path(member.filename).parts
                        if parts == (archive_root,):
                            continue
                        if len(parts) < 2 or ".." in parts or parts[0] != archive_root:
                            raise ValueError(f"unsafe archive path: {member.filename}")
                        if (member.external_attr >> 16) & 0o170000 == 0o120000:
                            raise ValueError(f"archive symlink: {member.filename}")
                        path = staging.joinpath(*parts[1:])
                        if member.is_dir():
                            path.mkdir(parents=True, exist_ok=True)
                        else:
                            path.parent.mkdir(parents=True, exist_ok=True)
                            with zf.open(member) as source, path.open("wb") as output:
                                shutil.copyfileobj(source, output)
                            path.chmod(((member.external_attr >> 16) & 0o777) or 0o644)
                if dest.exists():
                    raise ValueError(f"{dest} exists but lacks required source files")
                staging.rename(dest)
            ok(f"downloaded {name}")
        except (OSError, ValueError, zipfile.BadZipFile, urllib.error.URLError) as exc:
            err(f"could not download {name}: {exc}")
            complete = False
    return complete


TOOL_CHECKS = {
    "lspci": ("critical", "PCIe checks (package: pciutils)"),
    "setpci": ("critical", "link generation probe (package: pciutils)"),
    "nvidia-smi": ("critical", "GPU state / cmpunlocker verify"),
    "nvflash": ("optional", "VBIOS backup + serial (bundled in repo, or on PATH)"),
    "dmesg": ("optional", "kernel log scrape"),
    "journalctl": ("optional", "kernel log scrape fallback"),
    "vulkaninfo": ("optional", "Vulkan capability probe (package: vulkan-tools)"),
}
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):
    """Check root or working sudo for host setup and privileged inspection.

    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 - 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 boot_id():
    return (read_file("/proc/sys/kernel/random/boot_id") or "").strip()


def pending_state():
    try:
        return json.loads(PENDING_FILE.read_text())
    except (OSError, ValueError):
        return None


def mark_pending(stage, card=None):
    HARNESS_DIR.mkdir(exist_ok=True)
    PENDING_FILE.write_text(json.dumps({"stage": stage, "boot_id": boot_id(),
                                        "card": card}, indent=2))


def card_identity(bdf, device_id, serial=None):
    """Identify a card across boots without treating its GPU UUID as a serial."""
    query = run(["nvidia-smi", "-i", bdf, "--query-gpu=uuid",
                 "--format=csv,noheader"], timeout=30)
    uuid = query["out"].strip().splitlines()[0] if query["rc"] == 0 and query["out"].strip() else None
    if uuid and uuid.upper() in ("N/A", "UNKNOWN", "NONE"):
        uuid = None
    return {"serial": serial, "gpu_uuid": uuid, "device_id": device_id, "bdf": bdf}


def same_card(previous, current):
    """Return None when no stable identity is available for comparison."""
    if not isinstance(previous, dict):
        return None
    for key in ("serial", "gpu_uuid"):
        if previous.get(key) and current.get(key):
            return previous[key] == current[key]
    if previous.get("device_id") and previous["device_id"] != current["device_id"]:
        return False
    return None


def supported_versions():
    text = read_file(HARNESS_DIR / "cmpunlocker" / "driver" / "VERSION") or ""
    return [line.strip() for line in text.splitlines() if re.fullmatch(r"\d+\.\d+\.\d+", line.strip())]


def current_driver_version():
    r = run(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], timeout=30)
    if r["rc"] == 0:
        return r["out"].splitlines()[0].strip() if r["out"].strip() else None
    return None


def open_driver_active():
    version_text = read_file("/proc/driver/nvidia/version") or ""
    return "Open Kernel Module" in version_text


def preflight(auto_install=True, do_download=True):
    step("Preflight")
    report = {"tools": {}, "root": have_root()}
    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']}")
    if do_download:
        report["repos_ok"] = ensure_repos()
    versions = supported_versions()
    report["supported_driver_versions"] = versions
    if not INSTALL_SH.exists() or not VERIFY_SH.exists() or not versions:
        report["blocking_missing"] = ["cmpunlocker v0.3 source"]
        report["preflight_ok"] = False
        err("preflight blocked: cmpunlocker v0.3 source is missing or incomplete")
        return report
    pending = pending_state()
    if (pending and pending.get("stage") == "nvidia-driver-610-open"
            and pending.get("boot_id") == boot_id()):
        report["reboot_required"] = True
        warn(f"{pending['stage']} installed during this boot; cold shut down and power on, then rerun")
        return report
    driver_rebooted = bool(pending and pending.get("boot_id") != boot_id()
                           and pending.get("stage") == "nvidia-driver-610-open")
    version = current_driver_version()
    report["driver_version"] = version
    need_driver = not version or version not in versions or not open_driver_active()
    if driver_rebooted and need_driver:
        report["blocking_missing"] = ["nvidia-driver-610-open installed but inactive after reboot; inspect Secure Boot and driver logs"]
        report["preflight_ok"] = False
        err(report["blocking_missing"][0])
        return report
    if driver_rebooted:
        PENDING_FILE.unlink(missing_ok=True)
    pkgs = ["pciutils", "build-essential", "patch", "dkms", "curl",
            f"linux-headers-{run(['uname', '-r'])['out'].strip()}"]
    build_tools = ("lspci", "setpci", "gcc", "make", "patch", "curl")
    need_tools = not all(which(t) for t in build_tools)
    need_headers = not Path("/lib/modules", run(["uname", "-r"])["out"].strip(), "build").exists()
    if auto_install and (need_tools or need_headers or need_driver) and rel.get("ID") == "ubuntu":
        allowed, reason = sudo_precondition()
        report["sudo"] = {"ok": allowed, "why": reason}
        if allowed:
            update_log = HARNESS_DIR / "logs" / f"apt_update_{datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f')}.log"
            info(f"apt-get update log: {update_log}")
            update = run_logged_as_root(["apt-get", "update"], update_log, timeout=300)
            if update["rc"] != 0:
                err(f"apt-get update failed: {(update['err'] or update['out'])[-300:]}")
            else:
                to_install = pkgs + (["nvidia-driver-610-open"] if need_driver else [])
                info("installing: " + " ".join(to_install))
                install_log = HARNESS_DIR / "logs" / f"apt_install_{datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f')}.log"
                info(f"apt-get install log: {install_log}")
                install = run_logged_as_root(["apt-get", "install", "-y"] + to_install,
                                             install_log, timeout=1800)
                report["install"] = {"rc": install["rc"], "log": str(install_log),
                                     "tail": (install["out"] + install["err"])[-2000:]}
                if install["rc"] != 0:
                    err("package installation failed: " + report["install"]["tail"][-300:])
                elif need_driver:
                    mark_pending("nvidia-driver-610-open")
                    report["reboot_required"] = True
                    warn("driver installed; cold shut down and power on, then rerun")
        else:
            err(f"sudo unavailable: {reason}")

    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}
        info(f"{tool}: {path or 'missing'}")
    blocking = [t for t in build_tools if not which(t)]
    if not report.get("reboot_required") and (current_driver_version() not in versions or not open_driver_active()):
        blocking.append("supported active nvidia-open 610 driver")
    if need_headers and not Path("/lib/modules", run(["uname", "-r"])["out"].strip(), "build").exists():
        blocking.append("matching kernel headers")
    report["blocking_missing"] = blocking
    report["preflight_ok"] = not blocking and not report.get("reboot_required")
    if blocking and not report.get("reboot_required"):
        err("preflight blocked: " + ", ".join(blocking))
    elif report["preflight_ok"]:
        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, "aer_fatal_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()
        snap["aer_fatal_sum"] = sum(int(x) for x in re.findall(r"\d+", fatal))
    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):
    cmd = ["lspci", "-vvv", "-D", "-s", bdf]
    r = run(cmd, timeout=30)
    link = parse_link(r["out"]) if r["rc"] == 0 else {}
    if (not link.get("lnksta_width") or not link.get("lnksta_gen")) and os.geteuid() != 0:
        allowed, _ = sudo_precondition()
        if allowed:
            privileged = as_root(cmd, timeout=30)
            privileged_link = parse_link(privileged["out"]) if privileged["rc"] == 0 else {}
            fields = ("lnksta_width", "lnksta_gen")
            if sum(bool(privileged_link.get(field)) for field in fields) > sum(
                    bool(link.get(field)) for field in fields):
                r, link = privileged, privileged_link
    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, since):
    """Warn only for recent lines naming this endpoint; keep others separately."""
    since = since.replace("T", " ")
    text, source = "", None
    for command, label in (
            (["journalctl", "-k", "-b", "--since", since, "--no-pager", "-o", "short-precise"], "journalctl"),
            (["dmesg", "--ctime", "--since", since], "dmesg")):
        r = as_root(command, timeout=25)
        if r["rc"] == 0:
            text, source = r["out"], label
            break
    if source is None:
        return {"source": None, "related": [], "errors": [], "other_pcie": [],
                "note": "recent kernel log unavailable (root or journal access needed)"}
    tokens = (bdf.lower(), bdf[:-2].lower())
    related, errors, other = [], [], []
    for line in text.splitlines():
        low = line.lower()
        if any(token in low for token in tokens):
            related.append(line)
            if any(k in low for k in ("xid", "error", "aer:", "fail", "corrupted", "fallen off")):
                errors.append(line)
        elif any(k in low for k in ("aer:", "pcieport", "nvrm", "xid")):
            other.append(line)
    return {"source": source, "related": related[-200:], "errors": errors[-50:],
            "other_pcie": other[-200:], "note": f"{len(related)} card lines, {len(other)} other PCIe/GPU lines since {since}"}


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

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


def usable_serial(value):
    return bool(value and sanitize(value) == value and
                value.upper() not in ("N/A", "UNKNOWN", "NONE") and
                not re.fullmatch(r"0+", value))


def get_serial(nf, bdf):
    """Read the physical board serial; never substitute a GPU UUID or PCI address."""
    smi = run(["nvidia-smi", "-i", bdf, "--query-gpu=serial",
               "--format=csv,noheader"], timeout=30)
    if smi["rc"] == 0:
        serial = smi["out"].strip().splitlines()[0] if smi["out"].strip() else ""
        if usable_serial(serial):
            info(f"serial auto-read from nvidia-smi: {serial}")
            return serial, "nvidia-smi"

    detail = run(["nvidia-smi", "-q", "-i", bdf], timeout=30)
    if detail["rc"] == 0:
        match = re.search(r"^\s*Serial Number\s*:\s*(\S+)", detail["out"], re.I | re.M)
        if match and usable_serial(match.group(1)):
            info(f"serial auto-read from nvidia-smi details: {match.group(1)}")
            return match.group(1), "nvidia-smi-detail"

    vpd_cmd = ["lspci", "-vvv", "-D", "-s", bdf]
    for vpd in (run(vpd_cmd, timeout=30), as_root(vpd_cmd, timeout=30)):
        if vpd["rc"] == 0:
            match = re.search(r"\[SN\]\s+Serial number:\s*(\S+)", vpd["out"], re.I)
            if match and usable_serial(match.group(1)):
                info(f"serial auto-read from PCI VPD: {match.group(1)}")
                return match.group(1), "pci-vpd"

    if nf:
        q = as_root([nf, "--pcisegbus", bdf, "--query"], timeout=40)
        serials = {m.group(1) for m in re.finditer(
            r"Serial(?:\s*Number)?\s*[:=]\s*(\S+)", q["out"], re.I)
            if q["rc"] == 0 and usable_serial(m.group(1))}
        if q["rc"] == 0 and len(serials) == 1:
            serial = serials.pop()
            info(f"serial auto-read from card-scoped nvflash query: {serial}")
            return serial, "nvflash-pcisegbus-query"
    return None, "unavailable"


def normalize_bdf(value):
    match = re.search(r"([0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-7])$", value)
    return match.group(1).lower() if match else None


def memory_state(bdf, device_id):
    """Identify unlock from this card's reported memory, independent of Gen2 status."""
    r = run(["nvidia-smi", "--query-gpu=pci.bus_id,memory.total",
             "--format=csv,noheader,nounits"], timeout=30)
    if r["rc"] != 0:
        return {"status": "unavailable", "note": (r["err"] or r["out"])[-300:]}
    for line in r["out"].splitlines():
        fields = [field.strip() for field in line.split(",")]
        if len(fields) != 2 or normalize_bdf(fields[0]) != bdf:
            continue
        try:
            mib = int(fields[1])
        except ValueError:
            return {"status": "unavailable", "note": f"invalid memory total: {fields[1]}"}
        if device_id == "10de:20c2":
            status = "unlocked" if mib >= 60000 else "stock" if 7680 <= mib <= 8704 else "unexpected"
        elif device_id == "10de:2082":
            status = "unlocked" if 35000 <= mib < 60000 else "stock" if 9728 <= mib <= 10752 else "unexpected"
        else:
            status = "unsupported"
        return {"status": status, "memory_mib": mib}
    return {"status": "unavailable", "note": f"{bdf} absent from nvidia-smi"}


def ensure_unlocked(bdf, device_id, folder, card=None, allow_install=True):
    before = memory_state(bdf, device_id)
    result = {"before": before, "status": before["status"], "install_rc": None}
    pending = pending_state()
    if pending and pending.get("stage") == "cmpunlocker":
        if pending.get("boot_id") == boot_id():
            result.update(status="reboot-required", note="cold shut down and power on after cmpunlocker install")
            return result
        match = same_card(pending.get("card"), card or card_identity(bdf, device_id))
        if match is False:
            warn("previous unlock marker belongs to a different card; checking this card anew")
            PENDING_FILE.unlink(missing_ok=True)
        elif match is None and before["status"] != "unlocked":
            result.update(status="failed", note="cannot tie the previous unlock to this card; inspect the pending marker before retrying")
            return result
        elif before["status"] != "unlocked":
            result.update(status="failed", note="cmpunlocker installed but card is not unlocked after reboot")
            return result
        else:
            PENDING_FILE.unlink(missing_ok=True)
    if before["status"] == "unlocked":
        result.update(status="unlocked", note="unlocked memory is active")
        return result
    if before["status"] != "stock":
        result.update(status="failed", note="card memory/PCI ID is unavailable or unexpected; refusing to select an unlock profile")
        return result
    if not allow_install:
        result.update(status="failed", note="unlock inactive and --no-install was requested")
        return result
    if not INSTALL_SH.exists():
        result.update(status="failed", note="cmpunlocker install.sh is missing")
        return result
    allowed, reason = sudo_precondition()
    if not allowed:
        result.update(status="failed", note=f"root access required: {reason}")
        return result
    info("installing cmpunlocker for the detected card (this can take several minutes)")
    log_path = (folder / "extras" /
                f"cmpunlocker_install_{datetime.datetime.now().strftime('%Y%m%dT%H%M%S%f')}.log")
    info(f"cmpunlocker install log: {log_path}")
    install = run_logged_as_root(["bash", str(INSTALL_SH)], log_path, timeout=3600,
                                 cwd=INSTALL_SH.parent)
    result["install_rc"] = install["rc"]
    result["install_log"] = str(log_path)
    if install["rc"] != 0:
        result.update(status="failed", note=f"cmpunlocker installer failed (rc={install['rc']}); see {log_path}")
        return result
    mark_pending("cmpunlocker", card or card_identity(bdf, device_id))
    result.update(status="reboot-required", note="cmpunlocker installed; cold shut down and power on, then rerun")
    return result


def save_bios(nf, out_path, bdf):
    result = {"method": None, "path": None, "sha256": None, "size": 0, "ok": False, "note": ""}
    if out_path.exists():
        data = out_path.read_bytes()
        if len(data) > 512 and data.strip(b"\x00"):
            result.update(method="existing", path=str(out_path), size=len(data),
                          sha256=hashlib.sha256(data).hexdigest(), ok=True,
                          note="existing backup preserved")
        else:
            result["note"] = "existing backup is empty or invalid; move it before retrying"
        return result
    if nf:
        r = as_root([nf, "--pcisegbus", bdf, "--save", str(out_path)], timeout=90)
        if r["rc"] == 0 and out_path.exists() and out_path.stat().st_size > 512:
            data = out_path.read_bytes()
            if data.strip(b"\x00") == b"":
                result["note"] = "nvflash wrote an all-zero image"
            else:
                result.update(method="nvflash", path=str(out_path), size=len(data),
                              sha256=hashlib.sha256(data).hexdigest(), ok=True)
                return result
        else:
            result["note"] = f"card-scoped nvflash --save failed or image too small (rc={r['rc']}): {(r['err'] or r['out'])[-200:]}"
        out_path.unlink(missing_ok=True)
    else:
        result["note"] = "nvflash unavailable"

    rom = sysfs(bdf) / "rom"
    if as_root(["sh", "-c", f"echo 1 > {rom}"], timeout=15)["rc"] == 0:
        command = ["cat", str(rom)] if os.geteuid() == 0 else ["sudo", "-n", "cat", str(rom)]
        try:
            read = subprocess.run(command, capture_output=True, timeout=30)
            data = read.stdout if read.returncode == 0 else None
        except (OSError, subprocess.TimeoutExpired):
            data = None
        finally:
            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)
            result["note"] = (result["note"] + " | saved via sysfs rom").strip(" |")
            return result
        result["note"] += " | sysfs rom empty/all-zero"
    else:
        result["note"] += " | sysfs rom enable failed"
    return result


# ---------------------------------------------------------------- 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 run_verify_unlocker(bdf, device_id):
    if not VERIFY_SH.exists():
        return {"ran": False, "note": "harness/cmpunlocker source missing"}
    if not which("nvidia-smi"):
        return {"ran": False, "note": "nvidia-smi missing - unlock verify skipped"}
    kernel = run(["uname", "-r"])["out"].strip()
    inventory = Path("/lib/modules") / kernel / "updates" / "cmpunlocker" / "gpu_inventory"
    contents = read_file(inventory)
    if contents:
        entries = [line.split() for line in contents.splitlines()
                   if line.strip() and not line.lstrip().startswith("#")]
        if len(entries) != 1 or len(entries[0]) < 2 or normalize_bdf(entries[0][0]) != bdf or not entries[0][1].lower().endswith(device_id[-4:]):
            return {"ran": False, "status": "skipped-stale-inventory",
                    "note": f"installed gpu_inventory does not match {bdf} ({device_id}); live per-card memory and PCIe checks remain authoritative"}
    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:]}


# ---------------------------------------------------------------- optional Vulkan capability probe

def run_vulkan_probe():
    if not which("vulkaninfo"):
        return {"requested": True, "ran": False, "note": "vulkaninfo is missing"}
    r = run(["vulkaninfo", "--summary"], timeout=30)
    return {"requested": True, "ran": True, "rc": r["rc"],
            "note": "generic Vulkan capability listing; does not validate CMP VRAM",
            "log_tail": (r["out"] + r["err"]).splitlines()[-40:]}