"""Execute a percent-format source (# %% / # %% [markdown]) and write an executed .ipynb with outputs.

usage: python _m1_run.py src.py out.ipynb
Runs all code cells in ONE namespace (like a kernel), captures stdout/stderr per cell,
stores them as stream outputs.  Needs nothing but the stdlib (+ whatever the notebook imports).
"""
import sys, io, json, traceback, contextlib, os, time
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
src, out = sys.argv[1], sys.argv[2]
sys.path.insert(0, os.path.dirname(os.path.abspath(src)) + "/..")
lines = open(src).read().split("\n")
cells, cur, kind = [], [], None
def flush():
    global cur, kind
    if kind is not None:
        body = cur[:]
        while body and body[-1] == "": body.pop()
        cells.append((kind, body))
    cur, kind = [], None
for ln in lines:
    if ln.startswith("# %%"):
        flush(); kind = "markdown" if "[markdown]" in ln else "code"
    else:
        cur.append(ln)
flush()
ns = {"__name__": "__main__"}
nb_cells, count = [], 0
t0 = time.time()
for kind, body in cells:
    if kind == "markdown":
        text = "\n".join(l[2:] if l.startswith("# ") else l.lstrip("#") for l in body)
        nb_cells.append({"cell_type": "markdown", "metadata": {}, "source": text.splitlines(True)})
        continue
    count += 1
    code = "\n".join(body)
    so, se = io.StringIO(), io.StringIO()
    err = None
    with contextlib.redirect_stdout(so), contextlib.redirect_stderr(se):
        try:
            exec(compile(code, f"<cell {count}>", "exec"), ns)
        except BaseException:
            err = traceback.format_exc()
    outputs = []
    if so.getvalue(): outputs.append({"output_type": "stream", "name": "stdout", "text": so.getvalue().splitlines(True)})
    if se.getvalue(): outputs.append({"output_type": "stream", "name": "stderr", "text": se.getvalue().splitlines(True)})
    if err:
        outputs.append({"output_type": "error", "ename": "Error", "evalue": err.splitlines()[-1], "traceback": err.splitlines()})
    nb_cells.append({"cell_type": "code", "execution_count": count, "metadata": {}, "source": code.splitlines(True), "outputs": outputs})
    if err:
        print(err, file=sys.stderr)
        nb = {"cells": nb_cells, "metadata": {}, "nbformat": 4, "nbformat_minor": 4}
        json.dump(nb, open(out + ".FAILED", "w"), indent=1)
        sys.exit(1)
nb = {"cells": nb_cells,
      "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
                   "language_info": {"name": "python", "version": sys.version.split()[0]}},
      "nbformat": 4, "nbformat_minor": 4}
json.dump(nb, open(out, "w"), indent=1)
print(f"wrote {out}: {count} code cells in {time.time()-t0:.1f}s")