"""Generates completions for eval/train JSONL files through a running llama-server.

Prompts are rendered with the model's own HF chat template (thinking enabled) and sent as token ids, so the
server sees exactly the same input for every model. Output is appended to a JSONL file and the run resumes
where it stopped.

usage: gen.py --inp data/eval/gsm8k.jsonl --out results/orig/gsm8k.jsonl [--n 1] [--max-tokens 16384]
"""
import argparse
import json
import os
import threading
import time
import zlib
import urllib.request
from concurrent.futures import ThreadPoolExecutor

from transformers import AutoTokenizer

ap = argparse.ArgumentParser()
ap.add_argument("--inp", nargs="+", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--n", type=int, default=1, help="samples per prompt")
ap.add_argument("--max-tokens", type=int, default=16384)
ap.add_argument("--temperature", type=float, default=0.6)
ap.add_argument("--top-p", type=float, default=0.95)
ap.add_argument("--top-k", type=int, default=50)
ap.add_argument("--save-tokens", action="store_true", help="store generated token ids (for RL training)")
ap.add_argument("--workers", type=int, default=32)
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--system-file", default="", help="optional system prompt (text file); json {task: prompt} also works")
args = ap.parse_args()

tok = AutoTokenizer.from_pretrained(args.tokenizer)
SYSTEM = ""
if args.system_file:
    SYSTEM = open(args.system_file, encoding="utf-8").read()
    if args.system_file.endswith(".json"):
        SYSTEM = json.loads(SYSTEM)
items = []
for path in args.inp:
    with open(path, encoding="utf-8") as f:
        rows = [json.loads(l) for l in f]
    items += rows[: args.limit] if args.limit else rows

done = set()
if os.path.exists(args.out):
    with open(args.out, encoding="utf-8") as f:
        for l in f:
            r = json.loads(l)
            done.add((r["id"], r["sample"]))
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)

jobs = [(it, s) for it in items for s in range(args.n) if (it["id"], s) not in done]
print(f"{len(jobs)} jobs ({len(done)} already done)", flush=True)

lock = threading.Lock()
out_f = open(args.out, "a", encoding="utf-8")
stats = {"n": 0, "tok": 0, "t0": time.time()}


def run(job):
    it, s = job
    sysp = SYSTEM.get(it["task"], SYSTEM.get("*", "")) if isinstance(SYSTEM, dict) else SYSTEM
    ids = tok.apply_chat_template(([{"role": "system", "content": sysp}] if sysp else []) +
                                  [{"role": "user", "content": it["prompt"]}], tokenize=True,
                                  add_generation_prompt=True, return_dict=False)
    if isinstance(ids, dict):
        ids = ids["input_ids"]
    body = {
        "prompt": ids, "n_predict": args.max_tokens, "temperature": args.temperature, "top_p": args.top_p,
        "top_k": args.top_k, "min_p": 0.0, "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0,
        "dry_multiplier": 0.0, "xtc_probability": 0.0, "typical_p": 1.0,
        "seed": args.seed * 1000003 + zlib.crc32(it["id"].encode()) % 1000003 * 7 + s, "cache_prompt": False,
        "samplers": ["top_k", "top_p", "temperature"], "return_tokens": args.save_tokens,
    }
    for attempt in range(5):
        try:
            req = urllib.request.Request(f"http://localhost:{args.port}/completion", data=json.dumps(body).encode(),
                                         headers={"Content-Type": "application/json"})
            res = json.loads(urllib.request.urlopen(req, timeout=3600).read())
            break
        except Exception as e:  # server busy / transient
            if attempt == 4:
                raise
            time.sleep(5)
    rec = {"id": it["id"], "task": it["task"], "sample": s, "completion": "<think>" + res["content"],
           "n_tokens": res["tokens_predicted"], "truncated": res.get("stop_type") == "limit",
           "prompt_tokens": len(ids), "stop_type": res.get("stop_type")}
    if args.save_tokens:
        rec["tokens"] = res["tokens"]
    with lock:
        out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
        out_f.flush()
        stats["n"] += 1
        stats["tok"] += rec["n_tokens"]
        if stats["n"] % 50 == 0 or stats["n"] == len(jobs):
            dt = time.time() - stats["t0"]
            print(f"{stats['n']}/{len(jobs)}  {stats['tok'] / dt:.0f} tok/s  avg_len {stats['tok'] / stats['n']:.0f}  "
                  f"elapsed {dt / 60:.1f} min", flush=True)


with ThreadPoolExecutor(args.workers) as ex:
    list(ex.map(run, jobs))
out_f.close()