#!/usr/bin/env python3
"""Generate predictions from an OpenAI-compatible endpoint. The ONLY script here that needs a model.

Every published number reproduces from the cached generations without ever running this — see
harness/README.md. This exists so the cache can be regenerated, or a different model evaluated.

    PYBYTECODE_ENDPOINT=http://localhost:1234/v1 PYBYTECODE_MODEL=my-model \\
      ./generate.py --bench ../benchmarks/csn-3.12-licensed/bench.jsonl --out gen.jsonl

Resumable: an existing --out is read first and completed rows are skipped, so an interrupted run
costs nothing. Output is streamed and flushed per row, never accumulated in memory.
"""
from __future__ import annotations

import argparse
import json
import sys
import urllib.error
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import load_bench  # noqa: E402
from config import API_KEY, ENDPOINT, MODEL  # noqa: E402

INSTRUCTION = (
    "Decompile this Python 3.12 bytecode disassembly back into the original Python source code. "
    "Output only the source code."
)


def complete(prompt: str, temperature: float, max_tokens: int, seed: int | None) -> str:
    body = {
        "model": MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    if seed is not None:
        body["seed"] = seed
    req = urllib.request.Request(
        f"{ENDPOINT.rstrip('/')}/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"},
    )
    with urllib.request.urlopen(req, timeout=600) as r:
        return json.loads(r.read())["choices"][0]["message"]["content"]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--samples", type=int, default=1, help=">1 emits an `s` field per sample")
    ap.add_argument("--max-tokens", type=int, default=2048)
    ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()

    bench, _ = load_bench(a.bench)
    out_path = Path(a.out)
    done: set[tuple[int, int]] = set()
    if out_path.exists():
        for line in out_path.read_text().splitlines():
            if line.strip():
                r = json.loads(line)
                done.add((r["i"], r.get("s", 0)))
        print(f"resuming: {len(done)} generations already present", file=sys.stderr)

    todo = sorted(bench)[: a.limit or None]
    print(f"endpoint={ENDPOINT} model={MODEL} rows={len(todo)} samples={a.samples}", file=sys.stderr)

    with out_path.open("a") as f:
        for i in todo:
            for s in range(a.samples):
                if (i, s) in done:
                    continue
                prompt = f"{INSTRUCTION}\n\n{bench[i]['input']}"
                try:
                    got = complete(prompt, a.temperature, a.max_tokens, None if a.samples == 1 else s)
                except (urllib.error.URLError, OSError) as e:
                    raise SystemExit(
                        f"cannot reach {ENDPOINT}: {e}\n"
                        "Set PYBYTECODE_ENDPOINT to your OpenAI-compatible server, or use the "
                        "cached generations and skip this script entirely."
                    )
                rec = {"i": i, "got": got}
                if a.samples > 1:
                    rec["s"] = s
                f.write(json.dumps(rec) + "\n")
                f.flush()
            if i % 25 == 0:
                print(f"  {i}/{len(todo)}", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()