#!/usr/bin/env python3
"""CMP 170HX batch card checker - guided one-card-at-a-time workflow (23 cards).

Per card: detect GPU -> auto serial -> VBIOS backup -> PCIe link/AER/config checks
-> retrain + rescan tests -> dmesg scrape -> cmpunlocker verify.sh (read-only)
-> optional 5-min Vulkan VRAM stress (last step, ask yes/no) -> cards/<serial>/report.json.

Scoring (lenient): FAIL only on no-device/link-down; WARN on correctable AER/Xid/
width< x16/verify issues.
"""
import argparse
import json
import sys
import time
from pathlib import Path

import cxlib as X
from cxlib import (C, ask, detect_gpus, err, info, now_iso, ok, sanitize, step, warn)

TOTAL_DEFAULT = X.DEFAULT_TOTAL


# ---------------------------------------------------------------- scoring

def score(rep):
    fails, warns = [], []
    if not rep.get("bdf"):
        fails.append("no CMP 170HX device detected on the bus")
    link = rep.get("link", {})
    if rep.get("bdf"):
        if link.get("lspci_rc") not in (0, None):
            fails.append("lspci could not read the device")
        if link.get("lnksta_width") in ("x0", "x1", None) or not link.get("lnksta_speed"):
            fails.append(f"link down or untrained: {link.get('lnksta_raw') or 'no LnkSta'}")
        width = link.get("lnksta_width")
        if width and width != "x16":
            warns.append(f"link width {width} (capacitor mod expects x16 - mod likely "
                         f"missing/incomplete on this card)")
        if link.get("lnksta_gen") == 1:
            warns.append("link negotiated Gen1 only (cmpunlocker Gen2 patch/service not active?)")
    if rep.get("rescan", {}).get("ok") is False:
        fails.append("device lost after bus rescan")
    if rep.get("retrain", {}).get("ok") is False:
        fails.append("link did not survive retrain")
    aer = rep.get("aer_after", {})
    aer_before = rep.get("aer_before", {})
    if (aer.get("ce_active") or aer.get("parent_ce_active")
            or (aer.get("aer_correctable_sum", 0) > aer_before.get("aer_correctable_sum", 0))):
        warns.append("correctable AER errors present/incremented")
    if aer.get("ue_active"):
        fails.append(f"uncorrectable AER errors active: {aer['ue_active']}")
    dm = rep.get("dmesg", {})
    if dm.get("errors"):
        warns.append(f"{len(dm['errors'])} kernel error line(s) related to this card (see dmesg.log)")
    vu = rep.get("vulkan_test", {})
    if vu.get("requested") and vu.get("ran"):
        if vu.get("tool_errors"):
            warns.append("errors reported during Vulkan VRAM stress")
        if dm.get("errors") and vu.get("ran"):
            warns.append("kernel errors around VRAM stress window")
    if not rep.get("bios", {}).get("ok"):
        warns.append("VBIOS backup failed/incomplete")
    ver = rep.get("cmpunlocker_verify", {})
    if ver.get("ran") and ver.get("status") not in ("unlocked-ok",):
        warns.append(f"cmpunlocker verify status: {ver.get('status')}")

    result = "FAIL" if fails else ("WARN" if warns else "PASS")
    return result, fails, warns


# ---------------------------------------------------------------- one card

