#!/usr/bin/env python3
"""Generate creative-writing samples against an OpenAI-compatible endpoint.

Prompts: EQ-Bench creative-writing v3 set (prompts.json, their file, cited
in README). Generations at temperature 0.7/min_p 0.1 per EQ convention —
voice needs room; determinism would strangle it. Budget 6000 tokens per
piece: generous on purpose, every char is data.

Saves one JSONL line per piece: prompt metadata, full response, and the
reasoning trace. Server model ids (local file paths) are replaced with our
run label before saving — no machine info ever lands on disk.

Usage:
  BENCH_LABEL=file-a python3 generate.py --prompts prompts.json \\
      --out runs/file-a.jsonl [--ids 1,2,3] [--budget 6000] [--temp 0.7]
Env:
  BENCH_LABEL   run label recorded on every line (default: file)
  CREATIVE_UPSTREAM  chat-completions base URL
                (default: http://host.docker.internal:8080)
"""
import argparse
import json
import os
import time
import urllib.request

LABEL = os.environ.get("BENCH_LABEL", "file")
UPSTREAM = os.environ.get("CREATIVE_UPSTREAM", "http://host.docker.internal:8080")


def complete(messages, budget, temp, seed):
    body = json.dumps({
        "model": LABEL, "messages": messages, "max_tokens": budget,
        "temperature": temp, "min_p": 0.1, "seed": seed,
    }).encode()
    req = urllib.request.Request(UPSTREAM + "/v1/chat/completions", data=body,
                                 method="POST",
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=1200) as r:
        return json.loads(r.read())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompts", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--ids", default="")
    ap.add_argument("--budget", type=int, default=6000)
    ap.add_argument("--temp", type=float, default=0.7)
    ap.add_argument("--seed", type=int, default=1234)
    a = ap.parse_args()

    prompts = json.load(open(a.prompts))
    want = [w.strip() for w in a.ids.split(",") if w.strip()] or sorted(
        prompts, key=int)
    with open(a.out, "w") as out:
        for pid in want:
            p = prompts[pid]
            # user turn only: some models support no system role
            msgs = [{"role": "user", "content": p["writing_prompt"]}]
            t0 = time.time()
            resp = complete(msgs, a.budget, a.temp, a.seed)
            m = resp["choices"][0]["message"]
            out.write(json.dumps({
                "id": pid, "label": LABEL,
                "category": p.get("category"), "title": p.get("title"),
                "prompt": p["writing_prompt"],
                "text": m.get("content") or "",
                "thinking": m.get("reasoning_content") or "",
                "finish": resp["choices"][0].get("finish_reason"),
                "temp": a.temp, "seed": a.seed, "budget": a.budget,
                "elapsed": round(time.time() - t0, 1),
            }) + "\n")
            out.flush()
            print(f"  {pid} done", flush=True)


if __name__ == "__main__":
    main()