"""执行 notebook 的全部代码单元（按顺序、共享全局命名空间），用于零依赖自检。

用法::

    python scripts/run_notebook_cells.py [notebook.ipynb] [--record]

等价于在 Jupyter 中"全部运行"，但不需要 nbformat / nbclient：直接读 ipynb 的
JSON、拼接 code 单元、写入临时文件后交给 `python` 执行。退出码即执行结果；
``--record`` 额外把逐单元的标准输出写回 notebook 的输出字段（生成已执行副本）。
"""

from __future__ import annotations

import argparse
import contextlib
import io
import json
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent


def collect(notebook_path: Path) -> tuple[dict, list[str]]:
    notebook = json.loads(notebook_path.read_text(encoding="utf-8"))
    codes = ["".join(cell["source"])
             for cell in notebook["cells"] if cell["cell_type"] == "code"]
    return notebook, codes


def run_and_record(notebook: dict) -> None:
    """在本进程内逐单元执行并记录标准输出（写入 notebook 的 outputs 字段）。"""
    scope: dict = {"__name__": "__main__"}
    index = 0
    for cell in notebook["cells"]:
        if cell["cell_type"] != "code":
            continue
        buffer = io.StringIO()
        with contextlib.redirect_stdout(buffer):
            # 与 Jupyter 同语义：单元之间共享同一个全局命名空间
            exec(compile("".join(cell["source"]),      # pylint: disable=exec-used
                         f"<cell {index}>", "exec"), scope)
        text = buffer.getvalue()
        cell["execution_count"] = index + 1
        cell["outputs"] = ([{"output_type": "stream", "name": "stdout",
                             "text": text.splitlines(keepends=True)}] if text else [])
        index += 1


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="执行 notebook 代码单元并自检")
    parser.add_argument("notebook", nargs="?", type=Path,
                        default=ROOT / "mcm26b-核心代码.ipynb")
    parser.add_argument("--record", action="store_true",
                        help="在当前进程内执行并把输出写回 notebook（生成已执行副本）")
    args = parser.parse_args(argv)

    notebook, codes = collect(args.notebook)
    # 逐单元做语法检查，定位到具体单元
    for index, cell_source in enumerate(codes):
        try:
            compile(cell_source, f"<cell {index}>", "exec")
        except SyntaxError as exc:
            print(f"单元 {index} 语法错误：{exc}", file=sys.stderr)
            return 1
    total_lines = sum(len(cell.splitlines()) for cell in codes)
    print(f"语法检查通过：{len(codes)} 个代码单元，{total_lines} 行")

    if args.record:
        run_and_record(notebook)
        args.notebook.write_text(
            json.dumps(notebook, ensure_ascii=False, indent=1) + "\n",
            encoding="utf-8")
        print(f"已把逐单元输出写回 {args.notebook}")
        return 0

    with tempfile.TemporaryDirectory() as tmp:
        target = Path(tmp) / "notebook_flat.py"
        target.write_text("\n\n".join(codes), encoding="utf-8")
        process = subprocess.run([sys.executable, str(target)], cwd=ROOT,
                                 check=False)
    return process.returncode


if __name__ == "__main__":
    raise SystemExit(main())