zsh-theme: add multi-line support for long prompts, cleanup

note: this is a disaster

brethil.zsh-theme

@@ -5,15 +5,37 @@
 # 255: white
 # See the prompt expansion section in `info zsh` for more info
 
-is_byobu="$(env | grep BYOBU)"
-local USER_PROMPT="%(#.%K{160}%F{255}[root]%k%f.)" # root printed on a red background # TODO: add lightning here
-local JOB_PROMPT="%(1j.%B%K{202}%F{220} %j %k%b.)" # Shows jobs number on an orange background if there are background jobs
-if [[ -n "$is_byobu" ]]; then
-    local SHLVL_PROMPT="%(3L.%K{161}%F{255}%B $((SHLVL-1)) %f%b%k.)" # Shows SHLVL on a magenta background if SHLVL > 1 
-else 
-    local SHLVL_PROMPT="%(2L.%K{161}%F{255}%B %L %f%b%k.)" # Shows SHLVL on a magenta background if SHLVL > 1 
-fi
-local PATH_PROMPT="%F{white}[%F{green}%(4~:…/%3~:%~)%F{white}]%f" #  Shows last 3 items in path if there are more than 4
+function user_prompt(){
+  # root printed on a red background
+  local USER_PROMPT="%(#.%K{160}%F{255}[root]%k%f.)"
+  echo "$USER_PROMPT"
+  # if [[ -n $SSH_CONNECTION ]]; then
+  #     local host="${purple}%B%m%b${default}"
+  #     host="%U${host}%u"
+  #     echo "$host"
+  # fi
+}
+
+function job_prompt() {
+  local JOB_PROMPT="%(1j.%B%K{202}%F{220} %j %k%b.)" # Shows jobs number on an orange background if there are background jobs
+  echo "$JOB_PROMPT"
+}
+
+function shlvl_prompt() {
+  local SHLVL_PROMPT
+  if [[ -n "$BYOBU_PREFIX" ]]; then
+    SHLVL_PROMPT="%(3L.%K{161}%F{255}%B $((SHLVL-1)) %f%b%k.)" # Shows SHLVL on a magenta background if SHLVL > 1
+  else
+    SHLVL_PROMPT="%(2L.%K{161}%F{255}%B %L %f%b%k.)" # Shows SHLVL on a magenta background if SHLVL > 1
+  fi
+  echo "$SHLVL_PROMPT"
+}
+
+function path_prompt(){
+  # Shows last 3 items in path if there are more than 4
+  local PATH_PROMPT="%F{white}[%F{green}%(4~:…/%3~:%~)%F{white}]%f"
+  echo "$PATH_PROMPT"
+}
 
 # PREFIX/SUFFIX are added before/after `git_prompt_info`
 ZSH_THEME_GIT_PROMPT_PREFIX="%f%{$fg[yellow]%} "
@@ -67,20 +89,38 @@ function git_prompt() {
 function check_last_exit_code() {
   local LAST_EXIT_CODE=$?
   if [[ $LAST_EXIT_CODE -ne 0 ]]; then
-    local EXIT_CODE_PROMPT=' '
-    # EXIT_CODE_PROMPT+="%{$fg[red]%}-%{$reset_color%}"
-    # EXIT_CODE_PROMPT+="%{$fg_bold[red]%}$LAST_EXIT_CODE%{$reset_color%}"
+    local EXIT_CODE_PROMPT
     EXIT_CODE_PROMPT+="%B%F{88}-$LAST_EXIT_CODE-%f%b"
-    # EXIT_CODE_PROMPT+="%{$fg[red]%}-%{$reset_color%}"
     echo "$EXIT_CODE_PROMPT"
   fi
 }
 
-local return_code="%(?.."'$(check_last_exit_code)'")" # for RPS1
-local ret_status="%(?:%B%F{28} →%f%b:%{$fg[red]%} ✗%{$reset_color%})" # for PS1
+function return_code_RPS1() {
+  # For right-side prompt
+  local return_code="%(?..$(check_last_exit_code))" # for RPS1
+  echo "$return_code"
+}
 
+function ret_status() {
+  local ret_status="%(?:%B%F{28} →%f%b:%{$fg[red]%} x%{$reset_color%})"
+  echo "$ret_status"
+}
+
+_PROMPT_PROTO='$(virtualenv_info)$(user_prompt)$(path_prompt)$(git_prompt)$(job_prompt)$(shlvl_prompt)'
+
+function prompt_too_long(){
+  local zero='%([BSUbfksu]|([FK]|){*})'
+  local stripped="$(echo ${(S%%)_PROMPT_PROTO//$~zero/} | sed "s,\x1B\[[0-9;]*[a-zA-Z],,g" ) "
+  local prompt_len=${#stripped}
+  local max_len=$(($COLUMNS/2))
+  if [[ $prompt_len -ge $max_len ]]; then
+    echo true
+  fi
+}
+PROMPT="$(echo "$_PROMPT_PROTO")"
+NEWLINE=$'\n'
+PROMPT+='${$(prompt_too_long)/true/${NEWLINE}}'
+PROMPT+='$(ret_status) '
 
-# Prompt looks like this: <virtualenv> [user][path](git)[jobn][shlvl] retstatus     
-PROMPT='$(virtualenv_info)'"${USER_PROMPT}${PATH_PROMPT}"'$(git_prompt)'"${JOB_PROMPT}${SHLVL_PROMPT}${ret_status} "
 # Right prompt is just return code and time
-RPS1="${return_code} %F{255}[%*]%f" 
+RPS1='$(return_code_RPS1) %F{255}[%*]%f'