"""Apply the deliberation module to already generated reasoning (exactly what it would do online).

The module never changes the reasoning before it acts, and its probes are greedy (deterministic given the prefix),
so running it over recorded trajectories reproduces its online behaviour on those same trajectories. This gives a
fully paired comparison with the recorded baseline.

Module (settings chosen on dev data only):
  * every --every tokens of thinking: close at the last sentence end and read a short forced answer
  * commit when --agree successive answers are identical
  * if the reasoning hit the token limit: answer within the budget instead of being cut off
    (short-answer tasks: forced short answer at the limit; code / instruction tasks: thinking closed
     --free-reserve tokens before the limit, answer generated with the normal sampling settings)
Output: results/<out_tag>/*.jsonl in gen.py format (+ "module": commit/fallback/none) -> grade.py, compare.py

usage: posthoc_coc.py --src results/orig/main.jsonl results/orig/aime.jsonl --out-tag orig_dm
"""
import argparse
import json
import os
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, load_suite  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--src", nargs="+", required=True)
ap.add_argument("--out-tag", required=True)
ap.add_argument("--max-tokens", type=int, default=16384)
ap.add_argument("--every", type=int, default=2048)
ap.add_argument("--agree", type=int, default=3)
ap.add_argument("--free-reserve", type=int, default=1536)
ap.add_argument("--commit-tasks", nargs="*", default=["mmlupro", "bbh"],
                help="tasks where early commitment is allowed (dev: it hurts numeric math at a 16k budget)")
ap.add_argument("--only-tasks", nargs="*", default=None, help="process only these tasks (others copied unchanged)")
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")
ap.add_argument("--items", default=None, help="glob of item files (default: the eval suite)")
args = ap.parse_args()

THINK_END = 124902
URL = f"http://localhost:{args.port}/completion"
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_)
MATH = ("The final answer is $\\boxed{", 48, ["}$", "}\n"], "}$")
SHORT = {"math500": MATH, "aime": MATH, "gsm8k": MATH, "math": MATH,
         "mmlupro": ("The answer is (", 3, [")"], ")."), "bbh": ("The answer is ", 24, ["\n", ".\n"], ".")}
SHORT_IDS = {k: tok(v[0], add_special_tokens=False)["input_ids"] for k, v in SHORT.items()}
items = load_suite(args.items) if args.items else load_suite()
lock = threading.Lock()
stats = {"n": 0, "commit": 0, "fallback": 0, "t0": time.time()}


def post(body):
    for attempt in range(6):
        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 urllib.error.HTTPError as e:
            if e.code == 400:
                return None
            time.sleep(3)
        except Exception:
            time.sleep(3)
    return None


def extract(task, text):
    if task == "mmlupro":
        m = re.findall(r"answer is:?\s*\(?\**\s*([A-J])\b", text)
        a = m[-1] if m else None
    elif 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() or None
    else:
        b = last_boxed(text)
        a = re.sub(r"\s+", "", b) if b else None
    if a is None or "..." in a or a in ("?", "") or len(a) > 80:
        return None
    return a


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


def short_probe(ids, gen, cut, task):
    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, "return_tokens": True, "stop": stops})
    if pr is None:
        return None, None, 0
    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 text, extract(task, text), len(SHORT_IDS[task]) + len(pr["tokens"])


def run(rec):
    it = items[rec["id"]]
    task = it["task"]
    ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
                                  add_generation_prompt=True, return_dict=False)
    body = rec["completion"][len("<think>"):]
    thinking = body.split("</think>", 1)[0]
    gen = tok(thinking, add_special_tokens=False)["input_ids"]
    out = dict(rec)
    out["module"] = "none"
    if args.only_tasks is not None and task not in args.only_tasks:
        return out
    if task in SHORT and task in args.commit_tasks:
        prev, answers = 0, []
        for k in range(args.every, len(gen), args.every):
            cut = snap(gen, k, prev)
            if cut is None:
                continue
            prev = cut
            text, ans, n = short_probe(ids, gen, cut, task)
            answers.append(ans)
            tail = answers[-args.agree:]
            if len(tail) == args.agree and tail[0] is not None and all(a == tail[0] for a in tail):
                out.update(completion="<think>" + tok.decode(gen[:cut]) + "</think>" + text, n_tokens=cut + n,
                           truncated=False, module="commit", module_cut=cut)
                break
    if out["module"] == "none" and rec.get("truncated"):
        if task in SHORT:
            cut = snap(gen, args.max_tokens - 64)
            if cut:
                text, ans, n = short_probe(ids, gen, cut, task)
                if text is not None:
                    out.update(completion="<think>" + tok.decode(gen[:cut]) + "</think>" + text, n_tokens=cut + n,
                               truncated=False, module="fallback", module_cut=cut)
        else:
            cut = snap(gen, args.max_tokens - args.free_reserve)
            if cut:
                seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec.get("sample", 0) + 17
                pr = post({"prompt": ids + gen[:cut] + [THINK_END], "n_predict": args.free_reserve,
                           "temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "seed": seed,
                           "cache_prompt": False, "samplers": ["top_k", "top_p", "temperature"], "return_tokens": True})
                if pr is not None:
                    out.update(completion="<think>" + tok.decode(gen[:cut]) + "</think>" + pr["content"],
                               n_tokens=cut + len(pr["tokens"]), truncated=pr.get("stop_type") == "limit",
                               module="fallback", module_cut=cut)
    with lock:
        stats["n"] += 1
        stats["commit"] += out["module"] == "commit"
        stats["fallback"] += out["module"] == "fallback"
        if stats["n"] % 250 == 0:
            print(f"{stats['n']} done  commits {stats['commit']}  fallbacks {stats['fallback']}  "
                  f"{(time.time() - stats['t0']) / 60:.1f} min", flush=True)
    return out


for src in args.src:
    recs = [json.loads(l) for l in open(src, encoding="utf-8")]
    dst = os.path.join("results", args.out_tag, os.path.basename(src))
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    done = set()
    if os.path.exists(dst):
        done = {(json.loads(l)["id"], json.loads(l)["sample"]) for l in open(dst, encoding="utf-8")}
    todo = [r for r in recs if (r["id"], r["sample"]) not in done]
    print(src, len(todo), "records to process", flush=True)
    with open(dst, "a", encoding="utf-8") as f, ThreadPoolExecutor(args.workers) as ex:
        for out in ex.map(run, todo):
            f.write(json.dumps(out, ensure_ascii=False) + "\n")
            f.flush()
print("done", stats, flush=True)