"""Generate one reply per SENSUS Bench v1 evaluation case from an OpenAI-compatible endpoint.

    python3 sample.py --model accounts/<you>/models/<your-model> --out replies.jsonl

Settings match the recorded evaluation: temperature 0, at most 900 output tokens, thinking off.
The endpoint defaults to Fireworks; the key is read from FIREWORKS_API_KEY (or --key-env).
Licensed under the MIT License (LICENSE-CODE).
"""
import argparse
import json
import os
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor


def call(url, key, body):
    req = urllib.request.Request(url, data=json.dumps(body).encode(),
                                 headers={"content-type": "application/json", "authorization": "Bearer " + key})
    with urllib.request.urlopen(req, timeout=300) as r:
        return json.loads(r.read())


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--out", default="replies.jsonl")
    ap.add_argument("--eval", default="eval.jsonl")
    ap.add_argument("--url", default="https://api.fireworks.ai/inference/v1/chat/completions")
    ap.add_argument("--key-env", default="FIREWORKS_API_KEY")
    ap.add_argument("--max-tokens", type=int, default=900)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--keep-thinking", action="store_true", help="do not send reasoning_effort=none")
    a = ap.parse_args()
    key = os.environ[a.key_env]
    families = [json.loads(l) for l in open(a.eval, encoding="utf-8") if l.strip()]
    done = set()
    if os.path.exists(a.out):
        done = {json.loads(l)["id"] for l in open(a.out, encoding="utf-8") if l.strip()}

    def one(f):
        body = {"model": a.model, "messages": f["messages"], "temperature": 0, "max_tokens": a.max_tokens}
        if not a.keep_thinking:
            body["reasoning_effort"] = "none"
        for attempt in range(3):
            try:
                d = call(a.url, key, body)
                return {"id": f["id"], "reply": d["choices"][0]["message"]["content"],
                        "finish_reason": d["choices"][0].get("finish_reason")}
            except Exception as err:  # noqa: BLE001
                last = str(err)[:200]
                time.sleep(3 * (attempt + 1))
        return {"id": f["id"], "error": last}

    todo = [f for f in families if f["id"] not in done]
    with open(a.out, "a", encoding="utf-8") as out, ThreadPoolExecutor(a.workers) as pool:
        for row in pool.map(one, todo):
            out.write(json.dumps(row, ensure_ascii=False) + "\n")
            print(row["id"], "error" if "error" in row else "ok", flush=True)


if __name__ == "__main__":
    main()