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}%}'
)

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

fs-prompt-render-full() {
  setopt localoptions shortloops

  local separator="${zfg[faded]} ❯ "
  local sections=(
    "$(fs-prompt-exit-code)"
    "$(fs-prompt-user-and-hostname)"
    "$(fs-prompt-pwd)"
    "$(fs-prompt-git)"
    "$(fs-prompt-virtualenv)"
    "$(fs-prompt-jobs)"
    "$(fs-prompt-exec-time)"
  )
  echo "\n${(@pj.$separator.)sections:#}${zfg[rst]}"

  if ((PROMPT_EXIT_CODE == 0)); then
    echo -n "${zfg[faded]}"
  elif ((PROMPT_EXIT_CODE > 128 && PROMPT_EXIT_CODE < 160)); then
    echo -n "${zfg[white]}"
  else
    echo -n "${zfg[red]}"
  fi
  (($SHLVL > 1)) && printf '󰅂%.0s' {2..$SHLVL}
  echo -n "❯ ${zfg[rst]}"
}

fs-prompt-render-compact() {
  echo -n "${zfg[magenta]}❯ ${zfg[rst]}"
}

fs-prompt-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
}

fs-prompt-user-and-hostname() {
  local parts=()

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

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

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

fs-prompt-pwd() {
  print "${zfg[cyan]}${icons[folder]}%~"
}

fs-prompt-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="${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 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}"
}

fs-prompt-virtualenv() {
    [[ -n "$VIRTUAL_ENV" ]] && print "${zfg[green]}${icons[python]}${VIRTUAL_ENV:t}"
}

fs-prompt-jobs() {
  (($PROMPT_JOB_COUNT > 0)) && print "${zfg[magenta]}${icons[background]}%j"
}

fs-prompt-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]}${icons[clock]}${parts:#0*}  # only keep non-zero parts
}

# 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)'
}

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

  autoload -Uz add-zle-hook-widget
  # add-zle-hook-widget line-init   fs-prompt-zle-line-init
  add-zle-hook-widget line-finish fs-prompt-zle-line-finish
}

fs-setup-prompt