"""Commit-on-consensus (CoC), streaming version: an inference-time halting mechanism.

The main reasoning is ONE streamed request with exactly the sampling settings and seed of gen.py, so without a
commit the output is the ordinary generation. Every --chunk generated tokens (while the model is still thinking) a
probe runs on a separate slot: the reasoning so far is closed at the last sentence end with the model's own closing
pattern ("...." + </think>) and the model writes its final answer greedily. When --agree consecutive probes give the
same extracted answer, the reasoning is committed: the stream is dropped and the output becomes
reasoning[:cut] + </think> + that probe's answer. Code / instruction-following tasks are never probed.

--shadow: on consensus keep generating to the end and store the committed answer next to the full one, so one run
gives both outcomes on the same trajectory (the full arm is the ordinary baseline generation).
--record-all: never stop probing (with --shadow), store every probe for offline simulation of other settings.

usage: gen_coc2.py --inp FILE.. --out OUT.jsonl [--max-tokens 16384] [--chunk 1024] [--agree 2] [--shadow]
"""
import argparse
import json
import os
import queue
import re
import sys
import threading
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 last_boxed  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--inp", nargs="+", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--n", type=int, default=1)
ap.add_argument("--max-tokens", type=int, default=16384)
ap.add_argument("--chunk", type=int, default=1024)
ap.add_argument("--agree", type=int, default=2)
ap.add_argument("--probe-tokens", type=int, default=400)
ap.add_argument("--shadow", action="store_true")
ap.add_argument("--free-probes", action="store_true", help="let the model write its answer freely (slow)")
ap.add_argument("--fallback", action="store_true",
                help="if the reasoning hits the token limit, close it and answer instead of being cut off "
                     "(short-answer tasks: forced short answer within the last 64 tokens; other tasks: thinking is "
                     "closed --free-reserve tokens before the limit and the answer is generated normally)")
ap.add_argument("--free-reserve", type=int, default=1536)
ap.add_argument("--commit-tasks", nargs="*", default=None,
                help="tasks where early commitment is allowed (default: all probed tasks; the released module uses "
                     "mmlupro bbh, chosen on dev: numeric math answers get revised late and commits hurt there)")
ap.add_argument("--record-all", action="store_true")
ap.add_argument("--workers", type=int, default=24, help="main slots 0..W-1, probe slots W..2W-1 (server -np 2W)")
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("--seed", type=int, default=0)
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
args = ap.parse_args()

THINK_END = 124902
MATH_SUFFIX = "The final answer is $\\boxed{"  # math probes: short forced answer (a free answer re-derives everything)
URL = f"http://localhost:{args.port}/completion"
tok = AutoTokenizer.from_pretrained(args.tokenizer)
MATH_SUFFIX_IDS = tok(MATH_SUFFIX, add_special_tokens=False)["input_ids"]
# short forced-answer probes per task: (suffix text, max new tokens, stop strings, closing text)
SHORT = {"math": (MATH_SUFFIX, 48, ["}$", "}\n"], "}$"), "mmlupro": ("The answer is (", 3, [")"], ")."),
         "bbh": ("The answer is ", 24, ["\n", ".\n"], ".")}
SHORT["math500"] = SHORT["aime"] = SHORT["gsm8k"] = SHORT["math"]
SHORT_IDS = {k: tok(v[0], add_special_tokens=False)["input_ids"] for k, v in SHORT.items()}
ENDS = set()
for i_ in range(len(tok)):
    t_ = tok.convert_ids_to_tokens(i_)
    if t_ is not None and (t_.endswith(".") or t_.endswith("\u010a")):
        ENDS.add(i_)

items = []
for path in args.inp:
    items += [json.loads(l) for l in open(path, encoding="utf-8")]
done = set()
if os.path.exists(args.out):
    done = {(json.loads(l)["id"], json.loads(l)["sample"]) for l in open(args.out, encoding="utf-8")}
jobs = [(it, s) for it in items for s in range(args.n) if (it["id"], s) not in done]
print(len(jobs), "jobs", flush=True)
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
out_f = open(args.out, "a", encoding="utf-8")
lock = threading.Lock()
slots = queue.Queue()
for w in range(args.workers):
    slots.put(w)
stats = {"n": 0, "commit": 0, "t0": time.time(), "tok": 0}


