#!/usr/bin/env bash
# This script sets iTerm2 user-defined variables describing the state of the git
# repo in the current directory.
#
# For more information on the status bar, see:
# https://www.iterm2.com/3.3/documentation-status-bar.html
#
# Installation instructions for this script:
#
# bash: Place this in .bashrc.
# --------------------------------------------------------------------------------------
# function iterm2_print_user_vars() {
#   it2git
# }

# fish: Place this in ~/.config/fish/config.fish after the line
#       "source ~/.iterm2_shell_integration.fish".
# --------------------------------------------------------------------------------------
# function iterm2_print_user_vars
#   it2git
# end

# tcsh: Place this in .tcshrc
# --------------------------------------------------------------------------------------
# alias get_current_branch "bash -c '((git branch 2> /dev/null) | grep \* | cut -c3-)'"
# alias _iterm2_user_defined_vars 'it2git'

# zsh:Place this in .zshrc after "source /Users/georgen/.iterm2_shell_integration.zsh".
# --------------------------------------------------------------------------------------
# iterm2_print_user_vars() {
#   it2git
# }

GIT_BINARY=/usr/bin/git

# The counts follow libgit2's, which iTerm2 uses for a directory on this Mac, so a repository
# reads the same whichever way iTerm2 learns about it. Compatible with bash 3.2.

# Never take index.lock: this only looks at a repository someone may be using.
export GIT_OPTIONAL_LOCKS=0

# Reads `status --porcelain -b -z` once and sets BRANCH, PUSH_COUNT, PULL_COUNT, DIRTY, ADDS
# (untracked files), DELETES (deleted in the work tree) and UNMERGED (an array of paths).
# Every untracked file counts, not just its directory, as libgit2 lists them.
read_status() {
    BRANCH=""
    PUSH_COUNT=""
    PULL_COUNT=""
    ADDS=0
    DELETES=0
    UNMERGED=()
    local count=0 record xy path skip=0 header counts upstream ahead behind item
    while IFS= read -r -d '' record; do
        if ((skip)); then
            # The original path of a rename or copy.
            skip=0
            continue
        fi
        if [[ "$record" == "## "* ]]; then
            header="${record#\#\# }"
            counts=""
            if [[ "$header" == *"]" && "$header" == *" ["* ]]; then
                counts="${header##* [}"
                counts="${counts%]}"
                header="${header% \[*}"
            fi
            upstream=0
            if [[ "$header" == *...* ]]; then
                upstream=1
                header="${header%%...*}"
            fi
            if ((upstream)) && [[ "$counts" != "gone" ]]; then
                ahead=0
                behind=0
                local IFS_SAVED="$IFS"
                IFS=','
                for item in $counts; do
                    item="${item# }"
                    case "$item" in
                        "ahead "*) ahead="${item#ahead }" ;;
                        "behind "*) behind="${item#behind }" ;;
                    esac
                done
                IFS="$IFS_SAVED"
                PUSH_COUNT="$ahead"
                PULL_COUNT="$behind"
            fi
            if [[ "$header" == "HEAD (no branch)" ]]; then
                # Detached: libgit2 reports the full id.
                BRANCH=$("$GIT_BINARY" rev-parse HEAD 2>/dev/null)
            else
                BRANCH="$header"
            fi
            continue
        fi
        xy="${record:0:2}"
        path="${record:3}"
        case "$xy" in
            '!!') continue ;;
        esac
        count=$((count + 1))
        case "$xy" in
            '??') ADDS=$((ADDS + 1)); continue ;;
            DD|AU|UD|UA|DU|AA|UU) UNMERGED[${#UNMERGED[@]}]="$path"; continue ;;
        esac
        case "$xy" in
            # A rename in the work tree (of an intent-to-add file) has its R in the second column.
            R*|C*|?R|?C) skip=1 ;;
        esac
        # libgit2 finds no renames in the work tree, so the original of one renamed there (an
        # intent-to-add file) is deleted.
        if [[ "${xy:1:1}" == "D" || "${xy:1:1}" == "R" ]]; then
            DELETES=$((DELETES + 1))
        fi
    done < <("$GIT_BINARY" status --porcelain -b -z --untracked-files=all --ignore-submodules=all 2>/dev/null)
    if ((count > 0)); then
        DIRTY="dirty"
    else
        DIRTY="clean"
    fi
}

