#!/usr/bin/env python3
"""Run the BENCH-MATRIX blocks declared in matrix.toml through `bench.py session`.

    python3 matrix.py list                          # blocks, session counts, est. time
    python3 matrix.py plan B3 [--param k=v]         # the bench.py commands a block runs
    python3 matrix.py run B0,B1 [--dry-run]         # sequential, resumable
    python3 matrix.py status                        # progress for the current mod set

Resume: ~/onion-bench/results/matrix-state.json records each session (key
"<block>#<index>") with a hash of its exact command and the instance mods/ hash; a
session is skipped only when both still match. A failed session (non-zero exit: crash,
timeout, no join) is retried once, then the run stops. Running any block other than B0
warns loudly when B0 was not run under the current mod set. Stdlib only.
"""
import argparse
import glob
import hashlib
import json
import os
import re
import shlex
import subprocess
import sys
import time
import tomllib
from pathlib import Path

HERE = Path(__file__).resolve().parent
MATRIX = HERE / "matrix.toml"
OUT_DIR = HERE.parent / "proposed-configs" / "out"
GAME = Path.home() / "onion-bench" / "prism" / "instances" / "OnionBench" / ".minecraft"
STATE = Path.home() / "onion-bench" / "results" / "matrix-state.json"

SCENE_ORDER = ["s1", "s2", "s2s", "s3", "s4", "s5"]
S5_SECONDS = 10 + 120  # bench.py: at S1, wait 10 s, then the 120 s tour (all captured)
JVM_SETS = {"J0", "J1", "J2", "J3"}
DISPLAYS = {"desktop", "gamescope", "gamescope-headless"}
PARAM_RE = re.compile(r"\$\{(\w+)\}")

TOP_KEYS = {"defaults", "worlds", "trials", "missing_variants", "block"}
DEFAULT_KEYS = {"reps", "warmup", "capture", "display", "res", "world", "shader", "dh",
                "launch_overhead_s"}
BLOCK_KEYS = {"id", "title", "rule", "proxy_note", "manual_checks", "sessions", "params",
              "deferred", "manual", "interleave"}
SESSION_KEYS = {"variant", "plan", "reps", "repeat", "tag", "jvm", "heap", "shader",
                "shader_profile", "dh", "disable_mods", "add_mods", "world", "warmup",
                "capture", "manual", "note", "display", "res"}
MISSING_KEYS = {"id", "blocks", "contents"}
# Flags `bench.py session` accepts; bench_argv must emit only these.
BENCH_SESSION_FLAGS = {
    "--variant", "--plan", "--reps", "--warmup", "--capture", "--settle", "--label", "--jvm",
    "--heap", "--boot-timeout", "--timeout", "--any-order", "--kill-stale", "--dry-run",
    "--display", "--res", "--shader", "--shader-options", "--dh", "--disable-mods",
    "--add-mods", "--world-zip",
}


class MatrixError(Exception):
    pass


def log(msg):
    print(f"[matrix] {msg}", flush=True)


# ------------------------------------------------------------------ loading / validation

def load(path=MATRIX, out_dir=OUT_DIR):
    """Parse + validate matrix.toml. Raises MatrixError listing every problem."""
    with open(path, "rb") as f:
        m = tomllib.load(f)
    errors = validate(m, out_dir)
    if errors:
        raise MatrixError("matrix.toml invalid:\n  " + "\n  ".join(errors))
    return m


def split_plan(plan):
    if isinstance(plan, str):
        return [s.strip() for s in plan.split(",") if s.strip()]
    return list(plan)


def _pos_int(v):
    return isinstance(v, int) and not isinstance(v, bool) and v > 0


