"""Turns an answer written with a toolkit (kits/) into one self-contained file.

  three  : `import {...} from 'kit3d'` -> reads window.K3; the kit3d module is inserted before the page's module
           script and the import map gets 'three' and 'three/addons/' if they are missing;
  canvas : the kit2d script (global Kit) is inserted before the page's first script;
  html   : <link href="ui.css"> / <script src="ui.js"> are replaced by the inline kit (added if missing);
  svg    : the answer is a Python program using svgkit; it is run and the SVG it prints becomes the answer.
finalize(answer_text, cat) -> (text with one ```svg / ```html block, problem or None, source used for rendering)
"""
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import xml.dom.minidom

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

KITS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "kits")


def _read(name):
    return open(os.path.join(KITS, name), encoding="utf-8").read()


KIT3D, KIT2D, UICSS, UIJS = (_read(n) for n in ("kit3d.js", "kit2d.js", "ui.css", "ui.js"))
IMPORT_KIT3D = re.compile(r"""import\s+([^;'"]+?)\s+from\s+['"](?:\./)?kit3d(?:\.js)?['"]\s*;?""")
THREE_URL = "https://unpkg.com/three@0.160.0/build/three.module.js"
ADDONS_URL = "https://unpkg.com/three@0.160.0/examples/jsm/"


def _destructure(clause):
    """'{ a, b as c }' | '* as K' | 'K' | 'K, { a }' -> statements that read window.K3."""
    clause, out = clause.strip(), []
    m = re.fullmatch(r"\*\s+as\s+(\w+)", clause)
    if m:
        return f"const {m.group(1)} = window.K3;"
    d = re.match(r"(\w+)\s*(,|$)", clause)
    if d:
        out.append(f"const {d.group(1)} = window.K3;")
    b = re.search(r"\{(.*?)\}", clause, re.S)
    if b:
        names = []
        for part in b.group(1).split(","):
            part = part.strip()
            if part:
                a = re.fullmatch(r"(\w+)\s+as\s+(\w+)", part)
                names.append(f"{a.group(1)}: {a.group(2)}" if a else part)
        out.append("const { " + ", ".join(names) + " } = window.K3;")
    return " ".join(out)


def inline_three(src):
    if not IMPORT_KIT3D.search(src):
        return fix_importmap(src)
    src = IMPORT_KIT3D.sub(lambda m: _destructure(m.group(1)), src)
    src = fix_importmap(src)
    m = re.search(r"""<script\s+type=["']importmap["']\s*>(.*?)</script>""", src, re.S)
    if m:
        try:
            j = json.loads(m.group(1))
            imp = j.setdefault("imports", {})
            if "three" not in imp or "three/addons/" not in imp:
                imp.setdefault("three", THREE_URL)
                imp.setdefault("three/addons/", ADDONS_URL)
                src = src[:m.start(1)] + json.dumps(j) + src[m.end(1):]
        except ValueError:
            pass
    kit = '<script type="module">\n' + KIT3D + "\n</script>\n"
    m = re.search(r"""<script\s+type=["']module["']""", src)
    if not m:
        m = re.search(r"</body>", src, re.I)
    return src[:m.start()] + kit + src[m.start():] if m else src + kit


def inline_canvas(src):
    kit = "<script>\n" + KIT2D + "\n</script>\n"
    tag = re.search(r"""<script[^>]*\bsrc=["'](?:\./)?kit2d(?:\.js)?["'][^>]*>\s*</script>""", src, re.I)
    if tag:
        return src[:tag.start()] + kit + src[tag.end():]
    if "Kit." not in src or "window.Kit" in src:
        return src
    m = re.search(r"<script", src, re.I) or re.search(r"</body>", src, re.I)
    return src[:m.start()] + kit + src[m.start():] if m else src + kit


