##
## zsh prompt themes extension
## by Adam Spiers <adam@spiers.net>
##
## Load with `autoload -U promptinit; promptinit'.
## Type `prompt -h' for help.
##

prompt_themes=()
typeset -gU prompt_themes
typeset -g prompt_theme >/dev/null

promptinit () {
  emulate -L zsh
  setopt extendedglob
  local ppath='' name theme

  # Autoload all prompt_*_setup functions in fpath
  for theme in $^fpath/prompt_*_setup(N); do
    if [[ $theme == */prompt_(#b)(*)_setup ]]; then
      name="$match[1]"
      if [[ -r "$theme" ]]; then
        prompt_themes=($name $prompt_themes)
        autoload -U prompt_${name}_setup
      else
        print "Couldn't read file $theme containing theme $name."
      fi
    else
      print "Eh?  Mismatch between glob patterns in promptinit."
    fi
  done

  # Color definitions come in handy
  autoload -U colors
  colors

  # Variables common to all prompt styles
  prompt_newline=$'\n%{\r%}'
}

set_prompt() {
  emulate -L zsh
  local opt preview theme usage old_theme

  usage='Usage: prompt <options>
Options:
    -c              Show currently selected theme and parameters
    -l              List currently available prompt themes
    -p [<themes>]   Preview given themes (defaults to all)
    -h [<theme>]    Display help (for given theme)
    -s <theme>      Set and save theme
    <theme>         Switch to new theme immediately (changes not saved)

Use prompt -h <theme> for help on specific themes.'

  getopts "chlps" opt
  case "$opt" in
    c) if (( $+prompt_theme )); then
         print -n "Current prompt theme"
         (( $#prompt_theme > 1 )) && print -n " with parameters"
         print " is:\n  $prompt_theme"
       else
         print "Current prompt is not a theme."
       fi
       return
       ;;
    h) if [[ -n "$2" && -n $prompt_themes[(r)$2] ]]; then
         if functions prompt_$2_help >/dev/null; then
           print "Help for $2 theme:\n"
           prompt_$2_help
         else
           print "No help available for $2 theme."
         fi
         print "\nType \`prompt -p $2' to preview the theme, \`prompt $2'"
         print "to try it out, and \`prompt -s $2' to use it in future sessions."
       else
         print "$usage"
       fi
       ;;
    l) print Currently available prompt themes:
       print $prompt_themes
       return
       ;;
    p) if [[ -z "$prompt_theme[1]" ]]; then
         # Not using a prompt theme; save settings
         prompt_non_theme=( "$PS1" "$PS2" "$PS3" "$PS4" "$RPS1" )
         prompt_old_precmd="$(functions precmd)"
         prompt_old_preexec="$(functions preexec)"
       fi
       preview=( $prompt_themes )
       (( $#* > 1 )) && preview=( "$@[2,-1]" )
       for theme in $preview; do
         theme_args=( "$=theme" )
         [[ "$theme" == "$prompt_theme[*]" ]] && continue
         if [[ -z "$prompt_themes[(r)$theme_args[1]]" ]]; then
           print "\nUnknown theme: $theme_args[1]"
           continue
         fi
         print

         # The next line is a bit ugly.  It (perhaps unnecessarily)
         # runs the prompt theme setup function to ensure that if
         # the theme has a _preview function that it's been autoloaded.
         prompt_${theme_args[1]}_setup

         if functions prompt_${theme_args[1]}_preview >&/dev/null; then
           prompt_${theme_args[1]}_preview "${(@)theme_args[2,-1]}"
         else
           prompt_preview_theme "${(@)theme_args}"
         fi
       done
       print
       if [[ -z "$prompt_theme[1]" ]]; then
         PS1="$prompt_non_theme[1]"
         PS2="$prompt_non_theme[2]"
         PS3="$prompt_non_theme[3]"
         PS4="$prompt_non_theme[4]"
         RPS1="$prompt_non_theme[5]"
         if [[ -z "$prompt_old_precmd" ]]; then
           precmd () { }
         else
           eval "$prompt_old_precmd"
         fi
         if [[ -z "$prompt_old_preexec" ]]; then
           preexec () { }
         else
           eval "$prompt_old_preexec"
         fi
       else
         prompt_${prompt_theme[1]}_setup "${(@)prompt_theme[2,-1]}"
       fi
       ;;
    s) print "Set and save not yet implemented.  Please ensure your ~/.zshrc"
       print "contains something similar to the following:\n"
       print "  autoload -U promptinit"
       print "  promptinit"
       print "  prompt $*[2,-1]"
       ;;
    *) if [[ "$1" == 'random' ]]; then
         local random_themes
         if (( $#* == 1 )); then
           random_themes=( $prompt_themes )
         else
           random_themes=( "$@[2,-1]" )
         fi
         local i=$(( ( $RANDOM % $#random_themes ) + 1 ))
         argv=( "${=random_themes[$i]}" )
       fi
       if [[ -z "$1" || -z $prompt_themes[(r)$1] ]]; then
         print "$usage"
         return
       fi
       prompt_$1_setup "$@[2,-1]"
       prompt_theme=( "$@" )

       # Avoid screwing up the environment listing
       PSZZZZ=$reset_color
       RPSZZZZ=$reset_color
       PROMPTZZZZ=$reset_color
       RPROMPTZZZZ=$reset_color
       promptzzzz=$reset_color
       ;;
  esac
}

prompt () {
  local prompt_opts
  
  set_prompt "$@"
 
  (( $#prompt_opts )) &&
      setopt noprompt{bang,cr,percent,subst} prompt${^prompt_opts[@]}

  true
}

prompt_preview_theme () {
  print -n "$1 theme"
  (( $#* > 1 )) && print -n " with parameters \`$*[2,-1]'"
  print ":"
  prompt_${1}_setup "$@[2,-1]"
  precmd
  print -n -P "${PS1}"
  preexec
  print "command arg1 arg2 ... argn"
}

promptinit "$@"
