#!/usr/bin/env bash
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -euo pipefail

declare -r usage="\
Usage: ${0##*/} [--help] [--apply]

Checks Jupyter notebook (.ipynb) format using tensorflow_docs.tools.nbfmt.

Without --apply, exits 0 if all notebooks are formatted and 1 otherwise.
With --apply, reformats notebooks in place.
"

# Get the working directory to the repo root.
thisdir=$(dirname "${BASH_SOURCE[0]:?}")
repo_dir=$(git -C "${thisdir}" rev-parse --show-toplevel)
cd "${repo_dir}"

opt_apply=0
for arg in "$@"; do
    case "${arg}" in
        -h | --help)
            echo "${usage}"
            exit 0
            ;;
        --apply)
            opt_apply=1
            ;;
        *)
            echo "Unknown argument: '${arg}'" >&2
            echo "See '${0} --help' for usage." >&2
            exit 1
            ;;
    esac
done

if ! python3 -c "import tensorflow_docs.tools.nbfmt"; then
    echo "tensorflow-docs is not installed." >&2
    exit 2
fi

my_nbfmt() {
    python3 -m tensorflow_docs.tools.nbfmt --indent=1 "$@"
}

exit_code=0
if output=$(my_nbfmt --test . 2>&1); then
    echo "Notebooks are formatted."
elif (( opt_apply )); then
    my_nbfmt . &&
    echo "Reformatted notebooks." ||
    exit_code=$?
else
    echo "The following notebooks require formatting:" >&2
    grep '^-' <<<"${output}" >&2 || echo "${output}" >&2
    echo "Run 'check/nbformat --apply' to fix notebooks." >&2
    exit_code=1
fi

exit "${exit_code}"
