#!/bin/bash
RETURN_CODE="$?"
RESET="\[\e[0m\]"
BOLD="\[\e[1m\]"
R="\[\e[1;31m\]"
G="\[\e[1;32m\]"
Y="\[\e[1;33m\]"
B="\[\e[1;34m\]"
P="\[\e[1;35m\]"
C="\[\e[1;36m\]"
TIME="${G}\t${RESET}"
function prompt_command {
    if [[ $EUID == 0 ]] ; then
        local PROMPT_COLOR="$R"
        
        # getting hg and git info is disabled for root
        local HG_INFO=" ${P}-${RESET}"
        local GIT_INFO=" ${G}-${RESET}"
    else
        local PROMPT_COLOR="$G"
        
        # testing if current dir is a part of hg/git repo
        local TEST_DIR="$PWD"
            
        while true ; do
            [[ -d "$TEST_DIR/.hg" ]] && local HG_TRACKED=1 && break
            [[ -d "$TEST_DIR/.git" ]] && local GIT_TRACKED=1 && break
            [[ "$TEST_DIR" == '/' ]] && break
            TEST_DIR=$(dirname "$TEST_DIR")
        done
        
        if [[ $HG_TRACKED ]] ; then
            local SUM=$(LC_ALL=C hg sum 2>/dev/null)
            local PARENTS=($(grep -Po '(?<=parent: ).*(?=:)' <<< "$SUM"))
            local BRANCH=$(grep -Po '(?<=branch: ).*$' <<< "$SUM")
            local COMMIT=$(grep -Po '(?<=commit: ).*$' <<< "$SUM")
            local UPDATE=$(grep -Po '(?<=update: ).*$' <<< "$SUM")
            local REV=${PARENTS[0]}
            local MERGEREV=${PARENTS[1]}
            BRANCH=${BRANCH/default/}
            [ -z "$MERGEREV" ] || MERGEREV="+$MERGEREV"
            case "$COMMIT" in
                *'(clean)'*)
                    COMMIT=''
                    ;;
                *added*|*modified*|*deleted*)
                    COMMIT='!'
                    ;;
            esac
            case "$UPDATE" in
                *'(update)'*)
                    UPDATE='^'
                    ;;
                *'(merge)'*)
                    UPDATE=''
                    ;;
                *'(current)'*)
                    UPDATE=''
                    ;;
            esac
            local HG_INFO=" ${P}${BRANCH}☿${REV}${R}${UPDATE}${G}${MERGEREV}${R}${COMMIT}${RESET}"
        fi
        if [[ $GIT_TRACKED ]] ; then
            local GIT_STATUS=$(LC_ALL=C git status 2>/dev/null)
            local GIT_BRANCH=$(grep -Po '(?<=On branch ).*$' <<< "$GIT_STATUS")
            local GIT_REV=$(LC_ALL=C git rev-parse --short HEAD 2>/dev/null)
            
            if [[ $GIT_STATUS == *'Your branch is behind'* ]]; then
                local GIT_BEHIND="$R^"
            fi
            if [[ $GIT_STATUS == *'Your branch is ahead'* ]]; then
                local GIT_AHEAD="$Y^"
            fi
            if [[ $GIT_STATUS == *'Changes to be committed'* ]]; then
                local GIT_STAGED="$R!"
            fi
            if [[ $GIT_STATUS == *'Changes not staged'* ]]; then
                local GIT_UNSTAGED="$P?"
            fi
            
            local GIT_INFO=" ${G}${GIT_BRANCH}±${GIT_REV}${GIT_BEHIND}${GIT_AHEAD}${GIT_STAGED}${GIT_UNSTAGED}${RESET}"
        fi
    fi
    if [[ $VIRTUAL_ENV ]] ; then
        local VENV="${G}+"$(basename "$VIRTUAL_ENV")"${RESET}"
    fi
    if [[ $RETURN_CODE != 0 ]] ; then
        local RETCODE="${R}${RETURN_CODE}${RESET} "
    fi
    
    PS1="\
${BOLD}[${TIME} \u@\h \
${RETCODE}\W${VENV}${HG_INFO}${GIT_INFO}${BOLD}]\
${PROMPT_COLOR}\\\$${RESET} "
}
export VIRTUAL_ENV_DISABLE_PROMPT=1
export PROMPT_COMMAND=prompt_command