"""Commit-on-consensus (CoC) generation: an inference-time halting mechanism.

The reasoning is generated in chunks (--chunk tokens, KV cache of the slot is reused). After every chunk, while the
model is still thinking, a probe closes the thinking at the last sentence end (the model's own closing pattern:
"...." + </think>) and lets the model write its final answer greedily on a separate slot. When two consecutive probes
give the same extracted answer, the reasoning is committed there: the output becomes reasoning[:cut] + </think> +
that probe's answer. Code tasks are never probed (no cheap answer equivalence).

--no-probe runs the identical chunked generation without probes: the control for measuring CoC's own contribution
(both runs are token-for-token identical until CoC commits).

Output records have the same format as gen.py, plus "coc_cut" (None if not committed).
usage: gen_coc.py --inp data/quick/mini.jsonl --out results/coc/mini.jsonl --max-tokens 6144 [--no-probe]
"""
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("--max-tokens", type=int, default=16384)
ap.add_argument("--chunk", type=int, default=1024)
ap.add_argument("--agree", type=int, default=2, help="consecutive identical probe answers needed to commit")
ap.add_argument("--probe-tokens", type=int, default=400)
ap.add_argument("--no-probe", action="store_true")
ap.add_argument("--record-all", action="store_true",
                help="shadow mode that never stops probing: stores every probe (cut, answer, text) for offline "
                     "simulation of other chunk/agreement settings")
ap.add_argument("--shadow", action="store_true",
                help="on consensus, record the committed answer but keep generating to the end: one run yields both "
                     "outcomes on the same trajectory (exact pairing) and also training data")
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, EOS = 124902, 124900
tok = AutoTokenizer.from_pretrained(args.tokenizer)
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"] for l in open(args.out, encoding="utf-8")}
items = [it for it in items if it["id"] not in done]
print(len(items), "items", 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(f"http://localhost:{args.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(2)


def extract(task, text):
    """answer identity used for the consensus test"""
    if task in ("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(it):
    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"]
        probe = not args.no_probe and not task.startswith("code") and task not in ("humaneval", "mbpp", "ifeval")
        base_seed = args.seed * 1000003 + zlib.crc32(it["id"].encode()) % 1000003 * 7
        gen, stop, last_ans, prev_cut, commit, n_chunk, n_probe = [], None, None, 0, None, 0, 0
        plog = []
        while True:
            n_pred = min(args.chunk, args.max_tokens - len(gen))
            if n_pred <= 0:
                stop = "limit"
                break
            res = post({"prompt": ids + gen, "n_predict": n_pred, "temperature": args.temperature,
                        "top_p": args.top_p, "top_k": args.top_k, "min_p": 0.0, "repeat_penalty": 1.0,
                        "seed": base_seed + 7919 * n_chunk, "cache_prompt": True, "id_slot": w,
                        "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True})
            n_chunk += 1
            gen += res["tokens"]
            if res.get("stop_type") == "eos" or (res["tokens"] and res["tokens"][-1] == EOS):
                stop = "eos"
                break
            if not probe or THINK_END in gen:
                continue
            cut = snap(gen, prev_cut)
            if cut is None:
                continue
            prev_cut = cut
            pr = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.probe_tokens, "temperature": 0.0,
                       "cache_prompt": False, "id_slot": args.workers + w, "return_tokens": True})
            n_probe += 1
            ans = extract(task, pr["content"])
            if args.record_all:
                plog.append({"cut": cut, "ans": ans, "text": pr["content"], "n": len(pr["tokens"])})
                continue
            if ans is not None and ans == last_ans:
                commit = (cut, pr)
                if not args.shadow:
                    break
                probe = False
                continue
            last_ans = ans
        shadow = {}
        if commit and args.shadow:
            cut, pr = commit
            shadow = {"coc_completion": "<think>" + tok.decode(gen[:cut]) + "</think>" + pr["content"],
                      "coc_n_tokens": cut + len(pr["tokens"]), "coc_tokens": gen[:cut] + [THINK_END] + pr["tokens"]}
            commit = None
        if commit:
            cut, pr = commit
            text = tok.decode(gen[:cut]) + "</think>" + pr["content"]
            n_tok, trunc, stop_type = cut + len(pr["tokens"]), False, "commit"
        else:
            if gen and gen[-1] == EOS:
                gen = gen[:-1]
            text = tok.decode(gen)
            n_tok, trunc, stop_type = len(gen), stop == "limit", stop
        rec = {"id": it["id"], "task": task, "sample": 0, "completion": "<think>" + text, "n_tokens": n_tok,
               "truncated": trunc, "stop_type": stop_type, "coc_cut": commit[0] if commit else None,
               "probes": n_probe, "gen_tokens": len(gen), **shadow}
        if args.record_all:
            rec["probe_log"] = plog
        if shadow:
            rec["coc_cut"] = shadow["coc_n_tokens"]
        with lock:
            out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            out_f.flush()
            stats["n"] += 1
            stats["commit"] += bool(commit) or bool(shadow)
            stats["tok"] += len(gen)
            if stats["n"] % 20 == 0:
                el = time.time() - stats["t0"]
                print(f"{stats['n']}/{len(items)}  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, items))
print("done", stats, flush=True)