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

SCRIPT_DIR = Path(__file__).resolve().parent
HARNESS_DIR = SCRIPT_DIR / "harness"
CARDS_DIR = SCRIPT_DIR / "cards"
PROGRESS_FILE = SCRIPT_DIR / "progress.json"

REPOS = {
    "170th-Street": "https://github.com/amoghmunikote/170th-Street.git",
    "cmpunlocker": "https://github.com/amoghmunikote/cmpunlocker.git",
}
VERIFY_SH = HARNESS_DIR / "cmpunlocker" / "verify.sh"
GPU_IDS = ("20c2", "2082")
DEFAULT_TOTAL = 23
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):
    try:
        p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env)
        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):
    if os.geteuid() == 0:
        return run(cmd, timeout=timeout)
    if shutil.which("sudo"):
        r = run(["sudo", "-n"] + cmd, timeout=timeout)
        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 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 (not packaged: manual install)"),
    "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 preflight(auto_install=True, do_clone=True):
    step("Preflight")
    report = {"tools": {}, "root": have_root()}
    critical_missing, optional_missing = [], []

    for tool, (level, why) in TOOL_CHECKS.items():
        path = 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)

    if auto_install:
        pkgs = sorted({PKG_FOR[t] for t in critical_missing + optional_missing if t in PKG_FOR})
        if pkgs:
            if detect_pm():
                info(f"attempting to install: {' '.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)}")
            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))

    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 "nvflash" in critical_missing:
        warn("nvflash missing - VBIOS backup will fall back to sysfs ROM, serial to manual entry")
        print("     Download Linux x64 nvflash from https://www.nvidia.com/en-us/drivers/nvflash/")
        print("     and pass --nvflash /path/to/nvflash, or place it on PATH.")
    if "nvidia-smi" in critical_missing:
        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")

    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):
    return args.nvflash if getattr(args, "nvflash", None) else which("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")
    serial = ""
    while not serial:
        serial = ask("Enter serial number from the card label (folder name)")
        if not serial:
            warn("serial is required for the folder name")
    return serial, "manual-label"


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
    if nf:
        r = as_root([nf, "--save", str(out_path)], timeout=90)
        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"":
                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"nvflash --save failed rc={r['rc']}: {(r['err'] or r['out'])[-200:]}"
    else:
        result["note"] = "nvflash unavailable"

    for bdf in detect_gpus():
        rom = sysfs(bdf) / "rom"
        if as_root(["sh", "-c", f"echo 1 > {rom}"], timeout=15)["rc"] != 0:
            result["note"] += " | sysfs rom enable failed"
            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)
            result["note"] = (result["note"] + " | saved via sysfs rom").strip(" |")
            return result
        result["note"] += " | sysfs rom empty/all-zero"
    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():
    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:]}


# ---------------------------------------------------------------- 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