#!/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 info is disabled for root
        local HG_INFO=" ${P}-${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 TRACKED=1 && break
            [[ "$TEST_DIR" == '/' ]] && break
            TEST_DIR=$(dirname "$TEST_DIR")
        done
        
        if [[ "$TRACKED" ]] ; then
            local HG_INFO=" $(hg prompt $PROMPT_TEMPLATE)"
        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}${BOLD}]\
${PROMPT_COLOR}\\\$${RESET} "
}
export PROMPT_COMMAND=prompt_command