"""Validate or execute the plain-Python cells in the Quant60 Colab notebook."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


def load_notebook(path: Path) -> dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if payload.get("nbformat") != 4 or not isinstance(payload.get("cells"), list):
        raise ValueError("expected a valid nbformat 4 notebook")
    return payload


def code_cells(
    payload: dict[str, Any],
    *,
    include_optional: bool,
) -> list[tuple[int, str]]:
    output: list[tuple[int, str]] = []
    for index, cell in enumerate(payload["cells"]):
        if cell.get("cell_type") != "code":
            continue
        tags = cell.get("metadata", {}).get("tags", [])
        if "optional" in tags and not include_optional:
            continue
        source = cell.get("source")
        if not isinstance(source, list) or any(
            not isinstance(line, str) for line in source
        ):
            raise ValueError(f"cell {index} source must be a string array")
        output.append((index, "".join(source)))
    return output


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=Path)
    parser.add_argument("--execute", action="store_true")
    parser.add_argument("--include-optional", action="store_true")
    args = parser.parse_args(argv)

    payload = load_notebook(args.path.resolve())
    cells = code_cells(payload, include_optional=args.include_optional)
    for index, source in cells:
        compile(source, f"{args.path}:cell-{index}", "exec")
    if args.execute:
        namespace: dict[str, Any] = {"__name__": "__quant60_notebook__"}
        for index, source in cells:
            print(f"[quant60-notebook] executing cell {index}")
            exec(
                compile(source, f"{args.path}:cell-{index}", "exec"),
                namespace,
                namespace,
            )
    print(
        json.dumps(
            {
                "ok": True,
                "code_cells": len(cells),
                "executed": bool(args.execute),
                "optional_included": bool(args.include_optional),
            },
            sort_keys=True,
        )
    )
    return 0


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