prompt

#!/bin/bash

# Selects a colour for the prompt that wasn't used by the previous terminal
function init_prompt() {
    # 0=black, 1=red, 2=green, 3=yellow, 4=blue, 5=magenta, 6=cyan, 7=grey, 8-15=bright variants
    if [[ $SHELL == "/bin/zsh" ]]; then
        if [[ -f ~/.zsh_colour ]]; then
            export ZSH_COLOUR=$(cat ~/.zsh_colour)
        else
            export ZSH_COLOUR=0
        fi
        let 'ZSH_COLOUR = (ZSH_COLOUR + 1) % 4'
        echo "$ZSH_COLOUR" > ~/.zsh_colour
        if [[ "ZSH_COLOUR" -ge 2 ]]; then
            let 'ZSH_COLOUR += 3'
        else
            let 'ZSH_COLOUR += 1'
        fi
        precmd_functions+=(update_prompt)
    elif [[ $SHELL == "/bin/bash" ]]; then
        if [[ -f ~/.bash_colour ]]; then
            export BASH_COLOUR=$(cat ~/.bash_colour)
        else
            export BASH_COLOUR=0
        fi
        let 'BASH_COLOUR = (BASH_COLOUR + 1) % 4'
        echo $BASH_COLOUR > ~/.bash_colour
        if [[ "$BASH_COLOUR" -ge 2 ]]; then
            let 'BASH_COLOUR += 33'
        else
            let 'BASH_COLOUR += 31'
        fi
        export PROMPT_COMMAND=update_prompt
    else
        echo "Unknown shell. Will not configure prompt colour."
    fi
}

# Sets the prompt with the current directory, git branch and git status
function update_prompt() {
    local host=${HOSTNAME%.local}
    local user=$USER

    if [[ $SHELL == "/bin/zsh" ]]; then
        local colour="%F{$ZSH_COLOUR}"
        local reset="%f"
    else
        local colour="\[\e[${BASH_COLOUR}m\]"
        local reset="\[\e[0m\]"
    fi

    local repo=$(git rev-parse --show-toplevel 2> /dev/null)
    if [[ ! -z "$repo" ]]; then
        if [[ $SHELL == "/bin/zsh" ]]; then
            local red='%F{1}'
            local green='%F{2}'
            local yellow='%F{3}'
        else
            local red='\[\e[31m\]'
            local green='\[\e[32m\]'
            local yellow='\[\e[33m\]'
        fi

        local branch='HEAD'
        if [[ ! -z $(git branch) ]]; then
            # Get the current branch name or use "HEAD" if repo is newly initialized
            branch=$(git rev-parse --abbrev-ref HEAD)
        fi
        local stat=$(git status --porcelain)

        local warn=$green
        if [[ ! -z $(echo "$stat" | cut -c1-2 | grep 'U') ]]; then
            # Contains conflicts
            warn=$red
        elif [[ ! -z "$stat" ]]; then
            # Contains modifications
            warn=$yellow
        fi

        local pwd_phys=$(pwd -P)
        local base=$(basename "$repo")
        local path=$(echo "${pwd_phys#$repo}" | $MISC_ROOT/abbrev_path)

        if [[ $SHELL == "/bin/zsh" ]]; then
            export PROMPT="${colour}%n@%m${reset} ${base}${warn}(${branch})${reset}${path} ${colour}%#${reset} "
        else
            export PS1="${colour}${host}:${reset}${base}${warn}(${branch})${reset}${path} ${colour}${user}\$${reset} "
        fi
    else
        local path=$PWD
        if [[ "$path" =~ ^"$HOME"(/|$) ]]; then
            path="~${path#$HOME}"
        fi
        path=$(echo "$path" | $MISC_ROOT/abbrev_path)

        if [[ $SHELL == "/bin/zsh" ]]; then
            export PROMPT="${colour}%n@%m${reset} %1~ ${colour}%#${reset} "
        else
            export PS1="${colour}${host}:${reset}${path} ${colour}${user}\$${reset} "
        fi
    fi
}