#!/bin/bash
function prompt_command {
    local RETURN_CODE="$?"
    # dark, light, bold
    local Rd="\[\e[0;31m\]"; local R="\[\e[0;91m\]"; local Rb="\[\e[1;31m\]"
    local Gd="\[\e[0;32m\]"; local G="\[\e[0;92m\]"; local Gb="\[\e[1;32m\]"
    local Yd="\[\e[0;33m\]"; local Y="\[\e[0;93m\]"; local Yb="\[\e[1;33m\]"
    local Bd="\[\e[0;34m\]"; local B="\[\e[0;94m\]"; local Bb="\[\e[1;34m\]"
    local Pd="\[\e[0;35m\]"; local P="\[\e[0;95m\]"; local Pb="\[\e[1;35m\]"
    local Cd="\[\e[0;36m\]"; local C="\[\e[0;96m\]"; local Cb="\[\e[1;36m\]"
    if [[ $TERM == 'linux' ]] ; then
        R=$Rb
        G=$Gb
        Y=$Yb
        B=$Bb
        P=$Pb
        C=$Cb
    fi
    local BOLD="\[\e[1m\]"
    local RESET="\[\e[0m\]"
    local TIME="${G}\t${RESET}"
    if [[ $RETURN_CODE != 0 ]] ; then
        local RETCODE="${R}${RETURN_CODE}${RESET} "
    fi
    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 && $USE_HG_PROMPT ]] ; then
            local PROMPT_TEMPLATE="${P}{branch|quiet}{${Rd}{closed}}${P}☿{rev}{${B}+{tags|quiet|+}}{${Gd}*{bookmark}}{${R}{update}}{${G}+{rev|merge}}${R}{status|modified}${RESET}"
            local HG_INFO=" $(hg prompt $PROMPT_TEMPLATE)"
        fi
        if [[ $HG_TRACKED && ! $USE_HG_PROMPT ]] ; 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^"
                [[ $GIT_STATUS =~ behind[^$'\n']+([0-9]+)' 'commit ]] && GIT_BEHIND+=${BASH_REMATCH[1]}
            fi
            if [[ $GIT_STATUS == *'Your branch is ahead'* ]] ; then
                local GIT_AHEAD="$Y^"
                [[ $GIT_STATUS =~ ahead[^$'\n']+([0-9]+)' 'commit ]] && GIT_AHEAD+=${BASH_REMATCH[1]}
            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
    PS1="\
[${TIME} \u@\h \
${RETCODE}\W${VENV}${HG_INFO}${GIT_INFO}]\
${PROMPT_COLOR}\\\$${RESET} "
}
# first find prompt in hgrc, then check it's actually enabled
if grep -w '^prompt' ~/.hgrc &>/dev/null && hg help prompt &>/dev/null ; then
    USE_HG_PROMPT=1
fi
export VIRTUAL_ENV_DISABLE_PROMPT=1
export PROMPT_COMMAND=prompt_command