#!/usr/bin/bash

set -u

BLACK="\033[0;30m"
RED="\033[0;31m"
GREEN="\033[0;32m"
YELLOW="\033[0;33m"
BLUE="\033[0;34m"
PURPLE="\033[0;35m"
CYAN="\033[0;36m"
GRAY="\033[0;37m"
BOLD_GRAY="\033[1;30m"
BOLD_RED="\033[1;31m"
BOLD_GREEN="\033[1;32m"
BOLD_YELLOW="\033[1;33m"
BOLD_BLUE="\033[1;34m"
BOLD_PURPLE="\033[1;35m"
BOLD_CYAN="\033[1;36m"
BOLD_WHITE="\033[1;37m"
NC="\033[0m"

printHeader() {
    echo -e "${BOLD_RED}~~~~~~~~${BOLD_WHITE} Git pre-commit hook ${BOLD_RED}~~~~~~~~${NC}"
}

printFooter() {
    echo -e "${BOLD_RED}~~~~~~~~~~~${BOLD_WHITE} Hook finished ${BOLD_RED}~~~~~~~~~~~${NC}"
    echo
}

formatCode() {
    echo -e "${GREEN}Running${NC} isort"
    isort . --resolve-all-configs
    echo

    echo -e "${GREEN}Running${NC} black"
    black .
}

doPreCommit() {
    local STAGED_AND_MODIFIED_FILES
    STAGED_AND_MODIFIED_FILES="$(git diff --name-only --cached | sort -d | comm -12 - <(git diff --name-only | sort -d))"
    if [[ -z "$STAGED_AND_MODIFIED_FILES" ]]; then
        # No files were changed after staging, we can safely run formatters.
        echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Running code formatters.${NC}"
        formatCode
        echo

        STAGED_AND_MODIFIED_FILES="$(git diff --name-only --cached | sort -d | comm -12 - <(git diff --name-only | sort -d))"
        if [[ -z "$STAGED_AND_MODIFIED_FILES" ]]; then
            echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}No changes were made to staged files.${NC}"
            return 0
        fi

        # We formatted something, so we need to add these files back to stage.
        echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Adding refactored changes to stage.${NC}"
        echo -e "Files to add:"
        echo -e "$(echo "$STAGED_AND_MODIFIED_FILES" | sed -n "s/^/ > /p")"
        git add $STAGED_AND_MODIFIED_FILES
        echo
        return 0
    fi

    # There are some changes on staged files, which are not added to stage yet.
    # We don't want to proceed with this situation by default.

    if [ "${FORCE_PRE_COMMIT_WITH_UNSTAGED:-}" != true ]; then
        echo -e "${BOLD_YELLOW}WARNING${NC}: It seems that you have unstaged changes for files that are already staged:"
        echo -e "${BOLD_WHITE}$(echo "$STAGED_AND_MODIFIED_FILES" | sed -n "s/^/ > /p")${NC}"
        echo
        echo -e "If it's unintended, add these files to stage and retry."
        echo
        echo -e "Otherwise, if you consciously excluded these changes from stage, you can try to proceed anyway."
        echo -e "${BOLD_YELLOW}By doing this, you are confirming that you have carefully read the docs and checked the code.${NC}"
        echo -e "Rerun with 'FORCE_PRE_COMMIT_WITH_UNSTAGED=true git commit ...' to retry."
        return 1
    fi

    # We got permission from user to do dirty things.
    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Starting refactor with unstaged changes.${NC}"
    echo -e "Files both in staged and in changed:"
    echo -e "${BOLD_WHITE}$(echo "$STAGED_AND_MODIFIED_FILES" | sed -n "s/^/ > /p")${NC}"
    echo

    # Commit current stage, since we can't refactor both staged and unstaged files at once.
    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Temporarily saving staged changes as a commit.${NC}"
    git commit -m "Staged changes" --no-verify
    echo

    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Running code formatters on unstaged files to prevent merge conflicts later.${NC}"
    formatCode
    echo

    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Saving changed files to stash.${NC}"
    git stash -- $STAGED_AND_MODIFIED_FILES
    echo

    # We have completed work with unstaged changes, now we can get back the staged ones.
    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Rolling back temporary commit.${NC}"
    git reset --soft HEAD^
    echo

    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Running code formatters on commit files.${NC}"
    formatCode
    echo

    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Adding refactored changes to stage${NC}"
    STAGED_FILES="$(git diff --name-only --cached)"
    echo -e "Files to add:"
    echo -e "$(echo "$STAGED_FILES" | sed -n "s/^/ > /p")"
    git add $STAGED_FILES
    echo

    # Refactor completed, get back unstaged changes.
    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Restoring stashed file changes.${NC}"
    git stash pop
    echo

    echo -e "${BOLD_CYAN}--> ${BOLD_YELLOW}Completed.${NC}"
}

if [[ -v DISABLE_CODE_FMT ]]; then
    echo -e "${BOLD_YELLOW}DISABLE_CODE_FMT is set, skipping code formatting${NC}"
    echo
    exit 0
fi

RC=0
printHeader
doPreCommit || RC=1
printFooter
exit "$RC"
