dotfiles/config/zsh/prompt.zsh

typeset -gA zfg=(
  [rst]='%{%f%}'

  [black]='%{%F{black}%}'
  [red]='%{%F{red}%}'
  [green]='%{%F{green}%}'
  [yellow]='%{%F{yellow}%}'
  [blue]='%{%F{blue}%}'
  [magenta]='%{%F{magenta}%}'
  [cyan]='%{%F{cyan}%}'
  [white]='%{%F{white}%}'

  # %F{...} only supports the 8 basic colors by name
  [gray]='%{%F{8}%}'
  [bright]='%{%F{15}%}'
  [faded]='%{%F{240}%}'
)

render_prompt() {
  setopt localoptions shortloops

  local sections=(
    "$(render_exit_code)"
    "$(render_user_host)"
    "$(render_pwd)"
    "$(render_git)"
    "$(render_venv)"
    "$(render_jobs)"
    "$(render_exec_time)"
  )

  local separator=" ${zfg[faded]}❯ "
  echo "${(@pj.$separator.)sections:#}${zfg[gray]}"
  (($SHLVL > 1)) && printf '󰅂%.0s' {2..$SHLVL}
  echo -n "❯ ${zfg[rst]}"
}

render_exit_code() {
  ((PROMPT_EXIT_CODE == 0)) && return

  if ((PROMPT_EXIT_CODE > 128 && PROMPT_EXIT_CODE < 160)); then
    print "%{%B%}${zfg[white]}$(kill -l $PROMPT_EXIT_CODE)%{%b%}"
  else
    print "%{%B%}${zfg[red]}✘ $PROMPT_EXIT_CODE%{%b%}"
  fi
}

