"""Deterministic fixes of visual answers before rendering (no model call), measured on recorded runs.

  importmap  the page loads three.js as ES modules (imports from 'three', 'three/addons/...' or from the unpkg URL,
             whose addons import the bare name 'three') but has no import map: insert the import map of the task
             at the start of <head>.

usage: vis_autofix.py --src results/vis_rep/vis.jsonl --out results/vis_rep_fix/vis.jsonl
Records whose answer is changed get "autofix": [...]; the rest are copied unchanged (exactly paired).
"""
import argparse
import json
import os
import re
import sys

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

IMPORT_MAP = ('<script type="importmap">{"imports": {"three": "https://unpkg.com/three@0.160.0/build/three.module.js", '
              '"three/addons/": "https://unpkg.com/three@0.160.0/examples/jsm/"}}</script>')
USES_THREE_MODULE = re.compile(r"""from\s+['"](three|three/addons/[^'"]+|https://unpkg\.com/three@[^'"]+)['"]""")


def fix_importmap(src):
    if 'type="importmap"' in src or "type='importmap'" in src or not USES_THREE_MODULE.search(src):
        return src
    m = re.search(r"<head[^>]*>", src, re.I)
    if m:
        return src[:m.end()] + "\n" + IMPORT_MAP + src[m.end():]
    m = re.search(r"<script", src, re.I)
    return src[:m.start()] + IMPORT_MAP + "\n" + src[m.start():] if m else src


def autofix(text, kind):
    """Returns (new answer text, list of applied fixes)."""
    if kind != "html" or text is None:
        return text, []
    src = extract(text, kind)
    if not src:
        return text, []
    new = fix_importmap(src)
    if new == src:
        return text, []
    return text.replace(src, new), ["importmap"]


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--tasks", default="data/visual/tasks.jsonl")
    a = ap.parse_args()
    tasks = {json.loads(l)["id"]: json.loads(l) for l in open(a.tasks, encoding="utf-8")}
    os.makedirs(os.path.dirname(a.out), exist_ok=True)
    n = 0
    with open(a.out, "w", encoding="utf-8") as f:
        for l in open(a.src, encoding="utf-8"):
            r = json.loads(l)
            c = r["completion"]
            if "</think>" in c:
                head, fin = c.split("</think>", 1)
                new, fixes = autofix(fin, tasks[r["id"]]["kind"])
                if fixes:
                    r = {**r, "completion": head + "</think>" + new, "autofix": fixes}
                    n += 1
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{n} answers changed -> {a.out}")