#!/bin/bash
function prompt_command {
    local RETURN_CODE="$?"
    local RESET="\[\e[0m\]"
    local BOLD="\[\e[1m\]"
    local R="\[\e[1;31m\]"
    local G="\[\e[1;32m\]"
    local Y="\[\e[1;33m\]"
    local B="\[\e[1;34m\]"
    local P="\[\e[1;35m\]"
    local C="\[\e[1;36m\]"
    local PROMPT_TEMPLATE="${P}{branch|quiet}☿{rev}{$R{update}}{$G+{rev|merge}}$R{status|modified}${RESET}"
    local TIME="${G}\t${RESET}"
    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 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 HG_INFO=" $(hg prompt $PROMPT_TEMPLATE)"
        fi
        if [[ "$GIT_TRACKED" ]] ; then
            local GIT_STATUS=$(git status 2>/dev/null)
            
            if [[ $(echo ${GIT_STATUS} | grep "# Your branch is behind") ]]; then
                local GIT_UPDATE="$R^"
            fi
            if [[ $(echo ${GIT_STATUS} | grep "# Changes to be committed:") ]]; then
                local GIT_STAGED="$R!"
            fi
            if [[ $(echo ${GIT_STATUS} | grep "# Changes not staged for commit:") ]]; then
                local GIT_UNSTAGED="$P?"
            fi
            
            local GIT_INFO=" ${G}±${GIT_STAGED}${GIT_UNSTAGED}${GIT_UPDATE}${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 PROMPT_COMMAND=prompt_command