#!/usr/bin/env python3
"""Ivy Bridge (Gen7) driver test runner.

Reads manifest.toml, runs the tests against the driver trees the nix shell
points at, and diffs the result against baseline.json.
"""

import argparse
import json
import os
import re
import shlex
import shutil
import subprocess
import sys
import tomllib
from datetime import datetime
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent
LAB = ROOT.parent
MANIFEST = ROOT / "manifest.toml"
BASELINE = ROOT / "baseline.json"
RESULT = ROOT / "result"

C_OK, C_BAD, C_DIM, C_OFF = "\033[32m", "\033[31m", "\033[90m", "\033[0m"
if not sys.stdout.isatty():
    C_OK = C_BAD = C_DIM = C_OFF = ""


# ---------------------------------------------------------------- manifest

def expand(v: Any) -> Any:
    if isinstance(v, str):
        return os.path.expandvars(os.path.expanduser(v))
    if isinstance(v, dict):
        return {k: expand(x) for k, x in v.items()}
    if isinstance(v, list):
        return [expand(x) for x in v]
    return v


def load_manifest():
    with open(MANIFEST, "rb") as f:
        return expand(tomllib.load(f))


def select(man, only=None, group=None):
    tests = man["test"]
    if only:
        tests = [t for t in tests if t["name"] in only]
    if group:
        tests = [t for t in tests if t.get("group") in group]
    return tests


# ----------------------------------------------------------------- process

def run(cmd, outdir, tag, env=None, timeout=3600, cwd=None):
    full_env = dict(os.environ)
    if env:
        full_env.update({k: str(v) for k, v in env.items()})
    outdir.mkdir(parents=True, exist_ok=True)
    try:
        p = subprocess.run(
            cmd if isinstance(cmd, list) else shlex.split(cmd),
            capture_output=True, text=True, env=full_env, timeout=timeout,
            cwd=cwd or None,
        )
        text, rc = p.stdout + p.stderr, p.returncode
    except subprocess.TimeoutExpired:
        text, rc = "TIMEOUT\n", 124
    except FileNotFoundError as e:
        text, rc = f"NOT FOUND: {e}\n", 127
    (outdir / f"{tag}.log").write_text(text)
    return rc, text


def driver_env(cfg):
    icd = cfg.get("vk_driver_files", "")
    if not icd:
        return {}
    return {"VK_DRIVER_FILES": icd, "VK_ICD_FILENAMES": icd}


# ------------------------------------------------------------------ checks

def probe_env(cfg):
    """Everything the runner depends on, and whether it is actually there."""
    icd_dir = Path(cfg.get("icd_vendors", ""))
    vendors = sorted(p.stem for p in icd_dir.glob("*.icd")) if icd_dir.is_dir() else []
    cts = Path(cfg.get("cts_dir", ""))
    return [
        ("mesa build", cfg.get("mesa_build", ""), Path(cfg.get("mesa_build", "/x")).is_dir()),
        ("hasvk ICD", cfg.get("vk_driver_files", ""), Path(cfg.get("vk_driver_files", "/x")).is_file()),
        ("clvk build", cfg.get("clvk_build", ""), Path(cfg.get("clvk_build", "/x")).is_dir()),
        ("cts dir", cfg.get("cts_dir", ""), cts.is_dir()),
        ("deqp-vk", str(cts / "deqp-vk"), (cts / "deqp-vk").is_file()),
        ("icd vendors", f"{icd_dir}  [{' '.join(vendors) or 'empty'}]", bool(vendors)),
        ("ffmpeg", shutil.which("ffmpeg") or "-", bool(shutil.which("ffmpeg"))),
    ]


def check_config(cfg):
    """The Vulkan loader silently falls back to the system ICD when the file
    it is pointed at is missing, which quietly tests the wrong driver."""
    icd = cfg.get("vk_driver_files", "")
    if icd and not Path(icd).is_file():
        sys.exit(f"vk_driver_files missing: {icd}\n"
                 f"the loader would fall back to the system ICD — refusing to run")
    cts = cfg.get("cts_dir", "")
    if cts and not Path(cts).is_dir():
        sys.exit(f"cts_dir missing: {cts}")


