import json, os, re, shutil, subprocess, sys, time, threading

ROOT = os.environ.get("UFEXP_ROOT", "/tmp/ufexp")
HERE = os.path.dirname(os.path.abspath(__file__))
MODEL = "gpt-6.1-sol"
TIERS = ["default", "fast", "ultrafast"]
REPEATS = int(os.environ.get("REPEATS", "3"))

PROMPTS = {
    "write": "Write a tutorial of about 1500 English words explaining how git rebase works, with 3 short code examples. Output the whole tutorial directly as your final answer. Do not use any tools.",
    "fix": "The tests in test_textstats.py fail. Fix textstats.py (do not edit the tests) so that every test passes. Verify by running `python3 -m unittest -q`.",
    "fix_slow": "The tests in test_textstats.py fail. Fix textstats.py (do not edit the tests) so that every test passes. Verify by running `python3 -m unittest -q`.",
    "logs": "app.log in this directory is a large server log. Find the request_id of the slowest request (highest latency_ms) among lines with svc=billing and status=500. Answer with the request_id only.",
}
SRC = {"write": None, "fix": "fix", "fix_slow": "fix_slow", "logs": "logs"}

def prepare():
    with open(f"{HERE}/fix/test_textstats.py") as f:
        test = f.read()
    if test.count("SETUP_DELAY = 0") != 1:
        raise ValueError("expected exactly one SETUP_DELAY = 0 in the published fixture")
    for kind in ("fix", "fix_slow"):
        d = f"{ROOT}/src/{kind}"
        shutil.rmtree(d, ignore_errors=True)
        os.makedirs(d)
        for name in ("textstats.py", "test_textstats.py"):
            shutil.copy(f"{HERE}/fix/{name}", d)
        if kind == "fix_slow":
            with open(f"{d}/test_textstats.py", "w") as f:
                f.write(test.replace("SETUP_DELAY = 0", "SETUP_DELAY = 15"))
    d = f"{ROOT}/src/logs"
    shutil.rmtree(d, ignore_errors=True)
    os.makedirs(d)
    subprocess.run([sys.executable, f"{HERE}/logs/make_log.py"], cwd=d, check=True)

def check(kind, d, final):
    if kind == "write":
        n = len(final.split())
        return 1200 <= n <= 1900, f"{n} words"
    if kind in ("fix", "fix_slow"):
        r = subprocess.run(["python3", "-m", "unittest", "-q"], cwd=d, capture_output=True, text=True)
        # slow variant: don't re-wait, run with delay patched to 0
        return r.returncode == 0, (r.stderr.strip().splitlines() or [""])[-1]
    if kind == "logs":
        return "req-ANSWER42" in final, final.strip()[:40]

def one(kind, tier, rep):
    d = f"{ROOT}/work/{kind}-{tier}-{rep}"
    shutil.rmtree(d, ignore_errors=True); os.makedirs(d)
    if SRC[kind]:
        for f in os.listdir(f"{ROOT}/src/{SRC[kind]}"):
            shutil.copy(f"{ROOT}/src/{SRC[kind]}/{f}", d)
    cmd = ["codex", "exec", "--json", "-m", MODEL, "-c", f'service_tier="{tier}"',
           "-c", 'model_reasoning_effort="low"', "--skip-git-repo-check", "-s", "workspace-write",
           "-C", d, PROMPTS[kind]]
    t0 = time.time()
    p = subprocess.run(cmd, stdin=subprocess.DEVNULL, capture_output=True, text=True, timeout=900)
    wall = time.time() - t0
    events = [json.loads(l) for l in p.stdout.splitlines() if l.startswith("{")]
    open(f"{ROOT}/work/{kind}-{tier}-{rep}.jsonl", "w").write(p.stdout)
    usage = next((e["usage"] for e in events if e["type"] == "turn.completed"), {})
    cmds = sum(1 for e in events if e["type"] == "item.completed" and e["item"]["type"] == "command_execution")
    msgs = [e["item"]["text"] for e in events if e["type"] == "item.completed" and e["item"]["type"] == "agent_message"]
    final = msgs[-1] if msgs else ""
    ok, note = check(kind, d, final)
    if kind == "fix_slow":  # re-check quickly with the delay removed
        shutil.copy(f"{ROOT}/src/fix/test_textstats.py", d + "/test_textstats.py")
        ok, note = check("fix", d, final)
    return dict(kind=kind, tier=tier, rep=rep, wall=round(wall, 1), cmds=cmds, ok=ok, note=note, **usage)

def main(kinds):
    out = []
    for kind in kinds:
        for rep in range(REPEATS):
            res = {}
            def go(t):
                res[t] = one(kind, t, rep)
            th = [threading.Thread(target=go, args=(t,)) for t in TIERS]
            [x.start() for x in th]; [x.join() for x in th]
            for t in TIERS:
                out.append(res[t]); print(json.dumps(res[t]), flush=True)
            json.dump(out, open(f"{ROOT}/results.json", "w"), ensure_ascii=False, indent=1)

if __name__ == "__main__":
    kinds = sys.argv[1:]
    if kinds != ["--prepare-only"] and (not kinds or any(k not in PROMPTS for k in kinds)):
        sys.exit("usage: python3 run.py --prepare-only | <write|fix|fix_slow|logs> [...]")
    prepare()
    if kinds != ["--prepare-only"]:
        main(kinds)