#!/bin/zsh
#
# This file was written by Bart Trojanowski <bart@jukie.net>
#
# documented on my blog:
#   http://www.jukie.net/~bart/blog/tag/zsh
#
# references
#   http://www.zsh.org/mla/users/2006/msg01196.html
#   http://dotfiles.org/~frogb/.zshrc
#   http://kriener.org/articles/2009/06/04/zsh-prompt-magic

setopt prompt_subst
autoload colors
colors

autoload -Uz vcs_info

# -------------------------------
# define core prompt functions

# set some colors
for COLOR in RED GREEN BLUE YELLOW WHITE BLACK CYAN; do
    eval PR_$COLOR='%{$fg[${(L)COLOR}]%}'
    eval PR_BRIGHT_$COLOR='%{$fg_bold[${(L)COLOR}]%}'
done
PR_RST="%{${reset_color}%}"
PR_RESET="%{%b%s%u$reset_color%}"
PR_BG="%{%(?.$PR_RESET.%S)%}"

# set formats
# %b - branchname
# %u - unstagedstr (see below)
# %c - stangedstr (see below)
# %a - action (e.g. rebase-i)
# %R - repository path
# %S - path in the repository
FMT_BRANCH="${PR_GREEN}%b%u%c${PR_RST}" # e.g. master¹²
FMT_ACTION="(${PR_CYAN}%a${PR_RST}%)"   # e.g. (rebase-i)
FMT_PATH="%R${PR_YELLOW}/%S"              # e.g. ~/repo/subdir

# check-for-changes can be really slow.
# you should disable it, if you work with large repositories
zstyle ':vcs_info:*:prompt:*' check-for-changes true
zstyle ':vcs_info:*:prompt:*' unstagedstr   '¹'  # display ¹ if there are unstaged changes
zstyle ':vcs_info:*:prompt:*' stagedstr     '²'  # display ² if there are staged changes
zstyle ':vcs_info:*:prompt:*' actionformats "${FMT_BRANCH}${FMT_ACTION}//" "${FMT_PATH}"
zstyle ':vcs_info:*:prompt:*' formats       "${FMT_BRANCH}//"              "${FMT_PATH}"
zstyle ':vcs_info:*:prompt:*' nvcsformats   ""                             "%~"

function lprompt {
    local brackets=$1
    local color1=$2
    local color2=$3

    local bracket_open="${color1}${brackets[1]}${PR_BG}"
    local bracket_close="${color1}${brackets[2]}${PR_RESET}"

    local git='$vcs_info_msg_0_'
    local cwd="${color2}%B%1~%b"

    local vimode='${PR_VIMODE}'
    local vicol='${PR_VICOLOR}'

    #PROMPT="${PR_BG}${bracket_open}${git}${cwd}${bracket_close} ${vicol}${vimode}${PR_RESET} "
    PROMPT="${PR_BG}${bracket_open}${git}${cwd}${bracket_close} ${vicol}${vimode}${PR_RESET} "
}

function rprompt {
    local brackets=$1
    local color1=$2
    local color2=$3

    local bracket_open="${color1}${brackets[1]}${PR_RESET}"
    local bracket_close="${color1}${brackets[2]}${PR_RESET}"
    local colon="${color1}:"
    local at="${color1}@${PR_RESET}"

    local user_host="${color2}%n${at}${color2}%m"
    local vcs_cwd='${${vcs_info_msg_1_%%.}/$HOME/~}'
    local cwd="${color2}%B%20<..<${vcs_cwd}%<<%b"
    local inner="${user_host}${colon}${cwd}"

    RPROMPT="${PR_RESET}${bracket_open}${inner}${bracket_close}${PR_RESET}"
}

# -------------------------------
# host dependent colouring

#lprompt '[]' $PR_BRIGHT_BLACK $PR_WHITE
#rprompt '()' $PR_BRIGHT_BLACK $PR_WHITE

if [ $UID -eq 0 ]; then
    lprompt '<>' $PR_RED $PR_RED
    rprompt '<>' $PR_RED $PR_RED