# Compared whole, since a path may hold any character but NUL. There are seldom more than a few.
is_unmerged() {
    local unmerged
    for unmerged in ${UNMERGED[@]+"${UNMERGED[@]}"}; do
        [[ "$unmerged" == "$1" ]] && return 0
    done
    return 1
}

# Changes against HEAD, without rename detection, as libgit2's stats pass sees them: a
# typechange is a deletion and an addition, a conflict is not counted, and lines are counted
# for modified files only (binary files count none). Untracked files count as added. Sets
# FILES_ADDED, FILES_MODIFIED, FILES_DELETED, LINES_INSERTED and LINES_DELETED.
read_diff_stats() {
    FILES_ADDED=$ADDS
    FILES_MODIFIED=0
    FILES_DELETED=0
    LINES_INSERTED=0
    LINES_DELETED=0
    local letter path record ins del
    while IFS= read -r -d '' letter && IFS= read -r -d '' path; do
        if is_unmerged "$path"; then
            continue
        fi
        case "$letter" in
            A) FILES_ADDED=$((FILES_ADDED + 1)) ;;
            D) FILES_DELETED=$((FILES_DELETED + 1)) ;;
            T) FILES_ADDED=$((FILES_ADDED + 1)); FILES_DELETED=$((FILES_DELETED + 1)) ;;
            M) FILES_MODIFIED=$((FILES_MODIFIED + 1)) ;;
        esac
    done < <("$GIT_BINARY" -c diff.relative=false diff --no-ext-diff --no-renames -z --name-status HEAD -- 2>/dev/null)
    while IFS= read -r -d '' record; do
        # "INSERTED<TAB>DELETED<TAB>PATH", or "-<TAB>-<TAB>PATH" for a binary file.
        ins="${record%%$'\t'*}"
        record="${record#*$'\t'}"
        del="${record%%$'\t'*}"
        path="${record#*$'\t'}"
        if is_unmerged "$path"; then
            continue
        fi
        [[ "$ins" == "-" ]] && ins=0
        [[ "$del" == "-" ]] && del=0
        LINES_INSERTED=$((LINES_INSERTED + ins))
        LINES_DELETED=$((LINES_DELETED + del))
    done < <("$GIT_BINARY" -c diff.relative=false diff --no-ext-diff --no-renames -z --numstat --diff-filter=M HEAD -- 2>/dev/null)
}
function iterm2_begin_osc {
  printf "\033]"
}

function iterm2_end_osc {
  printf "\007"
}

function iterm2_set_user_var() {
  # Encoded before the sequence starts and written whole: tmux 2.6 to 3.0a give up on a sequence
  # that pauses for 100 ms partway and show the rest as text.
  local _iterm2_value
  _iterm2_value=$(printf "%s" "$2" | base64 | tr -d '\n')
  printf "\033]1337;SetUserVar=%s=%s\007" "$1" "$_iterm2_value"
}

git_poll () {
    read_status
    read_diff_stats

    iterm2_set_user_var gitBranch "$BRANCH"
    iterm2_set_user_var gitDirty "$DIRTY"
    iterm2_set_user_var gitPushCount "$PUSH_COUNT"
    iterm2_set_user_var gitPullCount "$PULL_COUNT"
    iterm2_set_user_var gitAdds "$ADDS"
    iterm2_set_user_var gitDeletes "$DELETES"
    iterm2_set_user_var gitLinesInserted "$LINES_INSERTED"
    iterm2_set_user_var gitLinesDeleted "$LINES_DELETED"
    iterm2_set_user_var gitFilesAdded "$FILES_ADDED"
    iterm2_set_user_var gitFilesModified "$FILES_MODIFIED"
    iterm2_set_user_var gitFilesDeleted "$FILES_DELETED"
}

# Outside a repository, or before its first commit (libgit2 then reports no state), every
# variable is cleared.
if ! "$GIT_BINARY" rev-parse --verify -q HEAD >/dev/null 2>&1; then
    iterm2_set_user_var gitBranch ""
    iterm2_set_user_var gitDirty ""
    iterm2_set_user_var gitPushCount ""
    iterm2_set_user_var gitPullCount ""
    iterm2_set_user_var gitAdds ""
    iterm2_set_user_var gitDeletes ""
    iterm2_set_user_var gitLinesInserted ""
    iterm2_set_user_var gitLinesDeleted ""
    iterm2_set_user_var gitFilesAdded ""
    iterm2_set_user_var gitFilesModified ""
    iterm2_set_user_var gitFilesDeleted ""
else
    git_poll
fi


