# Setup the prompt with pretty colors
setopt prompt_subst
autoload colors; colors;

hostcolor=green
[[ $(hostname) == "tardis" ]]  && hostcolor=red

precmd() {
  PROMPT='%T %n@%{$fg[$hostcolor]%}%m%(1j. $fg_bold[white]↵%{$fg_bold[red]%}%j.)%{$reset_color%}%-0>..>$(git_prompt_status)%>>
%(?..%{$fg_bold[white]%}%?)%{$reset_color%}$(vi_prompt_info)%{%(!.$fg[red]❰.$fg[green]❱)%1G%} '
  RPROMPT='%{$fg[green]%}%~%{$reset_color%}'
}

vi_prompt_info() {
  local vicmd="$fg_bold[green]❰$reset_color%1G"
  local viins="$fg_bold[blue]❱$reset_color%1G"
  local map=${KEYMAP:-viins}
  printf '%s' "%{${${map/vicmd/$vicmd}/(main|viins)/$viins}%}"
}

function zle-line-init zle-line-finish zle-keymap-select {
  zle reset-prompt
  zle -R
}

# zle -N zle-line-init
# zle -N zle-line-finish
zle -N zle-keymap-select

# reset zle on resize
TRAPWINCH() {
  zle && { zle reset-prompt; zle -R  }
}

# Get the status of the working tree
git_prompt_status() {
  local branch ahead behind 
  local added deleted modified renamed unmerged untracked dirty
  # Use porcelain status for easy parsing.
  local status_cmd="git status --porcelain -b"

  # Get current status.
  while IFS=$'\n' read line; do
    if [[ "$line" == \#\#\ * ]]; then
      [[ "$line" =~ '## (.*?)(\.\.\.|$)' ]] && branch=$match[1]
      [[ "$line" =~ 'ahead ([0-9]+)'  ]]  && ahead=$match[1]
      [[ "$line" =~ 'behind ([0-9]+)'  ]] && behind=$match[1]
    else
      # Count added, deleted, modified, renamed, unmerged, untracked, dirty.
      # T (type change) is undocumented, see http://git.io/FnpMGw.
      # index
      [[ "$line" == M[\ MTD]\ * ]]      && (( updated++ ))
      [[ "$line" == [AC][\ MTD]\ * ]]   && (( added++ ))
      [[ "$line" == D[\ MT]\ * ]]       && (( deleted++ ))
      [[ "$line" == R[\ MTD]\ * ]]      && (( renamed++ ))

      # work tree
      [[ "$line" == [\ MARCT]M\ * ]]    && (( modified++ ))
      [[ "$line" == [\ MARCT]D\ * ]]    && (( deleted_wt++ ))
      [[ "$line" == \?\?\ * ]]          && (( untracked++ ))

      # morge conflicts
      [[ "$line" == (AA|DD|U?|?U)\ * ]] && (( unmerged++ ))
    fi
  done < <(${(z)status_cmd} 2> /dev/null)

  local git_status=" %{$fg[yellow]%}"

  # Format branch
  if [[ -n $branch ]]; then
    git_status+="$branch"
  else
    git_status+="$(git rev-parse --short HEAD 2> /dev/null)"
    [[ $? -ne 0 ]] && return
  fi

  # Format upstream
  local upstream_str
  (( ahead  > 0 )) && upstream_str+="%{$fg[blue]%} >$ahead"
  (( behind > 0 )) && upstream_str+="%{$fg[blue]%} <$behind"
  git_status+="$upstream_str"

  # Format stashed
  stashed=$(git stash list | wc -l)
  if (( stashed > 0 )) then
    stashed_str+="%{$fg_bold[cyan]%} ⋎$stashed%{$reset_color%}"
  fi
  git_status+="$stashed_str"

  # Format index
  local index_str
  (( updated > 0 )) && index_str+="%{$fg[green]%} *$updated"
  (( added   > 0 )) && index_str+="%{$fg[green]%} +$added"
  (( deleted > 0 )) && index_str+="%{$fg[green]%} -$deleted"
  (( renamed > 0 )) && index_str+="%{$fg[green]%} ≈$renamed"
  git_status+="$index_str"

  # Format working tree
  local wt_str
  (( modified   > 0 )) && wt_str+="%{$fg[red]%} *$modified"
  (( deleted_wt > 0 )) && wt_str+="%{$fg[red]%} -$deleted_wt"
  (( untracked  > 0 )) && wt_str+="%{$fg[red]%} +$untracked"
  (( unmerged   > 0 )) && wt_str+="%{$fg[magenta]%} ♒$unmerged"
  git_status+="$wt_str"

  git_status+="%{$reset_color%}"

  echo $git_status
}

function print_if_fits() {
  local zero length

  zero='%([BSUbfksu]|([FB]|){*})'
  length=${#${(S%%)1//$~zero/}}
  echo "%-$length(l.$1.)"
}