def inline_ui(src):
    css, js = "<style>\n" + UICSS + "\n</style>", "<script>\n" + UIJS + "\n</script>"
    link = re.search(r"""<link[^>]*href=["'](?:\./)?ui\.css["'][^>]*>""", src, re.I)
    if link:
        src = src[:link.start()] + css + src[link.end():]
    else:  # the kit comes first so the page's own <style> can override it
        m = re.search(r"<head[^>]*>", src, re.I)
        src = src[:m.end()] + "\n" + css + src[m.end():] if m else css + src
    tags = list(re.finditer(r"""<script[^>]*\bsrc=["'](?:\./)?ui\.js["'][^>]*>\s*</script>""", src, re.I))
    if tags:  # the first reference becomes the inline kit, repeated references are dropped
        for t in reversed(tags[1:]):
            src = src[:t.start()] + src[t.end():]
        return src[:tags[0].start()] + js + src[tags[0].end():]
    m = re.search(r"</body>", src, re.I)
    return src[:m.start()] + js + "\n" + src[m.start():] if m else src + js


def run_svg_program(code, timeout=20):
    """Runs a svgkit program in a fresh folder; returns (svg or None, problem or None)."""
    d = tempfile.mkdtemp(prefix="svgkit_")
    try:
        shutil.copy(os.path.join(KITS, "svgkit.py"), d)
        path = os.path.join(d, "draw.py")
        with open(path, "w", encoding="utf-8") as f:
            f.write(code)
        try:
            r = sandbox.run([sys.executable, "-E", "-s", path], input=b"", capture_output=True, timeout=timeout, cwd=d,
                            creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        except subprocess.TimeoutExpired:
            return None, f"The program did not finish within {timeout} seconds."
        out, err = r.stdout.decode(errors="replace"), r.stderr.decode(errors="replace")
        if r.returncode != 0:
            return None, "The program failed:\n" + "\n".join(err.strip().splitlines()[-6:])
    finally:
        shutil.rmtree(d, ignore_errors=True)
    m = re.search(r"<svg.*</svg>", out, re.S)
    if not m:
        return None, "The program printed no <svg>...</svg> document (end it with print(s)). It printed:\n" + out[:300]
    try:
        xml.dom.minidom.parseString(m.group(0).encode("utf-8"))
    except Exception as e:
        return None, f"The printed SVG is not valid XML: {str(e)[:200]}"
    return m.group(0), None


def finalize(text, cat):
    """(final answer text, problem or None, source for rendering or None)."""
    if cat == "svg":
        code = extract_code(text)
        if code is None or "svgkit" not in code:
            src = extract(text, "svg")  # a direct SVG answer is still accepted
            if src:
                return f"```svg\n{src}\n```", None, src
            return text, "Your answer does not contain the Python program in a ```python code block.", None
        svg, problem = run_svg_program(code)
        return (f"```svg\n{svg}\n```" if svg else text), problem, svg
    src = extract(text, "html")
    if not src:
        return text, None, None
    new = {"three": inline_three, "canvas": inline_canvas, "html": inline_ui}[cat](src)
    return text.replace(src, new), None, new


CPLIB = _read("cplib.py").split('"""', 2)[2].lstrip()  # the library without its module docstring
CPLIB_IMPORT = re.compile(r"^[ \t]*(from\s+cplib\s+import\s+[^\n]+|import\s+cplib\b[^\n]*)$", re.M)


def inline_cplib(text):
    """Replaces `from cplib import ...` / `import cplib` in the answer's last ```python block by the library."""
    blocks = list(re.finditer(r"```(?:python|py|Python)?[ \t]*\n(.*?)```", text or "", re.S))
    if not blocks or not CPLIB_IMPORT.search(blocks[-1].group(1)):
        return text
    code = blocks[-1].group(1)
    module_style = re.search(r"^[ \t]*import\s+cplib\b", code, re.M)
    lib = CPLIB + ("\nimport types as _types\ncplib = _types.SimpleNamespace(**{k: v for k, v in dict(globals()).items() "
                   "if not k.startswith('_')})\n" if module_style else "")
    new = CPLIB_IMPORT.sub(lambda m: "# --- cplib (inlined) ---\n" + lib + "# --- end of cplib ---", code, count=1)
    new = CPLIB_IMPORT.sub("", new)
    s, e = blocks[-1].span(1)
    return text[:s] + new + text[e:]