#!/usr/bin/env python3
"""Check or strip outputs from Jupyter notebooks before they are committed.

Notebook outputs are where AI experiments leak: printed API responses, sample
customer records, tokens echoed by a debug cell, or megabytes of embeddings.
This script removes them (or fails CI when they are present) using only the
standard library, by editing the notebook JSON directly.

Usage:
    python3 scripts/strip_notebooks.py --check [PATH ...]
    python3 scripts/strip_notebooks.py --fix   [PATH ...] [--keep-metadata]

    PATH             Notebook files or directories. Defaults to the current directory.
    --check          Report notebooks that contain outputs or execution counts; change nothing.
    --fix            Clear outputs and execution counts and rewrite the files.
    --keep-metadata  With --fix, keep per-cell metadata (by default the volatile
                     "execution" and "collapsed"/"scrolled" keys are removed).

What --fix changes, per code cell:
    outputs -> [], execution_count -> null, and the volatile metadata keys above.
Markdown and raw cells, cell sources and notebook-level metadata (kernel, language)
are kept. Files are written with 1-space indentation, the format Jupyter itself uses,
so diffs stay small.

Requirements: Python 3.8+, no packages, no environment variables.

Exit codes:
    0  clean (--check) or files rewritten successfully (--fix)
    1  --check found notebooks with outputs
    2  usage error, unreadable file, or a file that is not notebook JSON
"""
import argparse
import json
import os
import sys

SKIP_DIRS = {".git", "node_modules", ".venv", "venv", "__pycache__", ".ipynb_checkpoints"}
VOLATILE_KEYS = ("execution", "collapsed", "scrolled", "ExecuteTime")


def find_notebooks(paths):
    found = []
    for root in paths:
        if os.path.isfile(root):
            found.append(root)
        elif os.path.isdir(root):
            for dirpath, dirnames, filenames in os.walk(root):
                dirnames[:] = sorted(d for d in dirnames if d not in SKIP_DIRS)
                found.extend(os.path.join(dirpath, f) for f in sorted(filenames) if f.endswith(".ipynb"))
        else:
            print("error: path not found: %s" % root, file=sys.stderr)
            sys.exit(2)
    return found


def dirty_cells(notebook, keep_metadata):
    count = 0
    for cell in notebook.get("cells", []):
        if cell.get("cell_type") != "code":
            continue
        meta = cell.get("metadata", {})
        if cell.get("outputs") or cell.get("execution_count") is not None or (
            not keep_metadata and any(k in meta for k in VOLATILE_KEYS)
        ):
            count += 1
    return count


def strip(notebook, keep_metadata):
    for cell in notebook.get("cells", []):
        if cell.get("cell_type") != "code":
            continue
        cell["outputs"] = []
        cell["execution_count"] = None
        if not keep_metadata:
            meta = cell.get("metadata", {})
            for key in VOLATILE_KEYS:
                meta.pop(key, None)


def main():
    parser = argparse.ArgumentParser(description="Check or strip Jupyter notebook outputs.")
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--check", action="store_true")
    mode.add_argument("--fix", action="store_true")
    parser.add_argument("--keep-metadata", action="store_true")
    parser.add_argument("paths", nargs="*", default=["."])
    args = parser.parse_args()

    notebooks = find_notebooks(args.paths)
    dirty = 0
    for path in notebooks:
        try:
            with open(path, encoding="utf-8") as handle:
                notebook = json.load(handle)
        except (OSError, ValueError) as exc:
            print("error: cannot read %s as notebook JSON: %s" % (path, exc), file=sys.stderr)
            return 2
        if not isinstance(notebook, dict) or "cells" not in notebook:
            print("error: %s has no 'cells' list; is it an nbformat 4 notebook?" % path, file=sys.stderr)
            return 2
        cells = dirty_cells(notebook, args.keep_metadata)
        if not cells:
            continue
        dirty += 1
        if args.check:
            print("%s: %d code cell(s) with outputs or execution data" % (path, cells))
            continue
        strip(notebook, args.keep_metadata)
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(notebook, handle, indent=1, ensure_ascii=False)
            handle.write("\n")
        print("stripped %s (%d cell(s))" % (path, cells))

    if args.check and dirty:
        print("\n%d notebook(s) need stripping. Run: python3 scripts/strip_notebooks.py --fix" % dirty)
        return 1
    print("%d notebook(s) checked, %d %s" % (len(notebooks), dirty, "stripped" if args.fix else "dirty"))
    return 0


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