def validate(m, out_dir=OUT_DIR):
    err = []
    for k in m:
        if k not in TOP_KEYS:
            err.append(f"unknown top-level key {k!r}")
    d = m.get("defaults", {})
    for k in d:
        if k not in DEFAULT_KEYS:
            err.append(f"defaults: unknown key {k!r}")
    worlds, trials = m.get("worlds", {}), m.get("trials", {})
    missing_ids = set()
    for mv in m.get("missing_variants", []):
        for k in mv:
            if k not in MISSING_KEYS:
                err.append(f"missing_variants {mv.get('id')}: unknown key {k!r}")
        if not mv.get("id") or not mv.get("contents"):
            err.append("missing_variants entry needs id and contents")
        missing_ids.add(mv.get("id"))
    _check_common(d, "defaults", worlds, trials, err)

    seen = set()
    for b in m.get("block", []):
        bid = b.get("id")
        where = f"block {bid}"
        if not bid:
            err.append("block without id")
            continue
        if bid in seen:
            err.append(f"{where}: duplicate id")
        seen.add(bid)
        for k in b:
            if k not in BLOCK_KEYS:
                err.append(f"{where}: unknown key {k!r}")
        if not b.get("rule"):
            err.append(f"{where}: missing decision rule")
        params = set(b.get("params", []))
        sessions = b.get("sessions", [])
        if not sessions and not (b.get("manual") or b.get("deferred")):
            err.append(f"{where}: no sessions (mark manual or deferred)")
        groups = set()
        for i, s in enumerate(sessions):
            sw = f"{where} session {i}"
            for k in s:
                if k not in SESSION_KEYS:
                    err.append(f"{sw}: unknown key {k!r}")
            for k in ("variant", "jvm", "plan"):
                v = s.get(k)
                for p in PARAM_RE.findall(v if isinstance(v, str) else ",".join(v or [])):
                    if p not in params:
                        err.append(f"{sw}: ${{{p}}} not declared in params")
            var = s.get("variant")
            if not var:
                err.append(f"{sw}: missing variant")
            elif not PARAM_RE.search(var) and not (Path(out_dir) / var).is_dir() \
                    and var not in missing_ids:
                err.append(f"{sw}: variant {var!r} not in {out_dir} nor in missing_variants")
            plan = split_plan(s.get("plan", []))
            if not plan:
                err.append(f"{sw}: empty plan")
            concrete = [p for p in plan if not PARAM_RE.search(p)]
            for p in concrete:
                if p not in SCENE_ORDER:
                    err.append(f"{sw}: unknown scene {p!r} (S6 is manual-only)")
            idx = [SCENE_ORDER.index(p) for p in concrete if p in SCENE_ORDER]
            if idx != sorted(set(idx)):
                err.append(f"{sw}: plan must be in S1->S5 order without duplicates")
            if "repeat" in s and not _pos_int(s["repeat"]):
                err.append(f"{sw}: repeat must be a positive int")
            if "heap" in s and not _pos_int(s["heap"]):
                err.append(f"{sw}: heap must be a positive int (MB)")
            jvm = s.get("jvm")
            if jvm is not None and jvm not in JVM_SETS and not PARAM_RE.search(jvm):
                err.append(f"{sw}: jvm {jvm!r} not in {sorted(JVM_SETS)}")
            for t in s.get("add_mods", []):
                if t not in trials:
                    err.append(f"{sw}: add_mods {t!r} not in [trials]")
            if not isinstance(s.get("disable_mods", []), list):
                err.append(f"{sw}: disable_mods must be a list")
            _check_common(s, sw, worlds, trials, err)
            g = (var, s.get("tag"))
            if g in groups:
                err.append(f"{sw}: variant {var!r} repeated without a distinct tag "
                           f"(analyze.py would merge them); use repeat or tag")
            groups.add(g)
    return err


def _check_common(s, where, worlds, trials, err):
    for k in ("reps", "warmup", "capture"):
        if k in s and not _pos_int(s[k]):
            err.append(f"{where}: {k} must be a positive int")
    if "display" in s and s["display"] not in DISPLAYS:
        err.append(f"{where}: display {s['display']!r} not in {sorted(DISPLAYS)}")
    if "res" in s and not re.fullmatch(r"\d+x\d+", str(s["res"])):
        err.append(f"{where}: res must be WxH")
    if "dh" in s and s["dh"] not in ("on", "off"):
        err.append(f"{where}: dh must be on/off")
    if "world" in s and s["world"] not in worlds:
        err.append(f"{where}: world {s['world']!r} not in [worlds]")


