"""Read a list of Jupyter Notebook files and append some installation preamble if needed."""

import sys

import nbformat as nbf

SKIP_TAG = "nbsetup-skip"
TAG = "differt-install-preamble"

CELL = nbf.NotebookNode(
    cell_type="code",
    execution_count=None,
    id="0",
    metadata={"tags": [TAG, "remove-cell", "skip-execution"]},
    outputs=[],
    source=(
        "# Run this cell to install DiffeRT and its dependencies, e.g., on Google Colab\n"
        "\n"
        "try:\n"
        "    import differt  # ruff: ignore[unused-import]\n"
        "except ImportError:\n"
        "    import sys  # ruff: ignore[unused-import]\n"
        "\n"
        "    !{sys.executable} -m pip install differt[all]"
    ),
)

if __name__ == "__main__":
    files_changed = False
    for ipath in sys.argv[1:]:
        ntbk = nbf.read(ipath, nbf.NO_CONVERT)

        if ntbk.metadata.get("nbsetup", {}).get("skip", False):
            continue

        cells = ntbk.cells

        if len(cells) > 1 and TAG in cells[0].get("metadata", {}).get("tags", []):
            cell = cells[0]

            if SKIP_TAG in cell["metadata"]["tags"]:
                continue  # Skip this cell, we don't want to overwrite it

            for field, value in CELL.items():
                if field not in cell or cell[field] != value:
                    cell[field] = value
                    files_changed = True
        else:
            cells.insert(0, CELL)
            files_changed = True

        nbf.write(ntbk, ipath)
    if files_changed:
        sys.exit(1)
