"""Generate fresh Swift responses; resumable by prompt hash and model identity."""
import os
os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import argparse
import json
import time
from common import RUN, BASELINE, records, read_json, write_json, sha256, stamp


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=str(BASELINE))
    ap.add_argument("--limit", type=int)
    ap.add_argument("--chunk", type=int, default=128)
    args = ap.parse_args()
    from transformers import AutoTokenizer
    from vllm import LLM
    data = RUN / "calibration"
    prompts = records(data / "prompts.jsonl")
    if args.limit:
        prompts = prompts[:args.limit]
    manifest = {"model": args.model, "config_sha256": sha256(os.path.join(args.model, "config.json")),
                "prompt_manifest_sha256": sha256(data / "manifest.json"),
                "thinking_cap": 2048, "nonthinking_cap": 1024, "seed": 15027,
                "temperature": 1.0, "top_p": .95, "top_k": 20,
                "max_model_len": 8192, "kv_cache_dtype": "fp8", "int8_activations": False}
    mp = data / "generation-manifest.json"
    if mp.exists():
        assert read_json(mp) == manifest, "Generation settings changed; use a separate run directory"
    else:
        write_json(mp, manifest)
    output = data / "gen.jsonl"
    done = {r["id"]: r for r in records(output)} if output.exists() else {}
    for p in prompts:
        if p["id"] in done:
            assert p["prompt_sha256"] == done[p["id"]]["prompt_sha256"]
    todo = [p for p in prompts if p["id"] not in done]
    if not todo:
        print("All requested examples already generated")
        return
    tok = AutoTokenizer.from_pretrained(args.model)
    llm = LLM(model=args.model, gpu_memory_utilization=.93, max_model_len=8192,
              max_num_seqs=32, max_num_batched_tokens=2048, kv_cache_dtype="fp8",
              mamba_ssm_cache_dtype="float16", language_model_only=True,
              enable_prefix_caching=False,
              compilation_config={"max_cudagraph_capture_size": 32, "custom_ops": ["+rms_norm", "+silu_and_mul"]})
    base = llm.get_default_sampling_params()
    t0 = time.monotonic()
    token_count = 0
    with output.open("a") as f:
        for start in range(0, len(todo), args.chunk):
            batch, inputs, params = [], [], []
            for p in todo[start:start + args.chunk]:
                kwargs = {"tools": p["tools"]} if p.get("tools") else {}
                text = tok.apply_chat_template(p["messages"], tokenize=False, add_generation_prompt=True,
                                               enable_thinking=p["think"], **kwargs)
                ids = tok.encode(text, add_special_tokens=False)
                if len(ids) > 6000:
                    raise ValueError(f"Prompt {p['id']} exceeds calibration context: {len(ids)}")
                sp = base.clone()
                sp.max_tokens = 2048 if p["think"] else 1024
                sp.temperature, sp.top_p, sp.top_k = 1.0, .95, 20
                sp.seed = 15027 + p["id"]
                batch.append((p, ids))
                inputs.append({"prompt_token_ids": ids})
                params.append(sp)
            results = llm.generate(inputs, params, use_tqdm=False)
            for (p, ids), result in zip(batch, results):
                c = result.outputs[0]
                r = {k: p[k] for k in ["id", "src", "think", "split", "prompt_sha256"]}
                r.update(prompt_ids=ids, output_ids=list(c.token_ids), finish=c.finish_reason)
                f.write(json.dumps(r) + "\n")
                token_count += len(c.token_ids)
            f.flush()
            elapsed = time.monotonic() - t0
            print(f"{start+len(batch)}/{len(todo)} examples; {token_count} output tokens; {token_count/elapsed:.1f} tok/s; {elapsed/60:.1f} min", flush=True)
    rows = records(output)
    write_json(data / "generation-summary.json", {
        "completed_at": stamp(), "examples": len(rows), "output_tokens": sum(len(r["output_ids"]) for r in rows),
        "truncated": sum(r["finish"] == "length" for r in rows),
        "note": "Calibration caps are not evaluation budgets; truncation rates are disclosed.",
    })


if __name__ == "__main__":
    main()