"""Build the reasoning-effort SFT set: Beam traces at low, medium and high effort for the same prompts, each
tagged the way Dory-Beam's chat template renders `reasoning_effort` (Dory's own `{reasoning effort: ...}`
convention, appended to the user turn).

  .venv/bin/python -I scripts/build_effort_dataset.py --efforts low,medium,high --pool xhigh:high --longest high

Inputs: data/effort_prompts.jsonl (the 20k prompts), low/high/xhigh generations in data/selfhost_effort/*.jsonl,
and the existing medium generations (data/gen, data/selfhost). Per prompt and level, one randomly drawn
verified-correct trace, or the longest with --longest (used for high, pooled from Beam's high and xhigh). Medium examples are tagged half the time and left untagged otherwise, so the untagged
default keeps behaving like medium. Writes data/final_effort/sft_effort.jsonl and stats.json.
"""

import argparse
import glob
import json
import random
import statistics
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
from verify import check  # noqa: E402



def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompts", default=str(ROOT / "data" / "effort_prompts.jsonl"))
    ap.add_argument("--effort-inputs", default=str(ROOT / "data" / "selfhost_effort" / "*.jsonl"))
    ap.add_argument("--medium-inputs", nargs="+",
                    default=[str(ROOT / "data" / "gen" / "*.jsonl"), str(ROOT / "data" / "selfhost" / "*.jsonl")])
    ap.add_argument("--untagged-medium", type=float, default=0.5)
    ap.add_argument("--out-dir", default=str(ROOT / "data" / "final_effort"))
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--efforts", default="low,medium,high", help="Beam effort levels to include")
    ap.add_argument("--pool", default="", help="use one Beam level's traces as candidates for another, e.g. xhigh:high")
    ap.add_argument("--longest", default="", help="levels that take the longest correct candidate instead of a random one")
    args = ap.parse_args()
    EFFORTS = tuple(args.efforts.split(","))
    pool = dict(x.split(":") for x in args.pool.split(",") if x)
    longest = set(filter(None, args.longest.split(",")))

    prompts = {json.loads(l)["id"]: json.loads(l) for l in open(args.prompts)}
    eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()}
    assert not set(prompts) & eval_ids

    gens = defaultdict(list)  # (id, effort) -> verified records
    teachers = defaultdict(set)
    files = sorted(glob.glob(args.effort_inputs))
    files += sorted(p for pat in args.medium_inputs for p in glob.glob(pat))
    for path in files:
        for line in open(path):
            r = json.loads(line)
            if "error" in r or r["id"] not in prompts:
                continue
            beam_eff = r.get("effort", "medium")
            if (beam_eff == "medium") != ("selfhost_effort" not in path):
                continue  # medium only from the original pool, other levels only from the effort runs
            eff = pool.get(beam_eff, beam_eff)
            if eff not in EFFORTS:
                continue
            r["beam_effort"] = beam_eff
            ok, _ = check(r["content"], r["gold"], r["answer_type"])
            r["correct"] = ok and r.get("finish_reason") == "stop" and bool(r.get("reasoning", "").strip())
            gens[(r["id"], eff)].append(r)
            teachers[eff].add(r.get("teacher", "beam-api"))

    rng = random.Random(args.seed)
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    stats = {e: {"generations": 0, "correct": 0, "examples": 0, "chars": [], "tokens": []} for e in EFFORTS}
    n_untagged = 0
    with open(out / "sft_effort.jsonl", "w") as f:
        for pid in sorted(prompts):
            for eff in EFFORTS:
                rs = gens.get((pid, eff), [])
                good = sorted((r for r in rs if r["correct"]), key=lambda r: len(r["reasoning"]))
                st = stats[eff]
                st["generations"] += len(rs)
                st["correct"] += len(good)
                for r in rs:
                    st["tokens"].append(r["usage"].get("completion_tokens", 0))
                if not good:
                    continue
                r = good[-1] if eff in longest else rng.choice(good)
                tagged = eff != "medium" or rng.random() >= args.untagged_medium
                n_untagged += not tagged
                user = r["prompt"] + (f"\n\n{{reasoning effort: {eff}}}" if tagged else "")
                st["examples"] += 1
                st["chars"].append(len(r["reasoning"]) + len(r["content"]))
                f.write(json.dumps({
                    "id": f"{pid}#{eff}", "prompt_id": pid, "effort": eff, "beam_effort": r["beam_effort"],
                    "tagged": tagged,
                    "domain": r["domain"], "teacher": r.get("teacher", "beam-api"),
                    "messages": [{"role": "user", "content": user},
                                 {"role": "assistant", "reasoning_content": r["reasoning"], "content": r["content"]}],
                    "gold": r["gold"], "chars": st["chars"][-1],
                }, ensure_ascii=False) + "\n")

    summary = {"prompts": len(prompts), "untagged_medium": n_untagged, "pool": pool, "longest": sorted(longest),
               "teachers": {e: sorted(t) for e, t in teachers.items()}}
    for e, st in stats.items():
        summary[e] = {"generations": st["generations"], "correct": st["correct"],
                      "accuracy": round(st["correct"] / max(1, st["generations"]), 3), "examples": st["examples"],
                      "chars_mean": round(statistics.mean(st["chars"])) if st["chars"] else None,
                      "tokens_mean": round(statistics.mean(st["tokens"])) if st["tokens"] else None}
    (out / "stats.json").write_text(json.dumps(summary, indent=1))
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    main()