"""Run the doc notebooks and hold each to the outputs it is stored with.

    python tests/run_notebooks.py notebooks/*.ipynb           # run and check
    python tests/run_notebooks.py --write notebooks/*.ipynb   # run and store the outputs

Each notebook runs in a fresh kernel (this Python's), in its own directory. A
cell that raises fails the run. Then every code cell's outputs must equal the
stored ones: its streams, and of each display the HTML when it has some (else
its plain text). What is left out of the comparison is what differs by design
between two runs: a widget's view (its model id is new each run) and the plain
text beside a display's HTML. Cells tagged ``service`` need a job service
(``WAI_QUANTUM_SERVICE``); without one they print that they skip, and their
outputs are not compared.

Exit status: 0 when every notebook ran and matched; 1 otherwise, with a diff.
"""

from __future__ import annotations

import argparse
import difflib
import json
import os
import sys
from pathlib import Path

import nbformat
from nbclient import NotebookClient

WIDGET_VIEW = "application/vnd.jupyter.widget-view+json"


def comparable(output: dict) -> dict:
    kind = output.get("output_type")
    if kind == "stream":
        return {"stream": output.get("name"), "text": "".join(output.get("text", ""))}
    if kind in ("display_data", "execute_result"):
        data = {k: ("".join(v) if isinstance(v, list) else v) for k, v in output.get("data", {}).items() if k != WIDGET_VIEW}
        if "text/html" in data:
            data.pop("text/plain", None)
        return {"display": data}
    if kind == "error":
        return {"error": output.get("ename"), "value": output.get("evalue")}
    return {"other": kind}


def merged(outputs: list) -> list:
    """Outputs as compared: consecutive stream chunks of one stream joined."""
    out: list = []
    for o in map(comparable, outputs):
        if out and "stream" in o and out[-1].get("stream") == o["stream"]:
            out[-1] = {"stream": o["stream"], "text": out[-1]["text"] + o["text"]}
        else:
            out.append(o)
    return out


def run(path: Path, timeout: int) -> nbformat.NotebookNode:
    nb = nbformat.read(path, as_version=4)
    client = NotebookClient(nb, timeout=timeout, kernel_name="python3", resources={"metadata": {"path": str(path.parent)}}, store_widget_state=False)
    client.execute()
    nb.metadata.pop("widgets", None)
    return nb


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--write", action="store_true", help="store the outputs instead of checking them")
    ap.add_argument("--timeout", type=int, default=600, help="seconds per cell")
    ap.add_argument("notebooks", nargs="+", type=Path)
    a = ap.parse_args()
    have_service = bool(os.environ.get("WAI_QUANTUM_SERVICE"))
    failed = 0
    cells = 0
    for path in a.notebooks:
        stored = nbformat.read(path, as_version=4)
        try:
            ran = run(path, a.timeout)
        except Exception as e:  # a cell raised, or the kernel died
            print(f"FAIL {path}: {type(e).__name__}: {e}")
            failed += 1
            continue
        if a.write:
            nbformat.write(ran, path)
            print(f"wrote {path}")
            continue
        bad = []
        for i, (s, r) in enumerate(zip(stored.cells, ran.cells)):
            if s.cell_type != "code":
                continue
            if "service" in s.metadata.get("tags", []) and not have_service:
                continue
            cells += 1
            want, got = merged(s.outputs), merged(r.outputs)
            if want != got:
                w = json.dumps(want, indent=1, ensure_ascii=False).splitlines()
                g = json.dumps(got, indent=1, ensure_ascii=False).splitlines()
                bad.append(f"cell {i}:\n" + "\n".join(list(difflib.unified_diff(w, g, "stored", "this run", lineterm=""))[:80]))
        if bad:
            failed += 1
            print(f"FAIL {path}: {len(bad)} cell(s) differ from their stored outputs")
            for b in bad:
                print(b)
        else:
            print(f"ok   {path}")
    if not a.write:
        print(f"notebooks: {len(a.notebooks) - failed} of {len(a.notebooks)} ran and matched, {cells} code cells compared" + ("" if have_service else " (no service: its cells skipped)"))
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())