"""Code read-out rescue for the deliberation module.

When the thinking of a code task hits the budget, the module closes it and lets the model answer. On hard problems the
model then writes an editorial-style "Solution Explanation" with pseudo-code and runs out of space before the program
(LCB v6: 37 of 112 budget read-outs had no Python block at all). Rescue: if the answer generated at the budget contains
no program that compiles, generate it again with the code block already opened ("```python" right after </think>).
Other records are copied unchanged, so the comparison with the source run is exactly paired. Tokens of the discarded
answer stay counted.
usage: code_rescue.py --src results/lcb_dm/lcb.jsonl --items data/lcb/lcb6.jsonl --out results/lcb_dmr/lcb.jsonl
"""
import argparse
import json
import os
import sys
import time
import urllib.request
import zlib
from concurrent.futures import ThreadPoolExecutor

from transformers import AutoTokenizer

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import extract_code  # noqa: E402

THINK_END = 124902
PREFILL = "```python\n"
CODE_TASKS = {"code_io", "humaneval", "mbpp", "code_fn", "lcb"}
SAMPLING = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "repeat_penalty": 1.0,
            "samplers": ["top_k", "top_p", "temperature"]}


def runnable(text):
    code = extract_code(text or "")
    if not code:
        return False
    try:
        compile(code, "<answer>", "exec")
        return True
    except Exception:
        return False


def close_block(content):
    text = PREFILL + content
    return text if "```" in content else text.rstrip() + "\n```"


def post(port, body):
    for attempt in range(5):
        try:
            req = urllib.request.Request(f"http://localhost:{port}/completion", data=json.dumps(body).encode(),
                                         headers={"Content-Type": "application/json"})
            return json.loads(urllib.request.urlopen(req, timeout=3600).read())
        except Exception:
            if attempt == 4:
                raise
            time.sleep(5)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--items", nargs="+", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--reserve", type=int, default=4096)
    ap.add_argument("--workers", type=int, default=32)
    ap.add_argument("--port", type=int, default=8080)
    ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
    args = ap.parse_args()
    tok = AutoTokenizer.from_pretrained(args.tokenizer)
    pre_ids = tok(PREFILL, add_special_tokens=False)["input_ids"]
    items = {}
    for path in args.items:
        for l in open(path, encoding="utf-8"):
            r = json.loads(l)
            items[r["id"]] = r

    def process(rec):
        it = items[rec["id"]]
        c = rec["completion"]
        if it["task"] not in CODE_TASKS or rec.get("module") != "fallback" or "</think>" not in c:
            return rec
        thinking, final = c[len("<think>"):].split("</think>", 1)
        if runnable(final):
            return rec
        ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
                                      add_generation_prompt=True, return_dict=False)
        gen = tok(thinking, add_special_tokens=False)["input_ids"]
        seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec.get("sample", 0) + 31
        r = post(args.port, {"prompt": ids + gen + [THINK_END] + pre_ids, "n_predict": args.reserve, "seed": seed,
                             "cache_prompt": False, "return_tokens": True, **SAMPLING})
        text = close_block(r["content"])
        return {**rec, "completion": "<think>" + thinking + "</think>" + text,
                "n_tokens": rec["n_tokens"] + len(pre_ids) + len(r["tokens"]), "module": "fallback_rescue",
                "rescued": runnable(text)}

    recs = [json.loads(l) for l in open(args.src, encoding="utf-8")]
    t0 = time.time()
    with ThreadPoolExecutor(args.workers) as ex:
        out = list(ex.map(process, recs))
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as f:
        for r in out:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    res = [r for r in out if r.get("module") == "fallback_rescue"]
    print(f"{len(out)} records, {len(res)} read-outs without a runnable program regenerated, "
          f"{sum(r['rescued'] for r in res)} now have one, {(time.time() - t0) / 60:.1f} min")