# ------------------------------------------------------------------ assets

def make_assets(cfg, outdir):
    w, h, n = cfg["asset_width"], cfg["asset_height"], cfg["asset_frames"]
    assets = outdir / "assets"
    assets.mkdir(parents=True, exist_ok=True)
    ref, src = assets / "ref.y4m", assets / "src.h264"

    if not ref.exists():
        rc, _ = run(f"ffmpeg -y -f lavfi -i testsrc2=size={w}x{h}:rate=30 "
                    f"-frames:v {n} -pix_fmt yuv420p {ref}", assets, "gen_ref")
        if rc != 0 or not ref.exists():
            sys.exit(f"could not generate ref.y4m — {assets}/gen_ref.log")

    if not src.exists():
        # B-frames on purpose: the decoder has to handle reordering
        rc, _ = run(f"ffmpeg -y -i {ref} -c:v libx264 -preset medium "
                    f"-bf 3 -b_strategy 2 -refs 4 -g 30 -crf 20 {src}",
                    assets, "gen_src")
        if rc != 0 or not src.exists():
            sys.exit(f"could not generate src.h264 — {assets}/gen_src.log")

    return ref, src


# ----------------------------------------------------------------- runners

def run_cts(t, cfg, outdir):
    qpa = outdir / "TestResults.qpa"
    cmd = cfg["cts_cmd"].format(filter=t["filter"], log=qpa)
    rc, text = run(cmd, outdir, "cts", env=driver_env(cfg), cwd=cfg.get("cts_dir"),
                   timeout=t.get("timeout", cfg.get("cts_timeout", 3600)))

    blob = text + (qpa.read_text(errors="ignore") if qpa.exists() else "")
    count = lambda w: len(re.findall(rf'StatusCode="{w}"', blob))

    passed = count("Pass")
    failed = count("Fail") + count("QualityWarning") + count("InternalError")
    skipped = count("NotSupported")

    if passed == 0 and failed == 0:  # some wrappers only print the summary
        m = re.search(r"Passed:\s*(\d+).*?Failed:\s*(\d+).*?Not supported:\s*(\d+)",
                      blob, re.S)
        if m:
            passed, failed, skipped = (int(x) for x in m.groups())

    aborted = "ABORTED" in blob
    return {"pass": passed, "fail": failed, "skip": skipped,
            "aborted": aborted, "rc": rc}


def last_error(text):
    """The line worth putting in the summary when ffmpeg gives up."""
    for line in reversed(text.strip().splitlines()):
        if re.search(r"assert|error|failed|invalid|not implemented|no such",
                     line, re.I) and "Task finished" not in line:
            return line.strip()[:120]
    return ""


def parse_psnr(text):
    m = re.search(r"average:([\d.]+)", text) or \
        re.search(r"PSNR.*?average:\s*([\d.]+)", text)
    return float(m.group(1)) if m else None


