#!/usr/bin/env python3
"""Runs every taicpp test (C++ unit binaries and python differential drivers) and prints one
summary table. Exit 1 if any step fails; nothing is skipped silently: a step whose prerequisite is
missing (binary not built, tmux absent) is reported as FAIL or SKIP with the reason.

usage (from taicpp/, after `make test-build`; `make test` does both):
    python3 tests/run_all.py [--only SUBSTR ...] [--list] [--timeout SECONDS] [-v]

Output of every step goes to build/test-logs/<step>.log. Nothing is written outside build/.
A step is a list of commands that all must exit 0; ("diff", A, B) runs two commands and requires
identical stdout. Counts in the table are numbers the step itself printed (checks, cases,
scenarios, tests); "-" means the step prints no count.
"""
import argparse
import os
import re
import shutil
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)            # taicpp/
REPO = os.path.dirname(ROOT)            # repository root (tai.py lives here)
BUILD = os.path.join(ROOT, "build")
LOGS = os.path.join(BUILD, "test-logs")
TAI = os.path.join(REPO, "tai.py")
PY = sys.executable
STEP_TIMEOUT = 1200                     # hard cap per step, seconds


def b(name):
    return os.path.join(BUILD, name)


def t(name):
    return os.path.join(HERE, name)


# (name, module, [commands], env or None).  A command is an argv list or ("diff", argv, argv).
STEPS = [
    # foundation
    ("foundation.unit", "foundation", [[b("foundation_test")]], None),
    ("foundation.parity", "foundation", [[PY, t("foundation_parity.py")]], None),
    ("foundation.repl_persist", "foundation", [[PY, t("repl_persist.py")]], None),
    ("foundation.tmux_ask", "foundation", [[PY, t("tmux_ask.py")]], None),
    ("store.smoke", "foundation", [[b("store_test")]], None),
    ("cli.unit", "cli", [[b("cli_test")]], None),
    ("util_config.unit", "util/config", [[b("util_config_test")]], None),
    # agent
    ("agent.diff", "agent", [[PY, t("agent_diff.py")]], None),
    ("agent.e2e", "agent", [[PY, t("agent_e2e.py")]], None),
    ("agent.sched", "agent", [[PY, t("agent_sched.py")]], None),
    # shell and util
    ("shell.unit", "shell", [[b("shell_test")]], None),
    ("shell.parity", "shell", [[PY, t("shell_parity.py")]], None),
    # cap and mcp
    ("capmcp.unit", "cap/mcp", [[b("capmcp/capmcp_unit")]], None),
    ("capmcp.wire_diff", "cap/mcp", [[PY, "-I", t("capmcp_diff.py"), b("capmcp/capmcp_probe")]], None),
    ("capmcp.cli_diff", "cap/mcp", [[PY, "-I", t("capmcp_cli_diff.py"), b("capmcp/capmcp_cli"), "--tai", TAI]], None),
    ("capmcp.catalog", "cap/mcp", [[b("capmcp/cap_model_catalog_test")]], None),
    ("fuzzy.unit", "util", [[b("fuzzy_test")]], None),
    ("termshot.unit", "termshot", [[b("termshot_test")]], None),
    ("whatsapp.unit", "whatsapp", [[b("whatsapp_test"), "tests/whatsapp_data"]], None),
    ("muse.flags", "llm", [[b("muse_flags_test")]], None),
    ("llm.history", "llm", [[b("llm_history_test")]], None),
    # orchestration (sub agent panels on the multiplexer)
    ("orch.unit", "orch", [[b("orch_test")]], None),
    ("orch.mplex_e2e", "orch", [[PY, "-I", t("orch_mplex_e2e.py"), b("detz"), b("orch_live")]], None),
    # misc
    ("misc.unit", "misc", [[b("misc/misc_test")]], None),
    ("misc.diff", "misc", [[PY, t("misc_diff.py")]], None),
    # tui and grr
    ("tui.spec", "tui", [[b("tui_spec")]], None),
    ("tui.demo_shot", "tui", [("diff", [PY, TAI, "--demo-shot"], [b("tui_shot")])], None),
    # mplex
    ("mplex.foundation", "mplex", [[b("mplex_foundation_test")]], {"TZ": "UTC"}),
    ("mplex.core", "mplex", [[PY, "-I", t("mplex_core_gen.py")], [b("mplex_core_test")]], None),
    ("mplex.model", "mplex", [[PY, "-I", t("mplex_model_gen.py"), b("mplex_model_cases.jsonl"), b("mplex_model_scratch")],
                              [b("mplex_model_test"), b("mplex_model_cases.jsonl"), b("mplex_model_scratch")]], None),
    ("mplex.copy", "mplex", [[PY, "-I", t("mplex_copy_gen.py"), REPO, b("mplex_copy_scenarios.txt"), "3600", "20251025"],
                             [b("mplex_copy_test"), b("mplex_copy_scenarios.txt")]], None),
    ("mplex.vt", "mplex", [[PY, "-I", t("mplex_vt_oracle.py"), TAI, b("mplex_vt_cases.txt")],
                           [b("mplex_vt_test"), TAI, b("mplex_vt_cases.txt")]], None),
    ("mplex.flow", "mplex", [[PY, "-I", t("mplex_flow_gen.py"), REPO, b("mplex_flow_corpus.jsonl")],
                             [b("mplex_flow_test"), b("mplex_flow_corpus.jsonl")]], None),
    ("mplex.flow_tmux", "mplex", [("diff", [b("mplex_flow_test"), "--tmux"], [PY, "-I", t("mplex_tmux_oracle.py"), REPO])], None),
    ("mplex.regress", "mplex", [[b("mplex_regress_test")]], None),
    ("mplex.spec", "mplex", [[b("mplex_spec_test"), REPO]], None),
    ("mplex.server", "mplex", [[b("mplex_server_test"), "--tai", TAI]], None),
    ("mplex.commands", "mplex", [[b("mplex_commands_test"), b("mplex_cli"), REPO]], None),
    ("mplex.commands_pty", "mplex", [[PY, "-I", t("mplex_commands_pty.py"), b("mplex_cli"), REPO, b("mplex_commands_pty_work")]], None),
    ("mplex.e2e", "mplex", [[PY, "-I", t("mplex_e2e.py"), REPO, "--cpp", b("mplex_cli")]], None),
]