def block_map(m):
    return {b["id"]: b for b in m.get("block", [])}


def missing_variant_ids(m):
    return {mv["id"] for mv in m.get("missing_variants", [])}


# ------------------------------------------------------------------ expansion / commands

def _subst(v, params):
    if not isinstance(v, str):
        return v
    return PARAM_RE.sub(lambda mo: params.get(mo.group(1), mo.group(0)), v)


def expand(m, block, params=None, overrides=None):
    """Flat list of concrete sessions for a block (repeat / interleave applied)."""
    params, overrides = params or {}, overrides or {}
    d = dict(m.get("defaults", {}))
    d.update({k: v for k, v in overrides.items() if v is not None})
    groups = []
    for si, s in enumerate(block.get("sessions", [])):
        merged = {k: v for k, v in d.items() if k != "launch_overhead_s"}
        merged.update(s)
        for k in ("display", "res"):  # CLI overrides beat per-session values too
            if overrides.get(k) is not None:
                merged[k] = overrides[k]
        merged["variant"] = _subst(merged["variant"], params)
        if merged.get("jvm") is not None:
            merged["jvm"] = _subst(merged["jvm"], params)
        merged["plan"] = [_subst(p, params) for p in split_plan(merged["plan"])]
        tag = s.get("tag")
        merged["label"] = f"{block['id']}.{tag}" if tag else block["id"]
        merged["source"] = si
        groups.append([dict(merged, run=r + 1) for r in range(s.get("repeat", 1))])
    out = []
    if block.get("interleave"):
        for r in range(max((len(g) for g in groups), default=0)):
            out += [g[r] for g in groups if r < len(g)]
    else:
        for g in groups:
            out += g
    for i, s in enumerate(out):
        s["index"] = i
        s["block"] = block["id"]
    return out


def unresolved(sess):
    """Placeholder names still present in a session (missing --param)."""
    vals = [sess["variant"], sess.get("jvm") or "", *sess["plan"]]
    return sorted({p for v in vals for p in PARAM_RE.findall(v)})


def resolve_path(p):
    p = os.path.expanduser(p)
    hits = sorted(glob.glob(p))
    return hits[-1] if hits else None


def bench_argv(m, sess, live_world=False):
    """`bench.py session` argv for one session, plus a list of missing inputs."""
    missing = []
    a = ["session", "--variant", sess["variant"], "--plan", ",".join(sess["plan"]),
         "--reps", str(sess.get("reps", 3)), "--label", sess["label"]]
    for k in ("warmup", "capture"):
        if k in sess:
            a += [f"--{k}", str(sess[k])]
    for k in ("display", "res", "shader", "dh", "jvm"):
        if sess.get(k) is not None:
            a += [f"--{k}", str(sess[k])]
    if sess.get("heap"):
        a += ["--heap", str(sess["heap"])]
    if sess.get("disable_mods"):
        a += ["--disable-mods", ",".join(sess["disable_mods"])]
    if sess.get("add_mods"):
        jars = []
        for t in sess["add_mods"]:
            pat = m["trials"][t]
            hit = resolve_path(pat)
            if hit is None:
                missing.append(f"trial jar {t}: {pat}")
                hit = os.path.expanduser(pat)
            jars.append(hit)
        a += ["--add-mods", ",".join(jars)]
    if sess.get("world") and not live_world:
        pat = m["worlds"][sess["world"]]
        hit = resolve_path(pat)
        if hit is None:
            missing.append(f"world {sess['world']}: {pat} (or run with --live-world)")
            hit = os.path.expanduser(pat)
        a += ["--world-zip", hit]
    return a, missing


def variant_missing(sess, out_dir=OUT_DIR):
    return not (Path(out_dir) / sess["variant"]).is_dir()


def session_seconds(sess, overhead):
    reps, w, c = sess.get("reps", 3), sess.get("warmup", 15), sess.get("capture", 30)
    scenes = sum(S5_SECONDS if p == "s5" else w + c for p in sess["plan"]) * reps
    return overhead + scenes


