typeset -gA _fs_prompt=(
  [is_transient]=true
  [transient_color]='%{%F{magenta}%}'
)

typeset -gA icons=(
  [user]=" "
  [folder]=" "
  [branch]=" "
  [tag]="⚑ "
  [detached]=" "
  [python]=" "
  [background]=" "
  [clock]=" "
)

fs_prompt_render_full() {
  local black='%{%F{0}%}' red='%{%F{1}%}'       green='%{%F{2}%}'   yellow='%{%F{3}%}' \
        blue='%{%F{4}%}'  magenta='%{%F{5}%}'   cyan='%{%F{6}%}'    white='%{%F{7}%}' \
        gray='%{%F{8}%}'  orange='%{%F{166}%}'  faded='%{%F{240}%}' \
        reset='%{%f%}'

  local sep_color="$faded" signal_color="$white" error_color="$red" reset="$reset"

  local sections=(
    "$(fs_prompt_exit_code $error_color $signal_color)"
    #                               root  user     separator
    "$(fs_prompt_user_and_hostname  $red  $yellow  $faded)"
    "$(fs_prompt_pwd $cyan)"
    #                branch  ahead  behind  untracked  staged  dirty    conflict  separator
    "$(fs_prompt_git $blue   $blue  $cyan   $white     $green  $orange  $red      $gray)"
    "$(fs_prompt_virtualenv $yellow)"
    "$(fs_prompt_jobs $magenta)"
    "$(fs_prompt_exec_time $gray)"
  )
  local sep="$sep_color ❯ "
  echo "\n${(@pj:$sep:)sections:#}$reset"

  if ((PROMPT_EXIT_CODE == 0)); then
    echo -n "$sep_color"
  elif ((PROMPT_EXIT_CODE > 128 && PROMPT_EXIT_CODE < 160)); then
    echo -n "$signal_color"
  else
    echo -n "$error_color"
  fi
  (($SHLVL > 1)) && printf '󰅂%.0s' {2..$SHLVL}
  echo -n "❯ $reset"
}

fs_prompt_render_compact() {
  echo -n "${_fs_prompt[transient_color]}❯ %{%f%}"
}

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

  local error_color="$1" signal_color="$2" bold="%{%B%}" no_bold="%{%b%}"

  if ((PROMPT_EXIT_CODE > 128 && PROMPT_EXIT_CODE < 160)); then
    print "${bold}${signal_color}$(kill -l $PROMPT_EXIT_CODE)${no_bold}"
  else
    print "${bold}${error_color}  ${PROMPT_EXIT_CODE}${no_bold}"
  fi
}

fs_prompt_user_and_hostname() {
  local root_color="$1" user_color="$2" at_color="${3:-$2}"
  local parts=()

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

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

  (($#parts)) && {
    local at="${at_color}@"
    print "${(pj:$at:)parts}"
  }
}

fs_prompt_pwd() {
  local color="$1"
  print "${color}${icons[folder]}%~"
}

fs_prompt_git() {
  local branch_color="$1"
  local ahead_color="$2"
  local behind_color="$3"
  local untracked_color="$4"
  local staged_color="$5"
  local dirty_color="$6"
  local conflict_color="$7"
  local sep_color="$8"

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

gitstatus="## some-longer-name...origin/tracking [ahead 5, behind 12]
M  staged_1
M  staged_2
 M dirty_1
 M dirty_2
 M dirty_3
DD conflict
?? untracked
"

  # 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="${icons[branch]}$fields[1]"
    local tracking=$fields[2]

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

    elif [[ $branch == *'no branch'* ]]; then
      # Dettached HEAD (also if a tag is checked out), branch line:
      #   ## HEAD (no branch)
      local icon="${icons[tag]}"
      local ref=$(command git describe --tags --exact-match HEAD 2>/dev/null)
      if [[ -z $ref ]]; then
        icon="${icons[detached]}"
        ref=$(command git describe --tags --long HEAD 2>/dev/null)
        [[ -n $ref ]] || ref=$(command git rev-parse --short HEAD 2>/dev/null)
      fi
      branch="${icon}%{%B%}${ref}%{%b%}"

    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 result=()
  {
    local sep="${sep_color}"
    local tracking=()
    (($ahead))      && tracking+="${ahead_color}${ahead}"      # ↑
    (($behind))     && tracking+="${behind_color}${behind}"    # ↓
    (($#tracking))  && result+="${(pj:$sep:)tracking} "

    (($conflicts))  && result+="${conflict_color} %{%B%}${conflicts}%{%b%}"

    local state=()
    (($staged))     && state+="${staged_color}+${staged}"
    (($dirty))      && state+="${dirty_color}${dirty}✶"
    (($#state))     && result+="${(pj:$sep:)state}"

    (($untracked))  && result+="${untracked_color}${untracked}?"
  }
  print "${branch_color}${branch} ${(j: :)result}"
}

fs_prompt_virtualenv() {
  [[ -n "$VIRTUAL_ENV" ]] || return
  local color="$1"
  print "${color}${icons[python]}${VIRTUAL_ENV:t}"
}

fs_prompt_jobs() {
  (($PROMPT_JOB_COUNT)) || return
  local color="$1"
  print "${color}${icons[background]}%j"
}

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

  local color="$1"
  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
  )
  local exec_time="${(@)parts:#0*}"           # only keep non-zero parts
  [[ -n "$exec_time" ]] && print "${color}${icons[clock]}${exec_time}"
}

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

# Hook triggered right before the prompt is drawn.
fs_prompt_precmd() {
  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}

  PS1='$(fs_prompt_render_full)'
}

# This hook is only used if transient prompt is enabled.
fs_prompt_zle_line_finish() {
  PS1='$(fs_prompt_render_compact)'
  zle reset-prompt
}

fs_setup_prompt() {
  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 preexec  fs_prompt_preexec
  add-zsh-hook precmd   fs_prompt_precmd

  if ${_fs_prompt[is_transient]}; then
    autoload -Uz add-zle-hook-widget
    add-zle-hook-widget line-finish fs_prompt_zle_line_finish
    PS2="${_fs_prompt[transient_color]}%_❯%{%f%} "
  else
    add-zle-hook-widget -d line-finish fs_prompt_zle_line_finish
  fi
}

fs_setup_prompt