# -*- shell-script -*-

autoload -Uz vcs_info
zstyle ':vcs_info:*' enable bzr darcs fossil git hg svn
zstyle ':vcs_info:*' formats '%F{4}{%F{2}%s%F{4}:%F{2}%b%u%c%F{4}}─%f'
zstyle ':vcs_info:*' actionformats '%F{4}{%F{2}%s%F{4}:%F{2}%b%u%c%F{4}|%F{2}%a%F{4}}─%f'
zstyle ':vcs_info:*' check-for-changes true
zstyle ':vcs_info:*' stagedstr '+'
zstyle ':vcs_info:*' unstagedstr '!'

function precmd {
    vcs_info

    local TERMWIDTH
    (( TERMWIDTH = ${COLUMNS} - 1 ))


    ###
    # Truncate the path if it's too long.

    PR_FILLBAR=""
    PR_PWDLEN=""

    local promptsize
    local pwdsize
    if [[ $ZSH_VERSION > 3.0.6 ]]; then
        promptsize=${#${(%):---(%n@%m:%l)---()--}}
        pwdsize=${#${(%):-%~}}
    else
        promptsize=${#${:---($USER@${HOST%%.*}:${TTY#/dev/})---()--}}
        pwdsize=${#$(pwd | perl -pe "s{$HOME}{~}")}
    fi

    if [[ "$promptsize + $pwdsize" -gt $TERMWIDTH ]]; then
            ((PR_PWDLEN=$TERMWIDTH - $promptsize))
    else
        PR_FILLBAR="\${(l.(($TERMWIDTH - ($promptsize + $pwdsize)))..${PR_HBAR}.)}"
    fi


    ###
    # Get APM info.

    if [[ -x "$(whence -p ibam)" ]]; then
        PR_APM_RESULT=$(ibam --percentbattery)
    elif [[ -x "$(whence -p apm)" ]]; then
        PR_APM_RESULT=$(apm)
    fi

    ###
    # Show virtual environment

    PR_VENV=""
    if [[ -n "$POETRY_ACTIVE" ]]; then
        # I'm not sure how to tell _which_ Poetry venv is running, so that
        # information gets left out.  Sigh.
        PR_VENV="${PR_VENV}${PR_BLUE}[${PR_GREEN}poetry${PR_BLUE}]${PR_SHIFT_IN}${PR_HBAR}${PR_SHIFT_OUT}"
    fi
    if [[ -n "$CONDA_DEFAULT_ENV" ]]; then
        PR_VENV="${PR_VENV}${PR_BLUE}[${PR_GREEN}conda${PR_BLUE}:${PR_GREEN}${CONDA_DEFAULT_ENV}${PR_BLUE}]${PR_SHIFT_IN}${PR_HBAR}${PR_SHIFT_OUT}"
    fi
    if [[ -n "$VIRTUAL_ENV" ]]; then
        PR_VENV="${PR_VENV}${PR_BLUE}[${PR_GREEN}venv${PR_BLUE}:${PR_GREEN}${VIRTUAL_ENV/#${HOME}/~}${PR_BLUE}]${PR_SHIFT_IN}${PR_HBAR}${PR_SHIFT_OUT}"
    fi
}


setopt extended_glob
preexec () {
    case "$TERM" in
        screen*|tmux*)
            local CMD=${1[(wr)^(<*|*=*|sudo|exec|-*)]}
            if [[ $(hostname -d) == 'work.com' ]]; then
                CMD="$(hostname -s):$CMD"
            fi
            echo -n "\ek$CMD\e\\"
            ;;
    esac
}


setprompt () {
    ###
    # Need this so the prompt will work.

    setopt prompt_subst


    ###
    # Reset some terminal properties.

    PR_RESET_TERM=""
    if [[ $ZSH_VERSION < 5.1 ]]; then
        # zsh versions before 5.1 don't understand bracketed paste mode,
        # but sometimes a program that _does_ will exit without disabling
        # it.  This ensures it's turned off so it doesn't interfere with
        # pasting into the terminal.
        PR_RESET_TERM="$PR_RESET_TERM"$'%{\e[?2004l%}'
    fi


    ###
    # See if we can use colors.

    autoload colors zsh/terminfo
    if [[ "$terminfo[colors]" -ge 8 ]]; then
        colors
    fi
    for color in RED GREEN YELLOW BLUE MAGENTA CYAN WHITE; do
        eval PR_$color='%{$fg[${(L)color}]%}'
        eval PR_LIGHT_$color='%{$fg[${(L)color}]%}'
        (( count = $count + 1 ))
    done
    PR_NO_COLOUR="%{$terminfo[sgr0]%}"


    ###
    # See if we can use extended characters to look nicer.

    if [[ "$(locale charmap)" = "UTF-8" ]]; then
        PR_SET_CHARSET=""
        PR_SHIFT_IN=""
        PR_SHIFT_OUT=""
        PR_HBAR="─"
        PR_ULCORNER="┌"
        PR_LLCORNER="└"
        PR_LRCORNER="┘"
        PR_URCORNER="┐"
    else
        if [[ $ZSH_VERSION > 3.0.6 ]]; then
            typeset -A altchar
        fi
        set -A altchar ${(s..)terminfo[acsc]}
        PR_SET_CHARSET="%{$terminfo[enacs]%}"
        PR_SHIFT_IN="%{$terminfo[smacs]%}"
        PR_SHIFT_OUT="%{$terminfo[rmacs]%}"
        PR_HBAR=${altchar[q]:--}
        PR_ULCORNER=${altchar[l]:--}
        PR_LLCORNER=${altchar[m]:--}
        PR_LRCORNER=${altchar[j]:--}
        PR_URCORNER=${altchar[k]:--}
    fi


    ###
    # Decide if we need to set titlebar text.

    # Note: old screen definition was:
    #
    # PR_TITLEBAR=$'%{\e_screen \005 (\005t) | %(0#.-=[ROOT]=- | .)%n@%m:%~ | ${COLUMNS}x${LINES} | %y\e\\%}'
    #
    # But I've switched to tmux _and_ some systems need a $TERM of
    # "screen" to work, with no other way to tell whether the command is
    # running in screen or tmux.  So I just assume that "screen" really
    # means tmux now.  If I do still need to run screen, I'll need to
    # update screen's hstatus stuff to match the way tmux does it.
    case $TERM in
        xterm*|mlterm|konsole*|gnome-*|vte*)
            PR_TITLEBAR=$'%{\e]0;%(0#.-=*[ROOT]*=- | .)%n@%m:%~ | ${COLUMNS}x${LINES} | %y\a%}'
            ;;
        screen*|tmux*)
            PR_TITLEBAR=$'%{\e_%(0#.-=[ROOT]=- | .)%n@%m:%~ | ${COLUMNS}x${LINES} | %y\e\\%}'
            ;;
        *)
            PR_TITLEBAR=''
            ;;
    esac


    ###
    # Decide whether to set a screen title
    case $TERM in
        screen*|tmux*)
            if [[ $(hostname -d) == 'work.com' ]]; then
                PR_STITLE=$'%{\ek%m:zsh\e\\%}'
            else
                PR_STITLE=$'%{\ekzsh\e\\%}'
            fi
            ;;
        *)
            PR_STITLE=''
            ;;
    esac


    ###
    # APM detection

    case "$OSTYPE" in
        linux*)
            if [[ -x "$(whence -p ibam)" ]]; then
                PR_APM='$PR_RED${${PR_APM_RESULT[(f)1]}[(w)-2]}%%(${${PR_APM_RESULT[(f)3]}[(w)-1]})$PR_LIGHT_BLUE:'
            elif [[ -x "$(whence -p apm)" ]]; then
                PR_APM='$PR_RED${PR_APM_RESULT[(w)5,(w)6]/\% /%%}$PR_LIGHT_BLUE:'
            else
                PR_APM=''
            fi
            ;;
        *)
            PR_APM=''
            ;;
    esac

    ###
    # Highlight certain machine names.
    case ${(%):-%m} in
        cermanius|capsule|cantina)
            HOST_COLOR=$PR_RED
            ;;
        *)
            HOST_COLOR=$PR_GREEN
            ;;
    esac

    ###
    # Finally, the prompt.

    PROMPT='$PR_RESET_TERM$PR_SET_CHARSET$PR_STITLE${(e)PR_TITLEBAR}\
$PR_CYAN$PR_SHIFT_IN$PR_ULCORNER$PR_BLUE$PR_HBAR$PR_SHIFT_OUT(\
$PR_GREEN%(0#.%SROOT%s.%n)$PR_GREEN@$HOST_COLOR%m$PR_GREEN:%l\
$PR_BLUE)$PR_SHIFT_IN$PR_HBAR$PR_CYAN$PR_HBAR${(e)PR_FILLBAR}$PR_BLUE$PR_HBAR$PR_SHIFT_OUT(\
$PR_MAGENTA%$PR_PWDLEN<...<%~%<<\
$PR_BLUE)$PR_SHIFT_IN$PR_HBAR$PR_CYAN$PR_URCORNER$PR_SHIFT_OUT\

