#!/bin/bash
#
# DESCRIPTION:
#
#   Set the bash prompt according to:
#    * the active virtualenv
#    * the branch/status of the current Git, Mercurial or Subversion repository
#    * the return value of the previous command
#    * one line prompt
#
# USAGE:
#
#   1. Save this file as ~/.bash_prompt
#   2. Add the following line to the end of your ~/.bashrc or ~/.bash_profile:
#        . ~/.bash_prompt
#
# LINEAGE:
#
#   Based on work by bradsokol
#
#       https://gist.github.com/bradsokol/2959821

# The various escape codes that we can use to color our prompt.
        RED="\[\033[0;31m\]"
     YELLOW="\[\033[1;33m\]"
     PURPLE="\[\033[1;35m\]"
      GREEN="\[\033[0;32m\]"
       BLUE="\[\033[1;34m\]"
       CYAN="\[\033[1;36m\]"
  LIGHT_RED="\[\033[1;31m\]"
LIGHT_GREEN="\[\033[1;32m\]"
      WHITE="\[\033[1;37m\]"
 LIGHT_GRAY="\[\033[0;37m\]"
 COLOR_NONE="\[\e[0m\]"

# THRESHOLD is expressed in seconds.
THRESHOLD=3

# various symbols, depending on used font in urxvt
GIT_PICT=""
HG_PICT=""
PY_PICT=""
JS_PICT=""


# Detect whether the current directory is a git repository.
function is_git_repository {
    git branch > /dev/null 2>&1
}

# Detect whether the current directory is a Mercurial repository.
function is_mercurial_repository {
    branch=$(hg branch 2>/dev/null)
    if [ -n "${branch}" ]; then
        return 0
    else
        return 1
    fi
}

# Detect whether the current directory is a Subversion repository.
function is_subversion_repository {
    test -d .svn
}

# Determine the branch/state information for this git repository.
function set_git_branch {
    # Capture the output of the "git status" command.
    git_status="$(git status 2> /dev/null)"

    # Set color based on clean/staged/dirty.
    if [[ ${git_status} =~ "working tree clean" ]]; then
        state="${GREEN}"
    elif [[ ${git_status} =~ "Changes to be committed" ]]; then
        state="${YELLOW}"
    else
        state="${LIGHT_RED}"
    fi

    # Set arrow icon based on status against remote.
    remote_pattern="# Your branch is (.*) of"
    if [[ ${git_status} =~ ${remote_pattern} ]]; then
        if [[ ${BASH_REMATCH[1]} == "ahead" ]]; then
            remote="↑"
        else
            remote="↓"
        fi
    else
        remote=""
    fi
    diverge_pattern="# Your branch and (.*) have diverged"
    if [[ ${git_status} =~ ${diverge_pattern} ]]; then
        remote="↕"
    fi

    # Get the name of the branch.
    gitsym=$(git symbolic-ref HEAD 2>/dev/null)
    if [[ $? == 0 ]]; then
        branch="${gitsym##refs/heads/}"
    fi

    # Set the final branch string.
    BRANCH="${state}${GIT_PICT} (${branch})${remote}${COLOR_NONE} "
}

# Determine the branch/state information for this Mercurial repository.
function set_mercurial_branch {
    # Get the name of the branch.
    branch=$(hg branch 2>/dev/null)
    # Default state
    state="${GREEN}"

    if [ -n "${branch}" ]; then
        branch="(${branch})"

        # Capture the output of the "hg status" command.
        hg_status="$(hg status | wc -l)"

        # Set color based on clean/staged/dirty.
        if [ "${hg_status}" -ne "0" ]; then
            state="${RED}"
        fi
    fi

    # Set the final branch string.
    BRANCH="${state}${HG_PICT} ${branch}${COLOR_NONE} "
}

