"""Assemble generated Beam traces into release/SFT files with summary stats.

Outputs (in --out-dir):
  traces_all.jsonl     every successful generation, correct or not, with metadata
                       (held-out eval prompts included, marked by eval_ids.txt)
  sft_verified.jsonl   verified-correct, untruncated, non-eval traces in "messages" form
                       (assistant turn carries reasoning_content, which Dory's chat
                       template renders as <think>...</think>)
  stats.json           counts, accuracy and token statistics per domain/band/effort
"""

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 band(r):
    if r["domain"] != "math":
        return "science"
    return f"math_d{min(int(r['difficulty'] or 0), 9)}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--inputs", nargs="+",
                    default=[str(ROOT / "data" / "gen" / "*.jsonl"), str(ROOT / "data" / "selfhost" / "*.jsonl")])
    ap.add_argument("--out-dir", default=str(ROOT / "data" / "final"))
    ap.add_argument("--max-per-prompt", type=int, default=1)
    # shortest: tersest style (run A). random: an unbiased draw from Beam's correct traces, which
    # keeps its self-checking on hard prompts (shortest-of-n favours lucky, unchecked solutions).
    ap.add_argument("--select", choices=["shortest", "random"], default="shortest")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--eval-ids", default=str(ROOT / "data" / "eval_ids.txt"))
    args = ap.parse_args()
    eval_ids = {line.strip() for line in open(args.eval_ids) if line.strip()}

    recs = []
    for path in sorted(p for pattern in args.inputs for p in glob.glob(pattern)):
        for line in open(path):
            r = json.loads(line)
            if "error" not in r:
                # Re-check with the current verifier so fixes apply retroactively.
                ok, pred = check(r["content"], r["gold"], r["answer_type"])
                r["pred"], r["correct"] = pred, ok and r.get("finish_reason") == "stop"
                recs.append(r)

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    by_id = defaultdict(list)
    for r in recs:
        by_id[r["id"]].append(r)

    # Per prompt, keep the shortest (or randomly drawn) verified-correct traces (length by
    # characters: SGLang n>1 requests only report total tokens, so per-sample counts are estimates).
    rng = random.Random(args.seed)
    n_sft = 0
    with open(out / "traces_all.jsonl", "w") as fa, open(out / "sft_verified.jsonl", "w") as fs:
        for r in recs:
            fa.write(json.dumps(r, ensure_ascii=False) + "\n")
        for pid, rs in by_id.items():
            if pid in eval_ids:
                continue
            good = [r for r in rs if r.get("correct") and r.get("reasoning", "").strip()]
            good.sort(key=lambda r: len(r["reasoning"]) + len(r["content"]))
            if args.select == "random":
                rng.shuffle(good)
            for r in good[: args.max_per_prompt]:
                n_sft += 1
                fs.write(
                    json.dumps(
                        {
                            "id": r["id"],
                            "domain": r["domain"],
                            "source": r["source"],
                            "effort": r["effort"],
                            "teacher": r.get("teacher", "beam-api"),
                            "difficulty": r.get("difficulty"),
                            "messages": [
                                {"role": "user", "content": r["prompt"]},
                                {
                                    "role": "assistant",
                                    "reasoning_content": r["reasoning"],
                                    "content": r["content"],
                                },
                            ],
                            "gold": r["gold"],
                            "n_samples": len(rs),
                            "n_correct": len(good),
                            "chars": len(r["reasoning"]) + len(r["content"]),
                            "completion_tokens": r["usage"].get("completion_tokens"),
                        },
                        ensure_ascii=False,
                    )
                    + "\n"
                )

    groups = defaultdict(list)
    for r in recs:
        split = "eval" if r["id"] in eval_ids else "train"
        groups[(r["effort"], split, band(r))].append(r)
        groups[(r["effort"], split, "ALL")].append(r)
    stats = {"generations": len(recs), "sft_examples": n_sft, "groups": {}}
    for (effort, split, b), rs in sorted(groups.items()):
        toks = [r["usage"].get("completion_tokens", 0) for r in rs]
        ok_toks = [r["usage"].get("completion_tokens", 0) for r in rs if r.get("correct")]
        stats["groups"][f"{effort}/{split}/{b}"] = {
            "n": len(rs),
            "accuracy": round(sum(r.get("correct", False) for r in rs) / len(rs), 3),
            "truncated": sum(r.get("finish_reason") == "length" for r in rs),
            "out_tokens_mean": round(statistics.mean(toks)),
            "out_tokens_median": round(statistics.median(toks)),
            "out_tokens_correct_mean": round(statistics.mean(ok_toks)) if ok_toks else None,
            "out_tokens_total": sum(toks),
        }
    (out / "stats.json").write_text(json.dumps(stats, indent=1))
    print(json.dumps(stats, indent=1))


if __name__ == "__main__":
    main()