def run_video(t, cfg, outdir, assets):
    ref, src = assets
    env = driver_env(cfg)
    op = t["op"]

    # h264_vulkan takes Vulkan frames, so the software frames have to be
    # uploaded explicitly — without this ffmpeg dies in the filter graph and
    # the encoder is never reached.
    dev = cfg.get("vulkan_device", "vk:0")
    hw = f"-init_hw_device vulkan={dev} -filter_hw_device vk"
    upload = "-vf format=nv12,hwupload"

    if op == "decode":
        rc, text = run(f"ffmpeg -v info -hwaccel vulkan -i {src} -f null -",
                       outdir, "decode", env=env)
        hits = re.findall(r"frame=\s*(\d+)", text)
        return {"frames": max((int(x) for x in hits), default=0),
                "errors": len(re.findall(r"\[h264.*?\] (error|Error|corrupt)", text)),
                "rc": rc}

    if op == "encode":
        out = outdir / "enc.h264"
        rc, text = run(f"ffmpeg -y {hw} -i {ref} {upload} "
                       f"-c:v h264_vulkan -qp 18 -g 1 {out}",
                       outdir, "encode", env=env)
        if rc != 0 or not out.exists() or out.stat().st_size == 0:
            return {"psnr": None, "min_psnr": t.get("min_psnr"), "rc": rc,
                    "err": last_error(text)}
        _, text = run(f"ffmpeg -i {out} -i {ref} -lavfi psnr -f null -", outdir, "psnr")
        return {"psnr": parse_psnr(text), "min_psnr": t.get("min_psnr"), "rc": rc}

    if op == "roundtrip":
        enc, dec = outdir / "rt.h264", outdir / "rt.y4m"
        rc, text = run(f"ffmpeg -y {hw} -i {ref} {upload} "
                       f"-c:v h264_vulkan -qp 18 {enc}",
                       outdir, "rt_encode", env=env)
        if rc != 0 or not enc.exists() or enc.stat().st_size == 0:
            return {"psnr": None, "min_psnr": t.get("min_psnr"), "rc": rc,
                    "err": last_error(text)}
        rc, text = run(f"ffmpeg -y -hwaccel vulkan -i {enc} {dec}",
                       outdir, "rt_decode", env=env)
        if rc != 0 or not dec.exists():
            return {"psnr": None, "min_psnr": t.get("min_psnr"), "rc": rc,
                    "err": last_error(text)}
        _, text = run(f"ffmpeg -i {dec} -i {ref} -lavfi psnr -f null -", outdir, "rt_psnr")
        return {"psnr": parse_psnr(text), "min_psnr": t.get("min_psnr"), "rc": 0}

    return {"rc": 1, "err": f"unknown video op: {op}"}


CL_PROBE = r'''
import json, sys
import numpy as np, pyopencl as cl

out = []
for p in cl.get_platforms():
    for d in p.get_devices():
        r = {"platform": p.name.strip(), "device": d.name.strip(), "ok": False}
        try:
            ctx = cl.Context([d])
            q = cl.CommandQueue(ctx)
            prg = cl.Program(ctx, open(sys.argv[1]).read()).build()
            a = np.arange(4096, dtype=np.float32)
            b = a * 2.0
            mf = cl.mem_flags
            ba = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=a)
            bb = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=b)
            bc = cl.Buffer(ctx, mf.WRITE_ONLY, a.nbytes)
            prg.vec_add(q, a.shape, None, ba, bb, bc)
            got = np.empty_like(a)
            cl.enqueue_copy(q, got, bc)
            q.finish()
            r["ok"] = bool(np.allclose(got, a + b))
            if not r["ok"]:
                r["err"] = "result mismatch"
        except Exception as e:
            r["err"] = f"{type(e).__name__}: {e}"[:200]
        out.append(r)
print("__CLPROBE__" + json.dumps(out))
'''

CL_KERNEL = """__kernel void vec_add(__global const float *a, __global const float *b,
                      __global float *c) {
    int i = get_global_id(0);
    c[i] = a[i] + b[i];
}
"""