render_user_host() {
  local parts=()

  # username in red if root, yellow if otherwise relevant
  if [[ $UID == 0 ]]; then
    parts+="${zfg[red]}%n"
  elif [[ $LOGNAME != $USER ]] || [[ -n $SSH_CONNECTION ]]; then
    parts+="${zfg[yellow]}%n"
  fi

  # hostname in yellow if relevant
  [[ -n $SSH_CONNECTION ]] && parts+="${zfg[yellow]}%m"

  (($#parts)) && {
    local separator="${zfg[gray]}@"
    print "${(pj:$separator:)parts}"
  }
}

render_pwd() {
  print "${zfg[cyan]}%~"
}

render_git() {
  local gitstatus  # local swallows git's exit code if not on its own line
  gitstatus=$(command git status --porcelain -b 2>/dev/null) || return

  # Sort through the status of files.
  local untracked=0 dirty=0 staged=0 conflicts=0 branch_line=''
  {
    while IFS='' read -r line; do
      case $line in
        \#\#*) branch_line=$line;;
        \?\?*) ((untracked++));;
        AA*|DD*|U?*|?U*) ((conflicts++));;
        *)
          [[ ${line:0:1} =~ '[MADRC]' ]] && ((staged++))
          [[ ${line:1:1} =~ '[MADRC]' ]] && ((dirty++))
          ;;
      esac
    done <<<$gitstatus
  }

  # Find out branch and upstream.
  local branch='' upstream='' ahead=0 behind=0
  {
    local fields=(${(s:...:)${branch_line#\#\# }})
    branch=$fields[1]
    local tracking=$fields[2]

    if [[ $branch == *'Initial commit on'* ]] \
    || [[ $branch == *'No commits yet on'* ]]; then
      # Branch name is last word in these possible branch lines:
      #   ## Initial commit on <branch>
      #   ## No commits yet on <branch>
      branch=${${(s: :)branch}[-1]}

    elif [[ $branch == *'no branch'* ]]; then
      # Dettached HEAD (also if a tag is checked out), branch line:
      #   ## HEAD (no branch)
      local tag=$(command git describe --tags --exact-match HEAD 2>/dev/null)
      if [[ -n $tag ]]; then
        branch="⚑$tag"
      else
        tag=$(command git describe --tags --long HEAD 2>/dev/null)
        if [[ -n $tag ]]; then
          branch="$tag"
        else
          branch="#$(command git rev-parse --short HEAD 2>/dev/null)"
        fi
      fi

    elif (($#fields > 1)); then
      # There is a tracking branch. Possibilites:
      #   ## <branch>...<upstream>
      #   ## <branch>...<upstream> [ahead 1]
      #   ## <branch>...<upstream> [behind 1]
      #   ## <branch>...<upstream> [ahead 1, behind 1]
      tracking=(${(s: [:)${tracking%]}})
      upstream=$tracking[1]
      if (($#tracking > 1)); then
        for e in ${(s:, :)tracking[2]}; do
          [[ $e == 'ahead '* ]]  && ahead=${e:6}
          [[ $e == 'behind '* ]] && behind=${e:7}
        done
      fi
    fi
  }

  local track_parts=()
  (($ahead > 0 ))  && track_parts+="${zfg[blue]}↑${ahead}"
  (($behind > 0 )) && track_parts+="${zfg[cyan]}${behind}↓"

  local state_parts=()
  (($staged > 0))  && state_parts+="${zfg[green]}+${staged}"
  (($dirty > 0 ))  && state_parts+="${zfg[red]}${dirty}✶"

  local separator="${zfg[gray]}"
  local gitinfo=("${zfg[blue]}${branch}")
  (($#track_parts > 0)) && gitinfo+="${(pj:$separator:)track_parts}"
  (($conflicts > 0 ))   && gitinfo+="${zfg[red]}${conflicts} "
  (($untracked > 0 ))   && gitinfo+="${zfg[white]}${untracked}?"
  (($#state_parts))     && gitinfo+="${(pj:$separator:)state_parts}"

  print "${(j: :)gitinfo}"
}

render_venv() {
    [[ -n "$VIRTUAL_ENV" ]] && print "${zfg[green]}${VIRTUAL_ENV:t}"
}

render_jobs() {
  (($PROMPT_JOB_COUNT > 0)) && print "${zfg[magenta]}%j bg"
}

render_exec_time() {
  (($PROMPT_EXEC_TIME <= 3)) && return  # don't print time if under 3s

  local parts=(
    "$((PROMPT_EXEC_TIME / 60 / 60 / 24))d"   # days
    "$((PROMPT_EXEC_TIME / 60 / 60 % 24))h"   # hours
    "$((PROMPT_EXEC_TIME / 60 % 60))m"        # minutes
    "$((PROMPT_EXEC_TIME % 60))s"             # seconds
  )
  print ${zfg[gray]}${parts:#0*}  # only keep non-zero parts
}

# Hook triggered when a command is about to be executed.
prompt_preexec_hook() {
  PROMPT_EXEC_START=$EPOCHSECONDS
}

# Hook triggered right before the prompt is drawn.
prompt_precmd_hook() {
  PROMPT_EXIT_CODE=$?  # this needs to be captured before anything else runs

  local stop=$EPOCHSECONDS
  local start=${PROMPT_EXEC_START:-$stop}
  PROMPT_EXEC_TIME=$((stop - start))
  unset PROMPT_EXEC_START  # needed because preexec is not always called

  local job_count='%j'; PROMPT_JOB_COUNT=${(%)job_count}
}

prompt_setup() {
  setopt NO_PROMPT_BANG PROMPT_CR PROMPT_PERCENT PROMPT_SP PROMPT_SUBST
  export PROMPT_EOL_MARK=''  # don't show % when a partial line is preserved
  export VIRTUAL_ENV_DISABLE_PROMPT=1  # we're doing it ourselves

  zmodload zsh/datetime   # so that $EPOCHSECONDS is available

  autoload -Uz add-zsh-hook
  add-zsh-hook precmd prompt_precmd_hook
  add-zsh-hook preexec prompt_preexec_hook

  PS1='$(render_prompt)'
}

prompt_setup