def spec_hash(argv, sess):
    blob = json.dumps({"argv": argv, "manual": bool(sess.get("manual"))}, sort_keys=True)
    return hashlib.sha256(blob.encode()).hexdigest()[:16]


def fmt_cmd(argv):
    return "python3 bench.py " + shlex.join(argv)


# ------------------------------------------------------------------ state

def mods_hash(mods_dir):
    """Same algorithm as bench.py mods_hash (name + size of every file in mods/)."""
    h = hashlib.sha256()
    files = sorted(p for p in Path(mods_dir).glob("*") if p.is_file())
    for p in files:
        h.update(f"{p.name}\0{p.stat().st_size}\n".encode())
    return h.hexdigest()


class State:
    def __init__(self, path):
        self.path = Path(path)
        self.data = {"version": 1, "sessions": {}, "b0": None}
        if self.path.exists():
            self.data.update(json.loads(self.path.read_text()))

    def save(self):
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        tmp.write_text(json.dumps(self.data, indent=1, sort_keys=True))
        tmp.replace(self.path)

    @staticmethod
    def key(sess):
        return f"{sess['block']}#{sess['index']}"

    def entry(self, sess):
        return self.data["sessions"].get(self.key(sess))

    def is_done(self, sess, spec, mods):
        e = self.entry(sess)
        return bool(e and e["status"] == "ok" and e["spec"] == spec and e["mods_hash"] == mods)

    def record(self, sess, spec, mods, status, rc, attempts):
        self.data["sessions"][self.key(sess)] = {
            "spec": spec, "mods_hash": mods, "status": status, "rc": rc,
            "attempts": attempts, "label": sess["label"], "variant": sess["variant"],
            "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}


def b0_warning(state, current):
    b0 = state.data.get("b0")
    if not b0:
        return "B0 (noise floor) has not been completed - win thresholds are undefined"
    if b0["mods_hash"] != current:
        return (f"MOD SET CHANGED since B0 (B0 under {b0['mods_hash'][:12]}, now "
                f"{current[:12]}): B0 MUST BE RE-RUN before results are comparable")
    return None


def banner(msg):
    bar = "!" * 78
    print(f"{bar}\n!! {msg}\n{bar}", file=sys.stderr, flush=True)


# ------------------------------------------------------------------ run

def default_runner(argv):
    return subprocess.run([sys.executable, str(HERE / "bench.py"), *argv], cwd=HERE).returncode


def run_blocks(m, block_ids, state, mods_dir, runner=default_runner, params=None,
               overrides=None, dry_run=False, include_manual=False, include_deferred=False,
               live_world=False, skip_blocked=False, out_dir=OUT_DIR):
    """Run blocks in order. Returns 0 ok, 1 session failed, 2 preflight problem."""
    bm = block_map(m)
    for bid in block_ids:
        if bid not in bm:
            raise MatrixError(f"unknown block {bid}; known: {', '.join(bm)}")
    # preflight: build every session's command, collect what can't run
    todo, problems = [], []
    for bid in block_ids:
        b = bm[bid]
        if b.get("deferred") and not include_deferred:
            log(f"{bid}: deferred ({b.get('proxy_note', 'fleet-only')}) - skipped")
            continue
        if b.get("manual"):
            log(f"{bid}: manual block - run by hand ({b.get('proxy_note', '')})")
            continue
        for s in expand(m, b, params, overrides):
            if s.get("manual") and not include_manual:
                log(f"{bid}#{s['index']} {s['label']}: manual - skipped")
                continue
            miss = [f"param ${{{p}}} (use --param {p}=...)" for p in unresolved(s)]
            if not miss and variant_missing(s, out_dir):
                miss.append(f"variant {s['variant']} not generated (see missing_variants)")
            argv, inputs = bench_argv(m, s, live_world)
            miss += inputs
            if miss:
                problems.append((s, miss))
                continue
            todo.append((s, argv))
    for s, miss in problems:
        log(f"BLOCKED {s['block']}#{s['index']} {s['label']}: " + "; ".join(miss))
    if problems and not (skip_blocked or dry_run):
        log("fix the above, or pass --skip-blocked (and/or --live-world)")
        return 2
    mods0 = mods_hash(mods_dir)
    if "B0" not in block_ids and any(s["block"] != "B0" for s, _ in todo):
        w = b0_warning(state, mods0)
        if w:
            banner(w)
    total = len(todo)
    for n, (s, argv) in enumerate(todo, 1):
        spec = spec_hash(argv, s)
        tagline = f"[{n}/{total}] {s['block']}#{s['index']} {s['label']}"
        if state.is_done(s, spec, mods0):
            log(f"{tagline}: done - skip")
            continue
        if dry_run:
            log(f"{tagline}: {fmt_cmd(argv)}")
            continue
        now = mods_hash(mods_dir)
        if now != mods0:
            banner(f"instance mods/ changed during the run ({mods0[:12]} -> {now[:12]}); "
                   "a session did not restore it - stopping")
            return 1
        rc = None
        for attempt in (1, 2):
            log(f"{tagline}: attempt {attempt}: {fmt_cmd(argv)}")
            rc = runner(argv)
            if rc == 0:
                break
            log(f"{tagline}: exit {rc}" + (" - retrying once" if attempt == 1 else ""))
        status = "ok" if rc == 0 else "failed"
        state.record(s, spec, mods0, status, rc, attempt)
        if s["block"] == "B0":
            _maybe_close_b0(m, state, mods0, params, overrides, live_world)
        state.save()
        if status != "ok":
            log(f"{tagline}: failed twice - stopping the run")
            return 1
    return 0


def _maybe_close_b0(m, state, mods, params, overrides, live_world):
    b = block_map(m).get("B0")
    if not b:
        return
    for s in expand(m, b, params, overrides):
        argv, _ = bench_argv(m, s, live_world)
        if not state.is_done(s, spec_hash(argv, s), mods):
            return
    state.data["b0"] = {"mods_hash": mods, "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}


# ------------------------------------------------------------------ list / plan / status

def block_summary(m, b, overhead, out_dir=OUT_DIR):
    ss = expand(m, b)
    auto = [s for s in ss if not s.get("manual") and not b.get("manual")]
    blocked = [s for s in auto if not unresolved(s) and variant_missing(s, out_dir)]
    secs = sum(session_seconds(s, overhead) for s in auto)
    return {"sessions": len(ss), "auto": len(auto), "manual": len(ss) - len(auto),
            "blocked": len(blocked), "seconds": secs}


def cmd_list(m, a):
    overhead = m.get("defaults", {}).get("launch_overhead_s", 60)
    print(f"| block | title | sessions | auto | manual | blocked (variant) | est. min | notes |")
    print("| --- | --- | --- | --- | --- | --- | --- | --- |")
    tot = {"sessions": 0, "auto": 0, "seconds": 0}
    for b in m.get("block", []):
        s = block_summary(m, b, overhead)
        notes = []
        if b.get("deferred"):
            notes.append("deferred")
        if b.get("manual"):
            notes.append("manual block")
        if b.get("params"):
            notes.append("params: " + ",".join(b["params"]))
        print(f"| {b['id']} | {b.get('title', '')} | {s['sessions']} | {s['auto']} | "
              f"{s['manual']} | {s['blocked']} | {s['seconds'] / 60:.0f} | {'; '.join(notes)} |")
        if not b.get("deferred"):
            for k in tot:
                tot[k] += s[k]
    print(f"\nnot deferred: {tot['sessions']} sessions ({tot['auto']} unattended), "
          f"est. {tot['seconds'] / 3600:.1f} h unattended")
    return 0


def cmd_plan(m, a):
    bm = block_map(m)
    for bid in a.blocks.split(","):
        if bid not in bm:
            raise MatrixError(f"unknown block {bid}")
        b = bm[bid]
        print(f"## {bid} {b.get('title', '')}" + (" (DEFERRED)" if b.get("deferred") else "")
              + (" (MANUAL)" if b.get("manual") else ""))
        print(f"rule: {b['rule']}")
        if b.get("proxy_note"):
            print(f"proxy: {b['proxy_note']}")
        for c in b.get("manual_checks", []):
            print(f"manual check: {c}")
        for s in expand(m, b, a.params, overrides(a)):
            argv, miss = bench_argv(m, s, a.live_world)
            pre = "# MANUAL: " if s.get("manual") else ""
            print(f"{pre}{fmt_cmd(argv)}")
            if s.get("shader_profile"):
                print(f"#   set Iris profile {s['shader_profile']} for {s['shader']} first")
            if s.get("note"):
                print(f"#   note: {s['note']}")
            if unresolved(s):
                print(f"#   needs --param {','.join(unresolved(s))}")
            if variant_missing(s) and not unresolved(s):
                print(f"#   BLOCKED: variant {s['variant']} not generated")
            for x in miss:
                print(f"#   missing: {x}")
        print()
    return 0


def cmd_status(m, a):
    state = State(a.state)
    mods = mods_hash(a.mods_dir) if Path(a.mods_dir).is_dir() else ""
    w = b0_warning(state, mods)
    if w:
        banner(w)
    print(f"mods hash {mods[:12]} | state {a.state}")
    print("| block | done | failed | pending | manual/deferred |")
    print("| --- | --- | --- | --- | --- |")
    for b in m.get("block", []):
        done = failed = pending = skipped = 0
        for s in expand(m, b, a.params, overrides(a)):
            if s.get("manual") or b.get("manual") or b.get("deferred"):
                skipped += 1
                continue
            argv, _ = bench_argv(m, s, a.live_world)
            e = state.entry(s)
            if state.is_done(s, spec_hash(argv, s), mods):
                done += 1
            elif e and e["status"] == "failed":
                failed += 1
            else:
                pending += 1
        print(f"| {b['id']} | {done} | {failed} | {pending} | {skipped} |")
    return 0


def cmd_run(m, a):
    state = State(a.state)
    return run_blocks(m, a.blocks.split(","), state, a.mods_dir, params=a.params,
                      overrides=overrides(a), dry_run=a.dry_run,
                      include_manual=a.include_manual, include_deferred=a.include_deferred,
                      live_world=a.live_world, skip_blocked=a.skip_blocked)


def overrides(a):
    return {"display": getattr(a, "display", None), "res": getattr(a, "res", None)}


def parse_params(items):
    out = {}
    for it in items or []:
        k, sep, v = it.partition("=")
        if not sep:
            raise MatrixError(f"--param needs name=value, got {it!r}")
        out[k] = v
    return out


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--matrix", default=str(MATRIX))
    ap.add_argument("--state", default=str(STATE))
    ap.add_argument("--mods-dir", default=str(GAME / "mods"))
    sub = ap.add_subparsers(dest="cmd", required=True)
    sub.add_parser("list", help="blocks, session counts, estimated time")
    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--param", action="append", default=[], help="name=value")
    common.add_argument("--display", choices=sorted(DISPLAYS))
    common.add_argument("--res")
    common.add_argument("--live-world", action="store_true",
                        help="don't restore world zips; use saves/onionbench as-is")
    p = sub.add_parser("plan", parents=[common], help="print the bench.py commands")
    p.add_argument("blocks")
    r = sub.add_parser("run", parents=[common], help="run blocks sequentially (resumable)")
    r.add_argument("blocks", help="comma list, e.g. B0,B1")
    r.add_argument("--dry-run", action="store_true")
    r.add_argument("--include-manual", action="store_true")
    r.add_argument("--include-deferred", action="store_true")
    r.add_argument("--skip-blocked", action="store_true",
                   help="skip sessions with missing variants/inputs instead of aborting")
    sub.add_parser("status", parents=[common], help="progress for the current mod set")
    a = ap.parse_args(argv)
    try:
        a.params = parse_params(getattr(a, "param", []))
        m = load(a.matrix)
        return {"list": cmd_list, "plan": cmd_plan, "run": cmd_run,
                "status": cmd_status}[a.cmd](m, a)
    except MatrixError as e:
        print(f"matrix.py: {e}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())