COUNT_RE = re.compile(r"(\d[\d,]*)\s+(checks?|cases?|tests?|scenarios?|assertions?|passed|ok|probes?|runs?)\b", re.I)
SLASH_RE = re.compile(r"\b(\d+)\s*/\s*(\d+)\b")


def count_of(text):
    """Best effort: the biggest number the step reported next to a unit word. None if none."""
    best = None
    for m in COUNT_RE.finditer(text):
        n = int(m.group(1).replace(",", ""))
        if best is None or n > best:
            best = n
    return best or None


def run_cmd(argv, env, log):
    start = time.time()
    try:
        done = subprocess.run(argv, cwd=ROOT, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                              stdin=subprocess.DEVNULL, timeout=STEP_TIMEOUT)
        out, code = done.stdout.decode("utf-8", "replace"), done.returncode
    except subprocess.TimeoutExpired as exc:
        out, code = (exc.stdout or b"").decode("utf-8", "replace") + "\nTIMEOUT after %ds\n" % STEP_TIMEOUT, 124
    except OSError as exc:
        out, code = "cannot start %s: %s\n" % (argv[0], exc), 127
    log.write("$ %s   (rc=%d, %.1fs)\n%s\n" % (" ".join(argv), code, time.time() - start, out))
    return code, out


def run_step(step):
    name, module, cmds, extra = step
    env = dict(os.environ)
    for key in [k for k in env if k.startswith("TAI_")]:
        env.pop(key, None)
    env.update(extra or {})
    start = time.time()
    status, note, text = "PASS", "", ""
    with open(os.path.join(LOGS, name + ".log"), "w") as log:
        for cmd in cmds:
            if isinstance(cmd, tuple):
                c1, o1 = run_cmd(cmd[1], env, log)
                c2, o2 = run_cmd(cmd[2], env, log)
                text += o1 + o2
                if c1 or c2:
                    status, note = "FAIL", "rc %d/%d" % (c1, c2)
                elif o1 != o2:
                    status, note = "FAIL", "outputs differ (%d vs %d bytes)" % (len(o1), len(o2))
                    log.write("DIFF: outputs differ\n")
                else:
                    text += "%d bytes identical\n" % len(o1)
            else:
                code, out = run_cmd(cmd, env, log)
                text += out
                if code:
                    status, note = "FAIL", "rc %d" % code
            if status == "FAIL":
                break
    if status == "PASS" and re.search(r"^SKIP", text, re.M):
        status, note = "SKIP", "reported SKIP (see log)"
    if status == "PASS" and re.search(r"\bFAIL\b|\b[1-9]\d* fail(ed|ures?)\b(?!=)", text):
        status, note = "FAIL", "exit 0 but output reports failures"
    return {"name": name, "module": module, "status": status, "note": note, "secs": time.time() - start,
            "count": count_of(text), "tail": text.strip().splitlines()[-12:]}


