#!/usr/bin/env python3

"""Filter Jupyter notebook cells by tag."""

import argparse
import sys

import nbformat


def main() -> None:
    """Filter Jupyter notebook cells by tag."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "-i",
        "--input",
        nargs="?",
        type=argparse.FileType("r"),
        default=sys.stdin,
        help="input file",
    )
    parser.add_argument(
        "-o",
        "--output",
        nargs="?",
        type=argparse.FileType("w"),
        default=sys.stdout,
        help="output file",
    )
    parser.add_argument(
        "-t",
        "--tag",
        nargs="?",
        type=str,
        default="exercise",
        help="tag to filter out",
    )
    parser.add_argument(
        "--show-celltoolbar",
        action="store_true",
        help="show the cell toolbar in output notebook (default: False)",
    )
    args = parser.parse_args()
    notebook = nbformat.read(args.input, as_version=nbformat.NO_CONVERT)
    filtered_cells = []
    for cell in notebook["cells"]:
        if "tags" in cell["metadata"] and args.tag in cell["metadata"]["tags"]:
            continue
        filtered_cells.append(cell)
    if not args.show_celltoolbar:
        metadata = notebook["metadata"]
        if "celltoolbar" in metadata:
            metadata.pop("celltoolbar")
    notebook["cells"] = filtered_cells
    nbformat.write(notebook, args.output)


if __name__ == "__main__":
    main()
