#   Set the bash prompt according to:
#    * the branch/status of the current git repository
#    * the return value of the previous command
#   AUTHOR: Scott Woods <scott@westarete.com> West Arete Computing
#   Based on work by halbtuerke and lakiolen.
#   http://gist.github.com/31967

# Unicode symbols https://github.com/pjmp/fancy-linux-prompt/blob/master/fancy-prompt.sh
# The various escape codes that we can use to color our prompt.
# Glued together by Robert Strutts 
     GIT_BRANCH_CHANGED_SYMBOL='+ changed'
     GIT_NEED_PULL_SYMBOL='⇣ do pull'
     GIT_NEED_PUSH_SYMBOL='⇡ needs push'
    
     BOLD="\\[$(tput bold)\\]"
     DIM="\\[$(tput dim)\\]"
     RESET="\\[$(tput sgr0)\\]"
     REVERSE="\\[$(tput rev)\\]"

        RED="\[\033[1;31m\]"
     YELLOW="\[\033[1;33m\]"
      GREEN="\[\033[0;32m\]"
       BLUE="\[\033[1;34m\]"
    MAGENTA="\[\033[1;35m\]"
       CYAN="\[\033[1;36m\]"
  LIGHT_RED="\[\033[1;31m\]"
LIGHT_GREEN="\[\033[1;32m\]"
 LIGHT_BLUE="\[\033[1;94m\]"
 LIGHT_CYAN="\[\033[1;96m\]"
      WHITE="\[\033[1;37m\]"
      BLACK="\[\033[1;30m\]"
 LIGHT_GRAY="\[\033[1;37m\]"
 
 BK_DEFAULT="\[\e[49m\]"
 BK_BLACK="\[\e[40m\]"
 BK_RED="\[\e[41m\]"
 BK_GREEN="\[\e[42m\]"
 BK_YELLOW="\[\e[43m\]"
 BK_BLUE="\[\e[44m\]"
 BK_LIGHT_RED="\[\e[101m\]"
 BK_LIGHT_GREEN="\[\e[102m\]"
 BK_LIGHT_YELLOW="\[\e[103m\]"
 BK_LIGHT_BLUE="\[\e[104m\]"
 BK_LIGHT_CYAN="\[\e[106m\]"
 BK_LIGHT_WHITE="\[\e[107m\]"
 BK_EXIT="\[\e[49m\]"
 
 COLOR_NONE="\[\e[0m\]"
 COLOR_DEFAULT="\[\e[39m\]"

[ -f ~/scripts/helper/.kube-ps1 ] && source ~/scripts/helper/.kube-ps1

function is_normal_theme() {
    # sudo apt-get install fonts-powerline
    
    if [ -f ~/.data/.unicode_support ]; then
     PS_SYMBOL='🐧'
     SEGMENT_SEPARATOR=$'\ue0b0'
     PL_BRANCH_CHAR=$'\ue0a0'         # 
    else
     PS_SYMBOL=''
     SEGMENT_SEPARATOR=' -> '
     PL_BRANCH_CHAR=''
    fi      
        
    if [ -f ~/.data/.simple_theme ]; then
        return 1
    else
        return 0
    fi      
}
 
__git_info() {
        # no .git directory
        [ -d .git ] || return

        local aheadN
        local behindN
        local branch
        local marks
        local stats

        # get current branch name or short SHA1 hash for detached head
        branch="$(git symbolic-ref --short HEAD 2>/dev/null || git describe --tags --always 2>/dev/null)"
        [ -n "$branch" ] || return  # git branch not found

        # how many commits local branch is ahead/behind of remote?
        stats="$(git status --porcelain --branch | grep '^##' | grep -o '\[.\+\]$')"
        aheadN="$(echo "$stats" | grep -o 'ahead \d\+' | grep -o '\d\+')"
        behindN="$(echo "$stats" | grep -o 'behind \d\+' | grep -o '\d\+')"
        [ -n "$aheadN" ] && marks+=" $GIT_NEED_PUSH_SYMBOL$aheadN"
        [ -n "$behindN" ] && marks+=" $GIT_NEED_PULL_SYMBOL$behindN"

        if ! is_normal_theme; then
            printf "%s" "$state[$branch]$WHITE$marks"
           return 
        fi   

        # print the git branch segment without a trailing newline
        # branch is modified?
        if [ -n "$(git status --porcelain)" ]; then
            printf "%s" "$RESET$bk$WHITE[$branch]$marks$state"
        else
            printf "%s" "$RESET$bk$WHITE[$branch]$marks$RESET$state"
        fi
    } 
 
# Detect whether the current directory is a git repository.
function is_git_repository {
  git branch > /dev/null 2>&1
}
 
