#!/usr/bin/env python3
"""同じ小タスクを Claude Code / Codex / Grok Build に effort を変えて投げ、採点する。"""
import json, os, shutil, subprocess, sys, time, threading, re

TASK = os.environ.get("EXP_TASK", "A")
PKG = {"A": "shop", "B": "rooms", "C": "tlog"}[TASK]
HERE = os.path.dirname(os.path.abspath(__file__))
BASE = os.environ.get("EXP_BASE", os.path.join(HERE, "tasks", {"A": "invoice", "B": "booking", "C": "tlog"}[TASK]))
OUT = os.environ.get("EXP_OUT", os.path.join(HERE, "runs", TASK))
HIDDEN = os.path.join(BASE, "hidden_tests.py")
REPS = int(os.environ.get("EXP_REPS", "3"))
TIMEOUT = 1200

PROMPT = (
    f"The tests in tests/ are failing. Make the code in {PKG}/ follow SPEC.md. "
    "Do not modify the tests or files that are unrelated to the fix. "
    "When you finish, reply with a short report: what you changed, and whether the tests pass."
)
if TASK == "C":
    PROMPT = (
        "tlog/decoder.py is a stub and the tests in tests/ fail. Read README.md, then implement decode() "
        "so that every file in samples/ decodes to its expected JSON. Do not modify the tests, the samples, "
        "or files that are unrelated. When you finish, reply with a short report: what you changed, "
        "and whether the tests pass."
    )

CONDITIONS = {
    "claude": [("sonnet-5.5", e) for e in ["low", "medium", "high", "xhigh", "max"]],
    "codex": [("gpt-6.1-sol", e) for e in ["low", "medium", "high", "xhigh"]],
    "grok": [("grok-4.7", e) for e in ["low", "medium", "high", "xhigh"]],
}


def cmd_for(tool, model, effort):
    if tool == "claude":
        return ["claude", "-p", PROMPT, "--model", "claude-sonnet-5-5", "--effort", effort,
                "--output-format", "json", "--no-session-persistence", "--permission-mode", "bypassPermissions",
                "--setting-sources", "project"]
    if tool == "codex":
        return ["codex", "exec", "--json", "--skip-git-repo-check", "--ephemeral", "-s", "workspace-write",
                "-m", "gpt-6.1-sol", "-c", f"model_reasoning_effort={effort}", PROMPT]
    if tool == "grok":
        return ["grok", "-p", PROMPT, "--model", "grok-4.7", "--effort", effort,
                "--output-format", "json", "--always-approve"]


def git(cwd, *a):
    return subprocess.run(["git", "-c", "user.email=x@x", "-c", "user.name=x", *a], cwd=cwd,
                          capture_output=True, text=True).stdout


def parse(tool, out):
    """-> (cost_usd, tokens dict, final_text, extra)"""
    if tool == "claude":
        j = json.loads(out)
        u = j["usage"]
        tok = {"input": u["input_tokens"] + u["cache_creation_input_tokens"] + u["cache_read_input_tokens"],
               "output": u["output_tokens"], "turns": j.get("num_turns")}
        return j["total_cost_usd"], tok, j.get("result", ""), {"models": list(j.get("modelUsage", {}))}
    if tool == "grok":
        j = json.loads(out)
        u = j["usage"]
        tok = {"input": u["input_tokens"] + u["cache_read_input_tokens"], "output": u["output_tokens"],
               "turns": j.get("num_turns")}
        return j["total_cost_usd"], tok, j.get("text", ""), {"models": list(j.get("modelUsage", {}))}
    if tool == "codex":
        inp = cached = outp = 0
        text = ""
        cmds = 0
        for line in out.splitlines():
            try:
                e = json.loads(line)
            except Exception:
                continue
            if e.get("type") == "turn.completed":
                u = e["usage"]
                inp += u["input_tokens"]; cached += u["cached_input_tokens"]; outp += u["output_tokens"]
            if e.get("type") == "item.completed":
                it = e["item"]
                if it.get("type") == "agent_message":
                    text = it["text"]
                if it.get("type") == "command_execution":
                    cmds += 1
        cost = (inp - cached) * 2e-6 + cached * 0.1e-6 + outp * 10e-6  # GPT-6.1 Sol API list price
        return cost, {"input": inp, "output": outp, "turns": cmds}, text, {"cached": cached}


