#!/usr/bin/env python
"""One vLLM process on one GPU: answer a shard of a problem file, then grade every answer.

Writes <out>.jsonl (one graded record per problem), <out>.texts.jsonl.gz (the full answers)
and <out>.done.json (settings, timing, throughput). A shard "i/k" takes problems i, i+k, ...
"""
import argparse
import gzip
import json
import os
import sys
import time

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import grading as G  # noqa: E402
from wsp import protocol as P  # noqa: E402

TYPES = {"temperature": float, "top_p": float, "top_k": int, "presence_penalty": float,
         "repetition_penalty": float, "max_tokens": int}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--n", type=int, default=1)
    ap.add_argument("--shard", default="0/1")
    ap.add_argument("--preset", default="think35", choices=sorted(P.PRESETS))
    ap.add_argument("--set", action="append", default=[], help="override a sampling value, e.g. temperature=1.0")
    ap.add_argument("--gpu-mem", type=float, default=0.92)
    ap.add_argument("--max-num-seqs", type=int, default=256)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0, help="per-request seed as in the paper; -1 for none")
    ap.add_argument("--with-vision", action="store_true", help="also load the vision encoder (not needed)")
    ap.add_argument("--prefix-cache", action="store_true",
                    help="vLLM prefix caching (off: prompts are ~150 tokens, answers ~30k, and on this hybrid "
                         "model it keeps extra recurrent-state copies)")
    a = ap.parse_args()

    cfg = dict(P.PRESETS[a.preset])
    for kv in a.set:
        k, v = kv.split("=", 1)
        cfg[k] = (v.lower() in ("1", "true", "on")) if k == "thinking" else TYPES[k](v)

    i, k = map(int, a.shard.split("/"))
    with open(a.data) as f:
        rows = [json.loads(line) for line in f if line.strip()][i::k]
    if a.limit:
        rows = rows[:a.limit]

    from vllm import LLM, SamplingParams
    t0 = time.time()
    extra = {} if a.with_vision else {"limit_mm_per_prompt": {"image": 0, "video": 0}}
    llm = LLM(model=a.model, dtype="bfloat16", gpu_memory_utilization=a.gpu_mem,
              max_model_len=cfg["max_tokens"] + 2048, seed=0, max_num_seqs=a.max_num_seqs,
              disable_log_stats=False, enable_prefix_caching=a.prefix_cache, **extra)
    tok = llm.get_tokenizer()
    prompts = [P.render(tok, r, cfg["thinking"]) for r in rows]
    sp = SamplingParams(n=a.n, temperature=cfg["temperature"], top_p=cfg["top_p"], top_k=cfg["top_k"],
                        presence_penalty=cfg["presence_penalty"], repetition_penalty=cfg["repetition_penalty"],
                        max_tokens=cfg["max_tokens"], seed=None if a.seed < 0 else a.seed)
    t1 = time.time()
    outs = llm.generate(prompts, sp)
    t2 = time.time()
    bad = sum(1 for o in outs for x in o.outputs if x.finish_reason not in ("stop", "length"))
    if bad:
        raise RuntimeError(f"{bad} generations neither stopped nor reached max_tokens")

    ntok = 0
    tmp = a.out + ".jsonl.tmp"
    with open(tmp, "w") as fr, gzip.open(a.out + ".texts.jsonl.gz", "wt") as ft:
        for r, o in zip(rows, outs):
            samples = []
            for j, s in enumerate(o.outputs):
                text = s.text
                n_t = len(s.token_ids)
                ntok += n_t
                rec = {"tokens": n_t, "finish": s.finish_reason, "cut": s.finish_reason == "length",
                       "closed": "</think>" in text,
                       "verbatim_loop_end": G.verbatim_loop_end(text),
                       "repetitive_end": G.repetitive_end(text),
                       "tail_dup8": round(G.tail_dup8(text), 4)}
                if r.get("kind") == "ifeval":
                    ok, follow = G.ifeval_follow(r, text, cfg["thinking"])
                    rec.update(correct=bool(ok), follow=follow, pred=None)
                else:
                    pred, nbox = G.extract(text)
                    rec.update(pred=pred, n_boxes=nbox, correct=G.is_correct(pred, r["answer"]))
                samples.append(rec)
                ft.write(json.dumps({"id": r["id"], "j": j, "text": text}) + "\n")
            fr.write(json.dumps({"id": r["id"], "source": r.get("source"), "answer": r.get("answer"),
                                 "samples": samples}) + "\n")
    os.replace(tmp, a.out + ".jsonl")
    note = os.path.join(a.model, "edit_note.json")
    done = {"model": a.model, "edit_note": json.load(open(note)) if os.path.exists(note) else None,
            "data": a.data, "shard": a.shard, "n": a.n, "preset": a.preset, "sampling": cfg, "seed": a.seed,
            "prompt_example": prompts[0] if prompts else None,
            "problems": len(rows), "samples": len(rows) * a.n, "tokens": ntok,
            "load_s": round(t1 - t0, 1), "gen_s": round(t2 - t1, 1),
            "tok_per_s": round(ntok / max(t2 - t1, 1e-9)), "grade_s": round(time.time() - t2, 1)}
    with open(a.out + ".done.json", "w") as f:
        json.dump(done, f, indent=1)
    print("DONE", json.dumps(done), flush=True)


if __name__ == "__main__":
    main()