# Determine the branch informatioin for this Subversion repository. No support
# for svn status, since that needs to hit the remote repository.
function set_subversion_branch {
    # Capture the output of the "svn info" command
    svn_info="$(svn info | egrep '^URL: ' 2> /dev/null)"

    # Get the name of the branch.
    branch_pattern="^URL: .*/(branches|tags)/([^/]+)"
    trunk_pattern="^URL: .*/trunk(/.*)?$"
    if [[ ${svn_info} =~ $branch_pattern ]]; then
        branch=${BASH_REMATCH[2]}
    elif [[ ${svn_info} =~ $trunk_pattern ]]; then
        branch='trunk'
    fi

    # Set the final branch string.
    BRANCH="(${branch}) "
}

# Return the prompt symbol to use, colorized based on the return value of the
# previous command.
function set_prompt_symbol {
    if test $1 -eq 0 ; then
        PROMPT_SYMBOL="${BLUE}\$${COLOR_NONE}"
    else
        PROMPT_SYMBOL="${LIGHT_RED}\$${COLOR_NONE}"
    fi
}

# Determine active Python virtualenv details.
function set_virtualenv {
    if test -z "$VIRTUAL_ENV" ; then
        PYTHON_VIRTUALENV=""
    else
        PYTHON_VIRTUALENV="${CYAN}[${PY_PICT} `basename \"$VIRTUAL_ENV\"`]${COLOR_NONE} "
    fi
}

function set_nodevirtenv {
    if test -z "$NODE_VIRTUAL_ENV" ; then
        NODE_VIRTUALENV=""
    else
        NODE_VIRTUALENV="${PURPLE}[${JS_PICT} `basename \"$NODE_VIRTUAL_ENV\"`]${COLOR_NONE} "
    fi
}

# Shamlesely stolen from: https://stackoverflow.com/a/31694983 and improve it
# with minimal time which is acceptable for executing command.
function timer_now {
    date +%s%N
}

function timer_start {
    timer_start=${timer_start:-$(timer_now)}
}

function timer_stop {
    local delta_us=$((($(timer_now) - $timer_start) / 1000))
    local us=$((delta_us % 1000))
    local ms=$(((delta_us / 1000) % 1000))
    local s=$(((delta_us / 1000000) % 60))
    local m=$(((delta_us / 60000000) % 60))
    local h=$((delta_us / 3600000000))

    timer_show=""

    if [[ $s -lt $THRESHOLD ]]; then
        unset timer_start
        return
    fi

    # Always show around 3 digits of accuracy
    if ((h > 0)); then
        timer_show=${h}h${m}m
    elif ((m > 0)); then
        timer_show=${m}m${s}s
    elif ((s >= 10)); then
        timer_show=${s}.$((ms / 100))s
    elif ((s > 0)); then
        timer_show=${s}.$(printf %03d $ms)s
    elif ((ms >= 100)); then
        timer_show=${ms}ms
    elif ((ms > 0)); then
        timer_show=${ms}.$((us / 100))ms
    else
        timer_show=${us}us
    fi
    unset timer_start
}

# Set the full bash prompt.
function set_bash_prompt {
    # Set the PROMPT_SYMBOL variable. We do this first so we don't lose the
    # return value of the last command.
    set_prompt_symbol $?

    # Set the PYTHON_VIRTUALENV variable.
    set_virtualenv

    # Set the NODE_PYTHON_VIRTUALENV variable.
    set_nodevirtenv

    # Set the BRANCH variable.
    if is_git_repository ; then
        set_git_branch
    elif is_subversion_repository ; then
        set_subversion_branch
    elif is_mercurial_repository ; then
        set_mercurial_branch
    else
        BRANCH=''
    fi

    timer_stop
    TIME=""
    if [[ -n "${timer_show}" ]]; then
        TIME="${LIGHT_GRAY}${timer_show}${COLOR_NONE} "
    fi

  # Set the bash prompt variable.
  PS1="${CYAN}\u@\h ${PYTHON_VIRTUALENV}${NODE_VIRTUALENV}"
  PS1="${PS1}${BLUE}\w${COLOR_NONE} ${BRANCH}${TIME}${PROMPT_SYMBOL} "
}

# start measuring time of execution
trap 'timer_start' DEBUG

# Tell bash to execute this function just before displaying its prompt.
PROMPT_COMMAND=set_bash_prompt