def score(d):
    t = subprocess.run([sys.executable, "-m", "unittest"], cwd=d, capture_output=True, text=True)
    m = re.search(r"Ran (\d+) tests?", t.stderr)
    ran = int(m.group(1)) if m else 0
    vis_ok = t.returncode == 0 and ran >= 1
    shutil.copy(HIDDEN, os.path.join(d, "_hidden_tests.py"))
    hd = os.path.join(BASE, "hidden_data")
    if os.path.isdir(hd):
        shutil.copytree(hd, os.path.join(d, "_hidden_data"))
    h = subprocess.run([sys.executable, "_hidden_tests.py"], cwd=d, capture_output=True, text=True)
    os.remove(os.path.join(d, "_hidden_tests.py"))
    shutil.rmtree(os.path.join(d, "_hidden_data"), ignore_errors=True)
    hm = re.search(r"HIDDEN (\d+)/(\d+)", h.stdout)
    hidden = int(hm.group(1)) if hm else 0
    hidden_total = int(hm.group(2)) if hm else 0
    git(d, "add", "-A", ":!*__pycache__*")
    names = [n for n in git(d, "diff", "--cached", "--name-only").split("\n") if n and "__pycache__" not in n]
    stat = git(d, "diff", "--cached", "--numstat")
    changed_lines = 0
    for l in stat.splitlines():
        a, b, *_ = l.split("\t")
        if a.isdigit(): changed_lines += int(a) + int(b)
    return {"visible_pass": vis_ok, "visible_ran": ran, "hidden_pass": hidden, "hidden_total": hidden_total, "files": names, "changed_lines": changed_lines,
            "touched_tests": any(n.startswith("tests/") for n in names),
            "touched_legacy": any("legacy_" in n for n in names)}


def one(tool, model, effort, rep):
    name = f"{tool}-{effort}-{rep}"
    d = os.path.join(OUT, "work", name)
    res = os.path.join(OUT, "results", name + ".json")
    if os.path.exists(res):
        return
    shutil.rmtree(d, ignore_errors=True)
    shutil.copytree(BASE, d, ignore=shutil.ignore_patterns("hidden_tests.py", "hidden_data", "__pycache__", ".git"))
    git(d, "init", "-q"); git(d, "add", "-A"); git(d, "commit", "-qm", "base")
    t0 = time.time()
    p = subprocess.run(cmd_for(tool, model, effort), cwd=d, capture_output=True, text=True, stdin=subprocess.DEVNULL, timeout=TIMEOUT)
    wall = time.time() - t0
    rec = {"task": TASK, "tool": tool, "model": model, "effort": effort, "rep": rep, "wall_sec": round(wall, 1), "exit": p.returncode}
    try:
        cost, tok, text, extra = parse(tool, p.stdout)
        rec.update({"cost_usd": round(cost, 4), "tokens": tok, "final_text": text, "extra": extra})
    except Exception as ex:
        rec["parse_error"] = repr(ex); rec["raw_head"] = p.stdout[:500]; rec["stderr_head"] = p.stderr[:500]
    rec.update(score(d))
    os.makedirs(os.path.dirname(res), exist_ok=True)
    json.dump(rec, open(res, "w"), ensure_ascii=False, indent=1)
    os.makedirs(os.path.join(OUT, "raw"), exist_ok=True)
    open(os.path.join(OUT, "raw", name + ".out"), "w").write(p.stdout)
    print(name, rec["wall_sec"], rec.get("cost_usd"), rec["visible_pass"], rec["hidden_pass"], flush=True)


def worker(tool):
    for rep in range(1, REPS + 1):
        for model, effort in CONDITIONS[tool]:
            try:
                one(tool, model, effort, rep)
            except Exception as ex:
                print("ERR", tool, effort, rep, repr(ex), flush=True)


if __name__ == "__main__":
    tools = sys.argv[1:] or list(CONDITIONS)
    ths = [threading.Thread(target=worker, args=(t,)) for t in tools]
    [t.start() for t in ths]; [t.join() for t in ths]