else
    case $HOST in
        xenon)
            lprompt '[]' $PR_BRIGHT_BLACK $PR_GREEN
            rprompt '()' $PR_YELLOW $PR_WHITE
            ;;
        oxygen*)
            lprompt '[]' $PR_WHITE $PR_GREEN
            rprompt '()' $PR_YELLOW $PR_WHITE
            ;;

        meson*)
            lprompt '<>' $PR_RED $PR_YELLOW
            rprompt '<>' $PR_RED $PR_WHITE
            ;;
        muon*)
            lprompt '<>' $PR_RED $PR_BLUE
            rprompt '<>' $PR_RED $PR_WHITE
            ;;
        quark*)
            lprompt '<>' $PR_RED $PR_GREEN
            rprompt '<>' $PR_RED $PR_WHITE
            ;;

        *)
            if [ "${$(hostname -f)#*.}" = "jukie.net" ]; then
                lprompt '[]' $PR_BRIGHT_BLACK $PR_WHITE
                rprompt '()' $PR_BRIGHT_BLACK $PR_WHITE
            else
                lprompt '{}' $PR_WHITE $PR_WHITE
                rprompt '()' $PR_WHITE $PR_WHITE
            fi
            ;;
    esac
fi

if [ -n "$debian_chroot" ]; then
    PROMPT="$bgc%{$fg[yellow]%}%B${debian_chroot}%b ${PROMPT}"
fi

# ------------------------------
# update the vcs_info_msg_ magic variables, but only as little as possible

# This variable dictates weather we are going to do the git prompt update
# before printing the next prompt.  On some setups this saves 10s of work.
PR_GIT_UPDATE=1

# called before command excution
# here we decide if we should update the prompt next time
function zsh_git_prompt_preexec {
        case "$(history $HISTCMD)" in 
            *git*)
                PR_GIT_UPDATE=1
                ;;
        esac
}
preexec_functions+='zsh_git_prompt_preexec'

# called after directory change
# we just assume that we have to update git prompt
function zsh_git_prompt_chpwd {
        PR_GIT_UPDATE=1
}
chpwd_functions+='zsh_git_prompt_chpwd'

# called before prompt generation
# if needed, we will update the prompt info
function zsh_git_prompt_precmd {
       if [[ -n "$PR_GIT_UPDATE" ]] ; then
               vcs_info 'prompt'
               PR_GIT_UPDATE=
       fi
}
precmd_functions+='zsh_git_prompt_precmd'

# ------------------------------
# handle vi NORMAL/INSERT mode change
PR_VIMODE="#"
PR_VICOLOR=${PR_BLUE}
function zle-line-init zle-keymap-select {
        PR_VIMODE="${${KEYMAP/vicmd/¢}/(main|viins)/$}"
        PR_VICOLOR="${${KEYMAP/vicmd/${PR_RED}}/(main|viins)/${PR_GREEN}}"
        zle reset-prompt
}
zle -N zle-line-init
zle -N zle-keymap-select

# ------------------------------
# this stuff updates screen and xterm titles as the command runs

case $TERM in
    xterm* | rxvt* | urxvt*)
        function zsh_term_prompt_precmd {
                print -Pn "\e]0;%n@%m: %~\a"
        }
        function zsh_term_prompt_preexec {
                local x="${${${1//\"/\\\"}//\$/\\\\\$}//\%/%%}"
                print -Pn "\e]0;%n@%m: %~  $x\a"
        }
        preexec_functions+='zsh_term_prompt_preexec'
        precmd_functions+='zsh_term_prompt_precmd'
        ;;
    screen*)
        function zsh_term_prompt_precmd {
                print -nR $'\033k'"zsh"$'\033'\\\

                print -nR $'\033]0;'"zsh"$'\a'
        }
        function zsh_term_prompt_preexec {
                local x="${${${1//\"/\\\"}//\$/\\\\\$}//\%/%%}"
                print -nR $'\033k'"$x"$'\033'\\\

                print -nR $'\033]0;'"$x"$'\a'
        }
        preexec_functions+='zsh_term_prompt_preexec'
        precmd_functions+='zsh_term_prompt_precmd'
        ;;
esac