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

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{202}%}'  faded='%{%F{240}%}' \
        reset='%{%f%}'

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

  #
  # First line
  #
  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"

  #
  # Second line
  #
  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"
}

# This is only used if transient prompt is enabled.
fs_prompt_render_compact() {
  echo -n "${_fs_prompt[transient_color]}❯ %{%f%}"
}

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

  local error_icon="$1" signal_icon="$2" error_color="$3" signal_color="$4"

  if ((PROMPT_EXIT_CODE > 128 && PROMPT_EXIT_CODE < 160)); then
    print "${signal_color}${signal_icon} $(kill -l $PROMPT_EXIT_CODE)"
  else
    print "%{%B%}${error_color}${error_icon} ${PROMPT_EXIT_CODE}%{%b%}"
  fi
}

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

  # username in red if root, yellow if otherwise relevant
  local user_name="${icon} %n"
  if [[ $UID == 0 ]]; then
    parts+="${root_color}${user_name}"
  elif [[ $LOGNAME != $USER ]] || [[ -n $SSH_CONNECTION ]]; then
    parts+="${user_color}${user_name}"
  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 icon="$1" color="$2"
  print "${color}${icon} %~"
}

fs_prompt_git() {
  typeset -A icons=([branch]="$1" [tag]="$2" [detached]="$3")
  typeset -A fg=(
    [branch]="$4"
    [ahead]="$5"
    [behind]="$6"
    [untracked]="$7"
    [staged]="$8"
    [dirty]="$9"
    [conflict]="$10"
    [sep]="$11"
  )

  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="${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="${fg[sep]}"
    local tracking=()
    (($ahead))      && tracking+="${fg[ahead]}⇡${ahead}"      # ↑ ⇡
    (($behind))     && tracking+="${fg[behind]}${behind}⇣"    # ↓ ⇣
    (($#tracking))  && result+="${(pj:$sep:)tracking}"

    (($conflicts))  && result+="${fg[conflict]} %{%B%}${conflicts}%{%b%}"

    local state=()
    (($staged))     && state+="${fg[staged]}+${staged}"
    (($dirty))      && state+="${fg[dirty]}${dirty}✶"
    (($#state))     && result+="${(pj:$sep:)state}"

    (($untracked))  && result+="${fg[untracked]}${untracked}?"
  }
  print "${fg[branch]}${branch} ${(j: :)result}"
}

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

fs_prompt_jobs() {
  (($PROMPT_JOB_COUNT)) || return
  local icon="$1" color="$2"
  print "${color}${icon} %j"
}

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

  local icon="$1" color="$2"
  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}${icon} ${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