def post(body):
    for attempt in range(5):
        try:
            req = urllib.request.Request(URL, 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(2)


def extract(task, text):
    a = _extract(task, text)
    if a is None or "..." in a or a in ("?", "") or len(a) > 80:  # placeholders never count as an answer
        return None
    return a


def _extract(task, text):
    a = _extract(task, text)
    if a is None or "..." in a or a in ("?", "") or len(a) > 80:  # placeholders never count as an answer
        return None
    return a


def _extract(task, text):
    if task == "mmlupro":
        m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", text) or re.findall(r"\\boxed\{\(?([A-J])\)?\}", text)
        return m[-1] if m else None
    if task == "bbh":
        m = re.findall(r"answer is:?\s*(.+)", text, re.I)
        a = m[-1] if m else (last_boxed(text) or "")
        a = re.sub(r"\\boxed\{(.*)\}", r"\1", a).strip().strip("*$ .").lower()
        return a or None
    b = last_boxed(text)
    if b is None and task == "gsm8k":
        nums = re.findall(r"-?\d[\d,]*\.?\d*", text)
        b = nums[-1] if nums else None
    return re.sub(r"\s+", "", b) if b else None


def snap(gen, lo):
    for j in range(len(gen), lo, -1):
        if gen[j - 1] in ENDS:
            return j
    return None


def run(job):
    it, s = job
    w = slots.get()
    try:
        ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
                                      add_generation_prompt=True, return_dict=False)
        task = it["task"]
        probing = not task.startswith("code") and task not in ("humaneval", "mbpp", "ifeval")
        short = task in SHORT and not args.free_probes
        if args.commit_tasks is not None and task not in args.commit_tasks:
            probing = False  # no early commitment for this task (the budget fallback still applies)
        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": True,
                "stream": True, "id_slot": w}
        req = urllib.request.Request(URL, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
        resp = urllib.request.urlopen(req, timeout=3600)
        gen, final, next_probe, prev_cut = [], None, args.chunk, 0
        probe_q, probe_res = queue.Queue(), []  # probes run in their own thread, in order

        def prober():
            while True:
                job_ = probe_q.get()
                if job_ is None:
                    return
                cut_, snap_ = job_
                if short:
                    suf, n_new, stops, close = SHORT[task]
                    pr = post({"prompt": ids + snap_ + [THINK_END] + SHORT_IDS[task], "n_predict": n_new,
                               "temperature": 0.0, "cache_prompt": False, "id_slot": args.workers + w,
                               "return_tokens": True, "stop": stops})
                    c_ = pr["content"]
                    if close == "}$":  # keep a balanced \boxed{...}
                        depth = 1
                        for j_, ch in enumerate(c_):
                            depth += (ch == "{") - (ch == "}")
                            if depth == 0:
                                c_ = c_[:j_]
                                break
                    c_ = c_.strip().rstrip(".").strip() if close == "." else c_.strip()
                    text = suf + c_ + close
                    toks = SHORT_IDS[task] + pr["tokens"]
                else:
                    pr = post({"prompt": ids + snap_ + [THINK_END], "n_predict": args.probe_tokens, "temperature": 0.0,
                               "cache_prompt": False, "id_slot": args.workers + w, "return_tokens": True})
                    text, toks = pr["content"], pr["tokens"]
                probe_res.append({"cut": cut_, "ans": extract(task, text), "text": text, "toks": toks})

        def fallback_answer(gen_):
            """answer within the budget when the reasoning was cut off by the limit"""
            if short:
                cut_ = snap(gen_[: args.max_tokens - 64], 0)
                if cut_ is None:
                    return None
                suf, n_new, stops, close = SHORT[task]
                pr = post({"prompt": ids + gen_[:cut_] + [THINK_END] + SHORT_IDS[task], "n_predict": n_new,
                           "temperature": 0.0, "cache_prompt": False, "id_slot": args.workers + w,
                           "return_tokens": True, "stop": stops})
                c_ = pr["content"]
                if close == "}$":
                    depth = 1
                    for j_, ch in enumerate(c_):
                        depth += (ch == "{") - (ch == "}")
                        if depth == 0:
                            c_ = c_[:j_]
                            break
                c_ = c_.strip().rstrip(".").strip() if close == "." else c_.strip()
                text_ = suf + c_ + close
                return {"cut": cut_, "text": text_, "n": len(SHORT_IDS[task]) + len(pr["tokens"])}
            cut_ = snap(gen_[: args.max_tokens - args.free_reserve], 0)
            if cut_ is None:
                return None
            pr = post({"prompt": ids + gen_[:cut_] + [THINK_END], "n_predict": args.free_reserve,
                       "temperature": args.temperature, "top_p": args.top_p, "top_k": args.top_k, "min_p": 0.0,
                       "seed": body["seed"] + 17, "cache_prompt": False, "id_slot": args.workers + w,
                       "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True})
            return {"cut": cut_, "text": pr["content"], "n": len(pr["tokens"])}

        pth = threading.Thread(target=prober, daemon=True)
        pth.start()
        commit, checked, closed = None, 0, False
        while True:
            line = resp.readline()
            if not line:
                break
            line = line.decode("utf-8").strip()
            if not line.startswith("data:"):
                continue
            ev = json.loads(line[5:])
            gen += ev.get("tokens", [])
            if ev.get("stop"):
                final = ev
                break
            if THINK_END in ev.get("tokens", []):
                closed = True
            if probing and not closed and (commit is None or args.record_all) and len(gen) >= next_probe:
                next_probe += args.chunk
                cut = snap(gen, prev_cut)
                if cut:
                    prev_cut = cut
                    probe_q.put((cut, list(gen[:cut])))
            # consensus check over finished probes
            while commit is None and checked < len(probe_res):
                checked += 1
                tail = probe_res[:checked][-args.agree:]
                if len(tail) == args.agree and tail[0]["ans"] is not None and all(p["ans"] == tail[0]["ans"] for p in tail):
                    commit = tail[-1]
            if commit is not None and not args.shadow:
                resp.close()  # drop the stream: the server stops generating for this slot
                break
        probe_q.put(None)
        if commit is None or args.shadow:
            pth.join(timeout=900)
            # probes that finished after the stream ended can still produce a (shadow) commit
            while commit is None and checked < len(probe_res):
                checked += 1
                tail = probe_res[:checked][-args.agree:]
                if len(tail) == args.agree and tail[0]["ans"] is not None and all(p["ans"] == tail[0]["ans"] for p in tail):
                    # only valid if that probe's cut was reached before thinking ended: always true by construction
                    commit = tail[-1]
        if gen and gen[-1] == 124900:  # the stream also returns the end-of-turn token
            gen = gen[:-1]
        full_text = "<think>" + tok.decode(gen)
        rec = {"id": it["id"], "task": task, "sample": s}
        truncated = bool(final) and final.get("stop_type") == "limit"
        if final is None and commit is None:
            truncated = len(gen) >= args.max_tokens
        fb = fallback_answer(gen) if (args.fallback and commit is None and truncated) else None
        coc = None
        if commit is not None:
            coc = {"completion": "<think>" + tok.decode(gen[:commit["cut"]]) + "</think>" + commit["text"],
                   "n_tokens": commit["cut"] + len(commit["toks"]), "cut": commit["cut"]}
        elif fb is not None:
            coc = {"completion": "<think>" + tok.decode(gen[:fb["cut"]]) + "</think>" + fb["text"],
                   "n_tokens": fb["cut"] + fb["n"], "cut": fb["cut"]}
            rec["fallback"] = True
        if args.shadow:
            rec.update(completion=full_text, n_tokens=len(gen),
                       truncated=bool(final) and final.get("stop_type") == "limit",
                       stop_type=final.get("stop_type") if final else None)
            if coc:
                rec.update(coc_completion=coc["completion"], coc_n_tokens=coc["n_tokens"], coc_cut=coc["cut"])
            if args.record_all:
                rec["probe_log"] = [{k: p[k] for k in ("cut", "ans", "text")} | {"n": len(p["toks"])}
                                    for p in probe_res]
        elif coc:
            rec.update(completion=coc["completion"], n_tokens=coc["n_tokens"], truncated=False, stop_type="commit",
                       coc_cut=coc["cut"])
        else:
            rec.update(completion=full_text, n_tokens=len(gen),
                       truncated=bool(final) and final.get("stop_type") == "limit",
                       stop_type=final.get("stop_type") if final else None, coc_cut=None)
        rec["probes"] = len(probe_res)
        with lock:
            out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            out_f.flush()
            stats["n"] += 1
            stats["commit"] += commit is not None
            stats["tok"] += len(gen)
            if stats["n"] % 20 == 0:
                el = time.time() - stats["t0"]
                print(f"{stats['n']}/{len(jobs)}  commits {stats['commit']}  {stats['tok'] / el:.0f} gen tok/s  "
                      f"{el / 60:.1f} min", flush=True)
    finally:
        slots.put(w)


with ThreadPoolExecutor(args.workers) as ex:
    list(ex.map(run, jobs))
print("done", stats, flush=True)