

set -eo pipefail

builtin cd "$(dirname "${BASH_SOURCE:-$0}")"
ROOT="$(git rev-parse --show-toplevel)"
builtin cd "$ROOT" || exit 1

BLACK_VERSION=$(black --version | head -n 1 | awk '{print $2}')
PYLINT_VERSION=$(pylint --version | head -n 1 | awk '{print $2}')

tool_version_check() {
    if [[ $2 != $3 ]]; then
        echo "Wrong $1 version installed: $3 is required, not $2."
        exit 1
    fi
}

tool_version_check "black" $BLACK_VERSION "23.3.0"
tool_version_check "pylint" $PYLINT_VERSION "2.8.2"

format_changed() {
    MERGEBASE="$(git merge-base origin/main HEAD)"

    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs -P 5 black
    fi
}

if [[ "$1" == '--files' ]]; then
   black "${@:2}"
elif [[ "$1" == '--all' ]]; then
   black fastchat
else
   format_changed
fi
echo 'FastChat Black: Done'

echo 'FastChat Pylint:'
pylint fastchat

if ! git diff --quiet &>/dev/null; then
    echo 'Reformatted files. Please review and stage the changes.'
    echo 'Changes not staged for commit:'
    echo
    git --no-pager diff --name-only

    exit 1
fi
