#!/usr/bin/env bash

###############################################################################
# Formats ipython notebooks with tensorflow-docs nbformat tool.
#
# Usage:
#     check/nbformat [--apply]
#
# Without '--apply', the diff that would be applied is printed and the exit
# status is 1 if there are any changes or else 0 if no changes are needed.
#
# With '--apply', the exit status is 0 and the changed files are actually
# reformatted.
#
################################################################################

only_print=1
for arg in "$@"; do
    if [[ "${arg}" == "--apply" ]]; then
        only_print=0
    else
        echo -e "\033[31mToo many arguments. Expected [--apply].\033[0m" >&2
        exit 1
    fi
done

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

pip show tensorflow-docs > /dev/null || exit 1

FORMAT_CMD="python3 -m tensorflow_docs.tools.nbfmt --indent=1"

# Test the notebooks
unformatted=$($FORMAT_CMD --test docs 2>&1 | grep "\- docs" || true)
needed_changes=0
if [ -n "${unformatted}" ]; then
    needed_changes=1
    if (( only_print == 0 )); then
        $FORMAT_CMD docs
    else
        echo -e "\033[31mThe following notebooks require formatting\033[0m."
        echo "${unformatted}"
    fi
fi

if (( needed_changes == 0 )); then
    echo -e "\033[32mNotebooks are formatted\033[0m."
elif (( only_print == 1 )); then
    echo -e "\033[33mNotebooks are not formatted. Please run 'check/nbformat --apply'\033[0m"
    exit 1
else
    echo -e "\033[33mReformatted changed notebooks\033[0m."
fi