def find_gpu(timeout=45):
    """Wait for a CMP 170HX to appear, offering rescan retries."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        bdfs = detect_gpus()
        if len(bdfs) == 1:
            return bdfs[0]
        if len(bdfs) > 1:
            warn(f"multiple CMP 170HX devices present: {', '.join(bdfs)}")
            return bdfs[0]
        time.sleep(2)
    warn("no CMP 170HX detected yet")
    if ask("trigger PCI bus rescan and retry? (y/N)", "n").lower().startswith("y"):
        if X.rescan_bus():
            ok("bus rescanned")
        else:
            err("rescan not permitted - reseat the card and press Enter")
            input()
        return find_gpu(timeout=30)
    return None


def process_card(index, total, args):
    header = f"Card {index}/{total}"
    step(header)
    bdf = find_gpu()
    if not bdf:
        err("no device - skipping this card")
        return None
    devid = ""
    for line in X.run(["lspci", "-D", "-s", bdf, "-nn"])["out"].splitlines():
        if "[10de:" in line:
            devid = line[line.index("[10de:") + 1:line.index("[10de:") + 10].rstrip("]")
    ok(f"device found at {bdf} ({devid})")

    nf = X.nvflash_bin(args)
    serial, serial_src = X.get_serial(nf, args)
    folder = X.CARDS_DIR / sanitize(serial)
    folder.mkdir(parents=True, exist_ok=True)
    info(f"folder: {folder}")

    rep = {
        "card_index": index, "total": total, "serial": serial, "serial_source": serial_src,
        "started_at": now_iso(), "bdf": bdf, "device_id": devid, "host": X.run(["uname", "-n"])["out"].strip(),
        "kernel": X.run(["uname", "-r"])["out"].strip(),
    }

    step(f"[{index}/{total}] VBIOS backup")
    rep["bios"] = X.save_bios(nf, folder / "bios.rom", args)
    (ok if rep["bios"]["ok"] else warn)(f"bios: {rep['bios']['method'] or 'FAILED'} "
                                        f"{rep['bios']['size']} bytes - {rep['bios']['note']}")

    step(f"[{index}/{total}] PCIe link + config")
    link, lspci_txt = X.check_link(bdf)
    rep["link"] = link
    (folder / "lspci.txt").write_text(lspci_txt)
    (folder / "config_space.hex").write_text(X.dump_config(bdf))
    info(f"LnkSta: {link.get('lnksta_raw') or 'n/a'} (setpci gen={link.get('setpci_current_gen')})")
    rep["aer_before"] = X.aer_snapshot(bdf)

    step(f"[{index}/{total}] retrain + rescan tests")
    rep["retrain"] = X.retrain_test(bdf)
    info(f"retrain: ok={rep['retrain']['ok']} - {rep['retrain']['note']}")
    rep["rescan"] = X.rescan_test(bdf)
    info(f"rescan: ok={rep['rescan']['ok']} - {rep['rescan']['note']}")
    rep["aer_after"] = X.aer_snapshot(bdf)
    link2, lspci2 = X.check_link(bdf)
    rep["link_after"] = link2
    if lspci2:
        (folder / "lspci_after.txt").write_text(lspci2)

    step(f"[{index}/{total}] kernel log scrape")
    dm = X.dmesg_capture(bdf)
    rep["dmesg"] = dm
    (folder / "dmesg.log").write_text("\n".join(dm.get("related", [])))
    info(f"{dm.get('note')} (source: {dm.get('source')})")
    if dm.get("errors"):
        for e in dm["errors"][-5:]:
            warn(e.strip()[:160])

    step(f"[{index}/{total}] nvidia-smi + cmpunlocker verify")
    rep["nvidia_smi"] = X.nvidia_smi_info()
    rep["cmpunlocker_verify"] = X.run_verify_unlocker()
    v = rep["cmpunlocker_verify"]
    (folder / "verify_unlocker.log").write_text(
        "\n".join(v.get("tail", [])) or (v.get("note") or "not run"))
    info(f"verify: {'ran rc=' + str(v.get('rc')) + ' -> ' + str(v.get('status')) if v.get('ran') else 'skipped - ' + str(v.get('note'))}")

    step(f"[{index}/{total}] optional Vulkan VRAM test (last step)")
    if args.vulkan == "never":
        rep["vulkan_test"] = {"requested": False, "ran": False, "note": "disabled via --vulkan never"}
    elif args.vulkan == "always":
        rep["vulkan_test"] = X.run_vulkan_test()
    else:
        if ask("Run 5-min Vulkan VRAM stress test now? (y/N)", "n").lower().startswith("y"):
            rep["vulkan_test"] = X.run_vulkan_test()
        else:
            rep["vulkan_test"] = {"requested": False, "ran": False, "note": "operator declined"}
            info("vulkan test skipped")

    result, fails, warns = score(rep)
    rep.update(finished_at=now_iso(), result=result, failures=fails, warnings=warns)
    (folder / "report.json").write_text(json.dumps(rep, indent=2))

    print()
    if result == "PASS":
        ok(f"{serial}: PASS")
    elif result == "WARN":
        warn(f"{serial}: WARN")
    else:
        err(f"{serial}: FAIL")
    for f in fails:
        err("  FAIL: " + f)
    for w in warns:
        warn("  WARN: " + w)
    info(f"report: {folder / 'report.json'}")
    return rep


# ---------------------------------------------------------------- progress

def load_progress():
    try:
        return json.loads(X.PROGRESS_FILE.read_text())
    except Exception:
        return {"done": []}


def save_progress(prog):
    X.PROGRESS_FILE.write_text(json.dumps(prog, indent=2))


def scoreboard(results):
    step("Scoreboard")
    counts = {"PASS": 0, "WARN": 0, "FAIL": 0}
    for r in results:
        if not r:
            continue
        counts[r["result"]] = counts.get(r["result"], 0) + 1
        color = {"PASS": C.GREEN, "WARN": C.YELLOW, "FAIL": C.RED}[r["result"]]
        print(f"  {color}{r['result']:<4}{C.RESET} #{r['card_index']:>2} "
              f"{r['serial']:<24} {r.get('bdf') or '-':<12} "
              f"{r.get('link', {}).get('lnksta_width') or '?':<4} "
              f"{r.get('link', {}).get('lnksta_speed') or '?'}")
    skipped = [r for r in results if r is None]
    print(f"\n  PASS={counts.get('PASS', 0)}  WARN={counts.get('WARN', 0)}  "
          f"FAIL={counts.get('FAIL', 0)}  skipped={len(skipped)}  "
          f"total_seen={len(results)}/{TOTAL_DEFAULT}")


# ---------------------------------------------------------------- demo

def demo(args):
    """Validate report pipeline without hardware: writes cards/DEMO-*."""
    step("Demo mode (no hardware)")
    args.vulkan = "never"
    args.serial = "DEMO-SERIAL-001"
    args.skip_bios = True
    real_find = globals()["find_gpu"]
    globals()["find_gpu"] = lambda timeout=0: "0000:0a:00.0"

    class FakeRun:
        pass

    orig_detect, orig_check = X.detect_gpus, X.check_link
    orig_aer, orig_dm, orig_cfg = X.aer_snapshot, X.dmesg_capture, X.dump_config
    orig_retrain, orig_rescan = X.retrain_test, X.rescan_test
    orig_smi, orig_verify = X.nvidia_smi_info, X.run_verify_unlocker
    X.detect_gpus = lambda: ["0000:0a:00.0"]
    X.check_link = lambda b: ({"lnksta_speed": "5GT/s", "lnksta_width": "x4", "lnksta_gen": 2,
                               "lnkcap_speed": "8GT/s", "lnkcap_width": "x16",
                               "lnksta_raw": "LnkSta: Speed 5GT/s, Width x4, TrErr-",
                               "setpci_current_gen": 2, "lspci_rc": 0}, "demo lspci output\n")
    X.aer_snapshot = lambda b: {"aer_sysfs_correctable": None, "aer_sysfs_fatal": None,
                                "aer_correctable_sum": 0, "ce_active": [], "ue_active": [],
                                "parent_bdf": None, "parent_ce_active": []}
    X.dmesg_capture = lambda b: {"source": "demo", "related": ["demo line"], "errors": [],
                                 "note": "demo"}
    X.dump_config = lambda b: "demo config space dump\n"
    X.retrain_test = lambda b: {"attempted": True, "ok": True, "note": "x4 -> x4",
                                "width_before": "x4", "width_after": "x4"}
    X.rescan_test = lambda b: {"attempted": True, "ok": True, "note": "re-enumerated"}
    X.nvidia_smi_info = lambda: {"available": False}
    X.run_verify_unlocker = lambda: {"ran": False, "note": "demo"}
    try:
        rep = process_card(1, TOTAL_DEFAULT, args)
    finally:
        globals()["find_gpu"] = real_find
        X.detect_gpus, X.check_link = orig_detect, orig_check
        X.aer_snapshot, X.dmesg_capture, X.dump_config = orig_aer, orig_dm, orig_cfg
        X.retrain_test, X.rescan_test = orig_retrain, orig_rescan
        X.nvidia_smi_info, X.run_verify_unlocker = orig_smi, orig_verify
    return rep


# ---------------------------------------------------------------- main

def main():
    global TOTAL_DEFAULT
    ap = argparse.ArgumentParser(description="CMP 170HX guided 23-card checker")
    ap.add_argument("--preflight", action="store_true", help="only run tool check/install/clone")
    ap.add_argument("--demo", action="store_true", help="run pipeline with fake hardware (self-test)")
    ap.add_argument("--cards", type=int, default=TOTAL_DEFAULT, help="total cards (default 23)")
    ap.add_argument("--start", type=int, default=1, help="first card index")
    ap.add_argument("--nvflash", help="path to nvflash binary")
    ap.add_argument("--serial", help="override serial for the next card (single-card run)")
    ap.add_argument("--skip-bios", action="store_true", help="skip VBIOS backup")
    ap.add_argument("--vulkan", choices=("ask", "always", "never"), default="ask",
                    help="5-min Vulkan VRAM test policy (default ask, runs as last step)")
    ap.add_argument("--no-install", action="store_true", help="do not auto-install missing packages")
    ap.add_argument("--no-clone", action="store_true", help="do not clone harness repos")
    ap.add_argument("--restart", action="store_true", help="ignore previous progress.json")
    ap.add_argument("--no-pause", action="store_true", help="do not wait for card insertion")
    args = ap.parse_args()

    TOTAL_DEFAULT = args.cards

    pf = X.preflight(auto_install=not args.no_install, do_clone=not args.no_clone)
    if args.preflight:
        return 0
    if args.demo:
        demo(args)
        return 0

    X.CARDS_DIR.mkdir(exist_ok=True)
    prog = {"done": []} if args.restart else load_progress()

    results = []
    for i in range(args.start, args.cards + 1):
        if str(i) in prog.get("done", []):
            info(f"card {i} already done (progress.json) - skipping")
            results.append(None)
            continue
        if not args.no_pause:
            ans = ask(f"Insert card {i}/{args.cards}, then press Enter "
                      f"(s=skip this card, q=quit)", "")
            if ans.lower() == "q":
                break
            if ans.lower() == "s":
                results.append(None)
                prog.setdefault("done", []).append(str(i))
                prog.setdefault("skipped", []).append(str(i))
                save_progress(prog)
                continue
        rep = process_card(i, args.cards, args)
        results.append(rep)
        if rep:
            prog.setdefault("done", []).append(str(i))
            save_progress(prog)
        else:
            if not ask("retry this card? (y/N)", "n").lower().startswith("y"):
                break

    if results:
        scoreboard(results)
    done = len([r for r in results if r])
    info(f"{done} card(s) processed. Output in {X.CARDS_DIR}")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print()
        warn("interrupted - progress saved in progress.json")
        sys.exit(130)