def run_opencl(t, cfg, outdir):
    """Each backend runs on its own, with a vendor dir holding only its ICD,
    so a backend that fails to load cannot be masked by another one."""
    outdir.mkdir(parents=True, exist_ok=True)
    kernel = outdir / "vec_add.cl"
    kernel.write_text(CL_KERNEL)
    probe = outdir / "probe.py"
    probe.write_text(CL_PROBE)

    icd_dir = Path(cfg.get("icd_vendors", ""))
    backends = t.get("backends", ["rusticl", "clvk", "pocl"])
    results = {}

    for name in backends:
        icd = icd_dir / f"{name}.icd"
        if not icd.exists():
            results[name] = {"ok": False, "err": f"no {name}.icd in {icd_dir}"}
            continue
        vend = outdir / "vendors" / name
        vend.mkdir(parents=True, exist_ok=True)
        target = vend / f"{name}.icd"
        if target.is_symlink() or target.exists():
            target.unlink()
        target.symlink_to(icd.resolve())

        env = dict(driver_env(cfg))
        env.update({"OCL_ICD_VENDORS": str(vend),
                    "OCL_ICD_ASSUME_ICD_EXTENSION": "1",
                    "PYOPENCL_COMPILER_OUTPUT": "1"})
        if name == "rusticl":
            env["RUSTICL_ENABLE"] = cfg.get("rusticl_enable", "crocus")

        rc, text = run([sys.executable, str(probe), str(kernel)], outdir,
                       f"cl_{name}", env=env, timeout=t.get("timeout", 600))
        m = re.search(r"__CLPROBE__(.*)", text)
        if rc != 0 or not m:
            results[name] = {"ok": False, "err": f"probe failed (rc={rc})"}
            continue
        devs = json.loads(m.group(1))
        if not devs:
            results[name] = {"ok": False, "err": "no device"}
            continue
        d = devs[0]
        results[name] = {"ok": d["ok"], "device": d["device"],
                         **({"err": d["err"]} if d.get("err") else {})}

    return {"backends": results,
            "ok_count": sum(1 for r in results.values() if r["ok"]),
            "total": len(backends),
            "rc": 0 if all(r["ok"] for r in results.values()) else 1}


def run_cmd(t, cfg, outdir):
    """Generic escape hatch: build (optionally) and run a probe from the lab
    tree.  Probes are compiled per run — checked-in binaries go stale as soon
    as the nix store is garbage collected."""
    outdir.mkdir(parents=True, exist_ok=True)
    cwd = t.get("cwd")
    cwd = str(LAB / cwd) if cwd and not os.path.isabs(cwd) else cwd
    fmt = lambda s: s.format(lab=LAB, out=outdir)

    if t.get("build"):
        rc, _ = run(fmt(t["build"]), outdir, "build", cwd=cwd, timeout=600)
        if rc != 0:
            return {"ok": False, "rc": rc, "err": "build failed"}

    rc, text = run(fmt(t["cmd"]), outdir, "cmd", env=driver_env(cfg),
                   cwd=cwd, timeout=t.get("timeout", 600))
    ok = rc == t.get("expect_rc", 0)
    if ok and t.get("match"):
        ok = bool(re.search(t["match"], text, re.S))
    return {"ok": ok, "rc": rc}


KINDS = {"cts": run_cts, "video": run_video, "opencl": run_opencl, "cmd": run_cmd}


# -------------------------------------------------------------- formatting

def fmt_one(r):
    if "pass" in r:
        s = f"{r['pass']}/{r['fail']}"
        return s + (f" {C_BAD}aborted{C_OFF}" if r.get("aborted") else "")
    if "frames" in r:
        return f"{r['frames']} frames, {r['errors']} errors"
    if "psnr" in r:
        if not r["psnr"]:
            why = f" — {r['err']}" if r.get("err") else ""
            return f"{C_BAD}failed{C_OFF}{C_DIM}{why}{C_OFF}"
        lo = r.get("min_psnr")
        flag = f" {C_BAD}<{lo}{C_OFF}" if lo and r["psnr"] < lo else ""
        return f"psnr {r['psnr']}{flag}"
    if "ok_count" in r:
        up = [n for n, b in r["backends"].items() if b["ok"]]
        down = [n for n, b in r["backends"].items() if not b["ok"]]
        s = "+".join(up) if up else "none"
        return f"{s} ok" + (f", {C_BAD}{'+'.join(down)} down{C_OFF}" if down else "")
    if "both_up" in r:  # pre-0.2 baselines
        return "both up" if r["both_up"] else f"{C_BAD}missing{C_OFF}"
    if "ok" in r:
        return "ok" if r["ok"] else f"{C_BAD}fail (rc={r.get('rc')}){C_OFF}"
    return "?"


def score(r):
    if "pass" in r:
        return r["pass"] - r["fail"] * 1000
    if "frames" in r:
        return r["frames"] - r["errors"] * 1000
    if "psnr" in r:
        return r["psnr"] or -1
    if "ok_count" in r:
        return r["ok_count"]
    if "both_up" in r:
        return 2 if r["both_up"] else 0
    if "ok" in r:
        return 1 if r["ok"] else 0
    return -1


