#!/usr/bin/env bash

################################################################################
# Based on the Cirq python-formatting script:
# https://github.com/quantumlib/Cirq/blob/master/check/format-incremental
#
# Formats python files that have been modified.
#
# Usage:
#     check/format-incremental [BASE_REVISION] [--apply] [--all]
#
# By default, the script runs flynt to convert old string literal formatting to f-strings
# and analyzes python files that have changed relative to the
# base revision and determines whether they need to be formatted. If any changes
# are needed, it prints the diff and exits with code 1, otherwise it exits with
# code 0.
#
# With '--apply', reformats the files instead of printing the diff and exits
# with code 0.
#
# With '--all', analyzes all python files, instead of only changed files.
#
# You can specify a base git revision to compare against (i.e. to use when
# determining whether or not a file is considered to have "changed"). For
# example, you can compare against 'origin/master' or 'HEAD~1'.
#
# If you don't specify a base revision, the following defaults will be tried, in
# order, until one exists:
#
#     1. upstream/master
#     2. origin/master
#     3. master
#
# If none exists, the script fails.
################################################################################

# Get the working directory to the repo root.
cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
cd "$(git rev-parse --show-toplevel)"


# Parse arguments.
only_print=1
only_changed=1
rev=""
for arg in $@; do
    if [[ "${arg}" == "--apply" ]]; then
        only_print=0
    elif [[ "${arg}" == "--all" ]]; then
        only_changed=0
    elif [ -z "${rev}" ]; then
        if [ "$(git cat-file -t ${arg} 2> /dev/null)" != "commit" ]; then
            echo -e "\033[31mNo revision '${arg}'.\033[0m" >&2
            exit 1
        fi
        rev="${arg}"
    else
        echo -e "\033[31mToo many arguments. Expected [revision] [--apply] [--all].\033[0m" >&2
        exit 1
    fi
done

if (( only_changed == 1 )); then
    # Figure out which branch to compare against.
    if [ -z "${rev}" ]; then
        if [ "$(git cat-file -t upstream/master 2> /dev/null)" == "commit" ]; then
            rev=upstream/master
        elif [ "$(git cat-file -t origin/master 2> /dev/null)" == "commit" ]; then
            rev=origin/master
        elif [ "$(git cat-file -t master 2> /dev/null)" == "commit" ]; then
            rev=master
        else
            echo -e "\033[31mNo default revision found to compare against. Argument #1 must be what to diff against (e.g. 'origin/master' or 'HEAD~1').\033[0m" >&2
            exit 1
        fi
    fi
    base="$(git merge-base ${rev} HEAD)"
    if [ "$(git rev-parse ${rev})" == "${base}" ]; then
        echo -e "Comparing against revision '${rev}'." >&2
    else
        echo -e "Comparing against revision '${rev}' (merge base ${base})." >&2
        rev="${base}"
    fi

    # Get the modified and added python files. Lines in git diff --name-status look like this:
    # M       qsimcirq/__init__.py
    # A       qsimcirq/added.py
    modified_files=$(git diff --name-status ${rev} -- | grep '\.py$' | grep -v '_pb2\.py$' | grep '^[MA].*$' | awk '{print $2}')

    # Get the moved python files. Lines in git diff --name-status look like this:
    # R100    qsimcirq/_doc.py    qsimcirq/_doc2.py
    moved_files=$(git diff --name-status ${rev} -- | grep '\.py$' | grep -v '_pb2\.py$' | grep '^R.*$' | awk '{print $3}')

    format_files="$modified_files $moved_files"

else
    echo -e "Formatting all python files." >&2
    format_files=$(find . -name "*.py" | grep -v "_pb2\.py$")
fi

any_files=0
for format_file in ${format_files}; do
  any_files=1
done

if (( any_files == 0 )); then
    echo -e "\033[32mNo files to format\033[0m."
    exit 0
fi

flynt_args=("--quiet" "-f")
if (( only_print == 1 )); then
    flynt_args+=("--dry-run")
fi

# Exclude setup.py due to issues in release pipeline (#351).
flynt ${format_files} "${flynt_args[@]}" -e setup\.py
FLYNTSTATUS=$?

echo "Running the black formatter..."

args=("--color")
if (( only_print == 1 )); then
    args+=("--check" "--diff")
fi

# Warn users about bug in black: https://github.com/psf/black/issues/1629
# Once that is fixed upstream, we can just do:
# black "${args[@]}" ${format_files}
# exit $?
LOGS="$(black "${args[@]}" ${format_files} 2>&1)"
BLACKSTATUS=$?
echo "${LOGS}"
if [[ "$BLACKSTATUS" == "123" && "$LOGS" =~ ^.*"INTERNAL ERROR: Black produced different code on the second pass of the formatter.".*$ ]]; then
  echo -e "\033[31mWarning, it seems we ran into https://github.com/psf/black/issues/1629. Typically, this can be fixed by adding a trailing comma. If you get stuck, please file a cirq issue on github.\033[0m"
fi

if [[ "$FLYNTSTATUS" != "0" || "$BLACKSTATUS" != "0"  ]]; then
  exit 1
fi
exit 0