# 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}"
    bk="$BK_GREEN"
  elif [[ ${git_status} =~ "Changes to be committed" ]]; then
    state="${YELLOW}"
    bk="$BK_YELLOW"
  else
    state="${RED}"
    bk="$BK_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="↑ needs push"
    else
      remote="↓ do pull"
    fi
  else
    remote="$PL_BRANCH_CHAR"
  fi
  diverge_pattern="# Your branch and (.*) have diverged"
  if [[ ${git_status} =~ ${diverge_pattern} ]]; then
    remote="↕ changed"
  fi
  
  BRANCH="${state}(git)${COLOR_NONE}${remote} "
  BRANCH+="$(__git_info bk)"
  if is_normal_theme; then
    BRANCH+="$BK_EXIT${SEGMENT_SEPARATOR}$RESET"
  fi    
}

 
# Return the prompt symbol to use, colorized based on the return value of the
# previous command.
function set_prompt_symbol() {
  if is_normal_theme; then
     ERRX="✘"
  else 
     ERRX="Opps!"
  fi    
  if test $1 -eq 0 ; then
      PROMPT_SYMBOL="\$${COLOR_NONE}"
  else
      PROMPT_SYMBOL="${RED}${ERRX} \$${COLOR_NONE}"
  fi
}

function pydir() {
    local pyt=$(dirname "$1")
    local pw=$(pwd)
    if [ "$pyt" == "$pw" ]; then
        return 0
    fi
    return 1    
}     
# Function to check if inside a Python virtual environment
function is_in_python_env() {
    # Check for VIRTUAL_ENV environment variable (most reliable method)
    if [ -n "$VIRTUAL_ENV" ]; then
        if pydir "$VIRTUAL_ENV"; then
           return 0
        fi   
    fi
    
    # Alternative check for venv/pipenv/conda environments
    if [ -n "$CONDA_DEFAULT_ENV" ]; then
        if pydir "$CONDA_DEFAULT_ENV"; then
            return 0
        fi    
    fi 
    if [ -n "$PIPENV_ACTIVE" ]; then
        return 0
    fi
   
    return 1
}

 
# 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 BRANCH variable.
  if is_git_repository ; then
    set_git_branch
  else
    BRANCH=''
  fi

  if is_in_python_env ; then
    PY="🐍"
  else
    PY=''
  fi  
  # Set the bash prompt variable.  
  
  if [ -n "$SSH_CLIENT" ] || [ -n "$SSH_TTY" ]; then
    _distro=$(awk '/^ID=/' /etc/*-release | awk -F'=' '{ print tolower($2) }')

    # set an icon based on the distro
    # make sure your font is compatible with https://github.com/lukas-w/font-logos
case $_distro in
    *kali*)                  ICON="ﴣ";;
    *arch*)                  ICON="";;
    *debian*)                ICON="";;
    *raspbian*)              ICON="";;
    *ubuntu*)                ICON="";;
    *elementary*)            ICON="";;
    *fedora*)                ICON="";;
    *coreos*)                ICON="";;
    *gentoo*)                ICON="";;
    *mageia*)                ICON="";;
    *centos*)                ICON="";;
    *opensuse*|*tumbleweed*) ICON="";;
    *sabayon*)               ICON="";;
    *slackware*)             ICON="";;
    *linuxmint*)             ICON="";;
    *alpine*)                ICON="";;
    *aosc*)                  ICON="";;
    *nixos*)                 ICON="";;
    *devuan*)                ICON="";;
    *manjaro*)               ICON="";;
    *rhel*)                  ICON="";;
    *macos*)                 ICON="";;
    *)                       ICON="";;
esac

    UPC="${RED}SSH:${WHITE}${PY} ${ICON} "
  else
    UPC="Local:${PY}"    
  fi
  
  # if you want kube prompt
  [ -f ~/scripts/helper/.kube-ps1 ] && PS1="$(kube_ps1)\r\n"
  [ ! -f ~/scripts/helper/.kube-ps1 ] && PS1=""   

  if [ -f ~/scripts/helper/theme ]; then
    source ~/scripts/helper/theme
  else
    UPC+="\u@\h" 
    if is_normal_theme; then
        PS1+="${BK_GREEN}$WHITE$UPC${GREEN}${SEGMENT_SEPARATOR}"
        WD="$WHITE$BK_BLUE\\w${COLOR_NONE}${BLUE}${SEGMENT_SEPARATOR}${WHITE}"
        PS1+="${WHITE}${BRANCH}${WD}$PS_SYMBOL${PROMPT_SYMBOL} "
    else
        WD="\\w${COLOR_NONE}${WHITE}"
        PS1+="$CYAN$UPC ${WHITE}${BRANCH} ${WD} ${PROMPT_SYMBOL} "
    fi  
  fi
}
# Tell bash to execute this function just before displaying its prompt.
PROMPT_COMMAND=set_bash_prompt