def comparable(a, b):
    """Result shapes must match, otherwise the scores mean different things."""
    keys = lambda r: next((k for k in ("pass", "frames", "psnr", "ok_count",
                                       "both_up", "ok") if k in r), None)
    ka, kb = keys(a), keys(b)
    return ka == kb or {ka, kb} == {"ok_count", "both_up"}


def below_floor(r):
    lo = r.get("min_psnr")
    return bool(lo) and (r.get("psnr") or 0) < lo


def has_regression(base, cur):
    for n, v in cur.items():
        if below_floor(v):
            return True
        old = base.get(n)
        if old and comparable(old, v) and score(v) < score(old):
            return True
    return False


def row(name, r, old):
    val = fmt_one(r)
    if old is None or not comparable(old, r):
        return f"  {C_DIM}new{C_OFF} {name:<28} {val}"
    cur_s, old_s = score(r), score(old)
    if cur_s > old_s:
        return f"  {C_OK}^{C_OFF} {name:<28} {val} (was {fmt_one(old)})"
    if cur_s < old_s:
        return f"  {C_BAD}v{C_OFF} {name:<28} {val} (was {fmt_one(old)})"
    return f"  = {name:<28} {val}"


def diff_report(base, cur, group_by="group"):
    if group_by != "group":
        return "\n".join(row(n, r, base.get(n)) for n, r in cur.items())
    groups = {}
    for n, r in cur.items():
        groups.setdefault(r.get("group", "other"), []).append((n, r))
    lines = []
    for g in sorted(groups):
        lines.append(f"\n{C_DIM}[{g}]{C_OFF}")
        lines += [row(n, r, base.get(n)) for n, r in groups[g]]
    return "\n".join(lines)


def strip_color(s):
    return re.sub(r"\033\[[0-9;]*m", "", s)


# ---------------------------------------------------------------- commands

def new_result_dir():
    d = RESULT / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    d.mkdir(parents=True, exist_ok=True)
    latest = RESULT / "latest"
    if latest.is_symlink() or latest.exists():
        latest.unlink() if latest.is_symlink() else shutil.rmtree(latest)
    latest.symlink_to(d.name)
    return d


def do_run(args):
    man = load_manifest()
    cfg = man["config"]
    check_config(cfg)

    tests = select(man, args.only, args.group)
    if not tests:
        sys.exit("no test matches")

    rundir = new_result_dir()
    assets = (make_assets(cfg, RESULT)
              if any(t["kind"] == "video" for t in tests) else (None, None))

    results = {}
    for t in tests:
        name = t["name"]
        print(f"{C_DIM}..{C_OFF} {name} ", end="", flush=True)
        fn = KINDS.get(t["kind"])
        if not fn:
            r = {"rc": 1, "err": f"unknown kind: {t['kind']}"}
        elif t["kind"] == "video":
            r = fn(t, cfg, rundir / name, assets)
        else:
            r = fn(t, cfg, rundir / name)
        r["group"] = t.get("group", "other")
        r["touched"] = t.get("touched", "")
        results[name] = r
        print(fmt_one(r))

    (rundir / "results.json").write_text(json.dumps(results, indent=2))

    base = json.loads(BASELINE.read_text()) if BASELINE.exists() else {}
    report = diff_report(base, results, args.group_by)
    (rundir / "summary.txt").write_text(strip_color(report))
    print("\n" + report)
    print(f"{C_DIM}-> {rundir}{C_OFF}")

    if args.update_baseline:
        BASELINE.write_text(json.dumps(results if not (args.only or args.group)
                                       else {**base, **results}, indent=2))
        print(f"{C_DIM}baseline updated{C_OFF}")

    return 1 if has_regression(base, results) else 0