$PR_CYAN$PR_SHIFT_IN$PR_LLCORNER$PR_BLUE$PR_HBAR$PR_SHIFT_OUT(\
%(?..$PR_LIGHT_RED%?$PR_BLUE:)\
${(e)PR_APM}$PR_YELLOW%D{%H:%M}\
$PR_LIGHT_BLUE:%(0#.$PR_RED.$PR_WHITE)%#$PR_BLUE)$PR_SHIFT_IN$PR_HBAR$PR_SHIFT_OUT\
$PR_CYAN$PR_SHIFT_IN$PR_HBAR$PR_SHIFT_OUT\
$PR_NO_COLOUR '

    RPROMPT=' $PR_SHIFT_IN$PR_CYAN$PR_HBAR$PR_BLUE$PR_HBAR$PR_SHIFT_OUT${vcs_info_msg_0_}$PR_VENV\
$PR_BLUE($PR_YELLOW%D{%a,%b%d}$PR_BLUE)$PR_SHIFT_IN$PR_HBAR$PR_CYAN$PR_LRCORNER$PR_SHIFT_OUT$PR_NO_COLOUR'

    PS2='$PR_CYAN$PR_SHIFT_IN$PR_HBAR$PR_SHIFT_OUT\
$PR_BLUE$PR_SHIFT_IN$PR_HBAR$PR_SHIFT_OUT(\
$PR_LIGHT_GREEN%_$PR_BLUE)$PR_SHIFT_IN$PR_HBAR$PR_SHIFT_OUT\
$PR_CYAN$PR_SHIFT_IN$PR_HBAR$PR_SHIFT_OUT$PR_NO_COLOUR '
}

setprompt