def snapshot():
    """Files under taicpp/ outside build/ and third_party/, to flag tests that leak outputs."""
    seen = set()
    for dirpath, dirnames, filenames in os.walk(ROOT):
        dirnames[:] = [d for d in dirnames if d not in ("build", "third_party", "__pycache__")]
        for name in filenames:
            seen.add(os.path.relpath(os.path.join(dirpath, name), ROOT))
    return seen


def main():
    global STEP_TIMEOUT
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", action="append", default=[])
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--timeout", type=int, default=STEP_TIMEOUT)
    ap.add_argument("-v", action="store_true")
    opts = ap.parse_args()
    STEP_TIMEOUT = opts.timeout
    steps = [s for s in STEPS if not opts.only or any(o in s[0] for o in opts.only)]
    if opts.list:
        for s in steps:
            print(s[0])
        return 0
    os.makedirs(LOGS, exist_ok=True)
    before = snapshot()
    results = []
    for s in steps:
        print("running %-26s" % s[0], end=" ", flush=True)
        r = run_step(s)
        print("%s %.1fs" % (r["status"], r["secs"]), flush=True)
        if r["status"] != "PASS":
            print("\n".join("    | " + line for line in r["tail"]))
        results.append(r)
    print()
    print("%-26s %-12s %-6s %8s %8s  %s" % ("step", "module", "status", "count", "secs", "note"))
    print("-" * 90)
    for r in results:
        print("%-26s %-12s %-6s %8s %8.1f  %s" % (r["name"], r["module"], r["status"],
              "-" if r["count"] is None else r["count"], r["secs"], r["note"]))
    print("-" * 90)
    npass = sum(r["status"] == "PASS" for r in results)
    nfail = sum(r["status"] == "FAIL" for r in results)
    nskip = sum(r["status"] == "SKIP" for r in results)
    total = sum(r["count"] or 0 for r in results)
    leaked = sorted(snapshot() - before)
    if leaked:
        print("WARNING: tests wrote outside build/: " + ", ".join(leaked[:10]))
    print("steps: %d  pass: %d  fail: %d  skip: %d  reported counts summed: %d  logs: %s" %
          (len(results), npass, nfail, nskip, total, os.path.relpath(LOGS, ROOT)))
    if nfail:
        print("RESULT: FAILED: " + ", ".join(r["name"] for r in results if r["status"] == "FAIL"))
        return 1
    print("RESULT: ok" + (" (with skips)" if nskip else ""))
    return 0


if __name__ == "__main__":
    sys.exit(main())