def do_capture(args):
    man = load_manifest()
    cfg = man["config"]
    check_config(cfg)

    prof = man.get("capture", {}).get(args.profile)
    if not prof:
        sys.exit(f"profile not found: {args.profile}\n"
                 f"available: {', '.join(man.get('capture', {}))}")
    t = next((x for x in man["test"] if x["name"] == args.name), None)
    if not t:
        sys.exit(f"test not found: {args.name}")

    outdir = new_result_dir() / args.name / "capture"
    outdir.mkdir(parents=True, exist_ok=True)

    env = dict(driver_env(cfg))
    env.update(prof.get("env", {}))
    print(f"{C_DIM}profile: {args.profile} — {prof.get('desc', '')}{C_OFF}")
    for k, v in env.items():
        print(f"{C_DIM}  {k}={v}{C_OFF}")

    if t["kind"] == "cts":
        cmd = cfg["cts_cmd"].format(filter=t["filter"], log=outdir / "TestResults.qpa")
        cwd = cfg.get("cts_dir")
    elif t["kind"] == "video":
        ref, src = make_assets(cfg, RESULT)
        dev = cfg.get("vulkan_device", "vk:0")
        cmd = (f"ffmpeg -v verbose -hwaccel vulkan -i {src} -f null -"
               if t.get("op") == "decode" else
               f"ffmpeg -y -v verbose -init_hw_device vulkan={dev} "
               f"-filter_hw_device vk -i {ref} -vf format=nv12,hwupload "
               f"-c:v h264_vulkan -qp 18 -g 1 {outdir}/enc.h264")
        cwd = None
    else:
        sys.exit(f"capture does not support kind={t['kind']}")

    wrapper = prof.get("wrapper", "").format(out=outdir)
    rc, _ = run(f"{wrapper} {cmd}".strip(), outdir, "capture", env=env, cwd=cwd)

    print(f"rc={rc}  ->  {outdir}")
    for f in sorted(outdir.iterdir()):
        print(f"  {f.name}  {f.stat().st_size} B")
    return rc


def do_doctor(args):
    cfg = load_manifest()["config"]
    bad = 0
    for label, value, ok in probe_env(cfg):
        mark = f"{C_OK}ok {C_OFF}" if ok else f"{C_BAD}MISSING{C_OFF}"
        bad += not ok
        print(f"  {mark:<12} {label:<14} {value}")
    if bad:
        print(f"\n{C_DIM}enter the lab shell with 'nix develop' so the paths "
              f"get exported, or override them in the environment{C_OFF}")
    return 1 if bad else 0


def do_list(args):
    man = load_manifest()
    groups = {}
    for t in man["test"]:
        groups.setdefault(t.get("group", "other"), []).append(t)
    for g in sorted(groups):
        print(f"\n{C_DIM}[{g}]{C_OFF}")
        for t in groups[g]:
            print(f"  {t['name']:<28} {C_DIM}{t.get('touched', '')}{C_OFF}")
    print(f"\n{C_DIM}capture profiles:{C_OFF}")
    for k, v in man.get("capture", {}).items():
        print(f"  {k:<12} {C_DIM}{v.get('desc', '')}{C_OFF}")


def do_clean(args):
    if RESULT.exists():
        shutil.rmtree(RESULT)
        print("result/ removed")


def main():
    ap = argparse.ArgumentParser(prog="ivb")
    sub = ap.add_subparsers(dest="cmd", required=True)

    r = sub.add_parser("run", help="run tests and diff against the baseline")
    r.add_argument("--only", nargs="+", metavar="NAME")
    r.add_argument("--group", nargs="+", metavar="G")
    r.add_argument("--group-by", choices=["group", "none"], default="group")
    r.add_argument("--update-baseline", action="store_true")
    r.set_defaults(fn=do_run)

    c = sub.add_parser("capture", help="run one test under a low-level profile")
    c.add_argument("name")
    c.add_argument("--profile", default="bat")
    c.set_defaults(fn=do_capture)

    sub.add_parser("doctor", help="check the driver trees the lab points at").set_defaults(fn=do_doctor)
    sub.add_parser("list", help="list tests and capture profiles").set_defaults(fn=do_list)
    sub.add_parser("clean", help="remove result/").set_defaults(fn=do_clean)

    args = ap.parse_args()
    sys.exit(args.fn(args) or 0)


if __name__ == "__main__":
    main()