#compdef module

# Main dispatcher
_module()
{

  _arguments \
    '(-h --help)'{--help,-h}'[display usage info]'                                     	\
    '(- *)'{--version,-v}'[module command version]'                                	\
    '(-A --all)'{--all,-A}'[list, avail and spider will report hidden modules]'         \
    '(-D --debug)'{--debug,-D}'[Send program tracing information to stderr]'           	\
    '(-I --ignore_cache)'{--ignore_cache,-I}'[Treat the cache file(s) as out-of-date]'  \
    '(-T --trace)'{--trace,-T}'[Report major internal actions by Lmod]'                 \
    '(-d --default)'{--default,-d}'[List default modules only when used with avail]'   	\
    '(-f --force)'{--force,-f}'[Force removal of a sticky module or save an empty collection]' \
    '(-q --quiet)'{--quiet,-q}'[disable verbose messages]'                             	\
    '(-r --regexp)'{--regexp,-r}'[lua regular expression matching]'                     \
    '(-s --style)'{--style,-s}'[avail output style]'                                    \
    '(-t --terse)'{--terse,-t}'[display avail, list and spider output in short format]'	\
    '(-w --width)'{--width,-w}'[use given width for max term width]'                  	\
    '--config[display full Lmod configuration]'                                         \
    '--config_json[display full Lmod configuration in json format]'                     \
    "--dumpname[Report the string Lmod]"                                                \
    "--dumpversion[Report Lmod's version]"                                              \
    "--gitversion[Report Lmod's git version]"                                           \
    '--latest[Load latest (ignore default)]'                                            \
    '--location[Print the file location only when using show]'                          \
    '--miniConfig[display configuration differences from default]'                      \
    '--mt[display the module table]'                                                    \
    '--no_redirect[send messages to stderr]'                                            \
    '--pin_version[pin versions when restoring a collection]'                           \
    '--pod[generate pod format for manpage generation]'                                 \
    '--raw[Print raw contents of a modulefile when using show]'                         \
    '--redirect[send messages to stdout]'                                               \
    '--show_hidden[avail and spider will report hidden modules]'                        \
    '--style[]'                                                                         \
    '--timer[report run times]'                                                         \
    '--terse_show_extensions[Include module extensions in terse avail output]'          \
    '*::module command:_module_command'
}

_module_command()
{
  local curcontext="$curcontext"
  local -a _module_cmds

  _module_cmds=(
    "add:load a module into the shell environment"
    "avail:list all available modules"
    "describe:module collection contents"
    "disable:disable a module collection"
    "help:print the usage of each sub-command"
    "keyword:search for a given keyword in modules"
    "list:list loaded modules"
    "load:load a module into the shell environment"
    "mcc:module collection contents"
    "overview:list all available modules by name only no versions"
    "purge:unload all loaded modules"
    "refresh:refresh all non-persistent components of loaded modules"
    "restore:load a module collection"
    "rm:remove a module from the shell environment"
    "save:save a collection"
    "show:display information about a module"
    "spider:List all possible modules"
    "swap:swap loaded a loaded module with another module"
    "try-load:load a module into the shell environment, no warnings if not found"
    "unload:remove a module from the shell environment"
    "unuse:remove a directory from MODULEPATH"
    "update:reload all loaded modules"
    "use:add a directory to MODULEPATH"
    "whatis:display module information"
  )

  if (( CURRENT == 1 )); then
    _describe -t commands 'module command' _module_cmds || compadd "$@"
  else
    local curcontext="$curcontext"

    cmd="${${_module_cmds[(r)$words[1]:*]%%:*}}"
    # Deal with any aliases
    case $cmd in
      add|try-load) cmd="load";;
      rm) cmd="unload";;
      display) cmd="show";;
    esac

    if (( $#cmd )); then
      local update_policy
      curcontext="${curcontext%:*:*}:module-${cmd}:"
      zstyle -s ":completion:${curcontext}:" cache-policy update_policy
      _call_function ret _module_$cmd || _message 'no more arguments'
    else
      _message "unknown module command: $words[1]"
    fi

    return ret
  fi
}

# Fills the available modules cache
_module_loaded_modules()
{
  _loaded_modules=(${$(module -q -t --no_redirect list 2>&1 > /dev/null | sed ' /^ *$/d; /:$/d; s#/*$##g; s/ *<[A-Za-z][A-Za-z]*>//')})
}

# Fills the available modules cache
_module_available_modules()
{
  if [[ -n $MODULEPATH ]];  then
    _available_modules=(${$(module -q -t --no_redirect avail 2>&1 > /dev/null | sed ' /:$/d; s/(@.*)//g; s#/*$##g; s/ *<[A-Za-z][A-Za-z]*>//')})
  fi
}

# Fills the available modules cache
_module_spider_list()
{
  _spider_list=(${$(module -q -t --no_redirect spider 2>&1 > /dev/null)})
}

_module_restore()
{
  _savelist_list=(${$(module -q -t --no_redirect savelist 2>&1 > /dev/null)})
  compadd "$@" -a -- _savelist_list
}

_module_disable()
{
  _savelist_list=(${$(module -q -t --no_redirect savelist 2>&1 > /dev/null)})
  compadd "$@" -a -- _savelist_list
}

_ml_mcc()
{
  _savelist_list=(${$(module -q -t --no_redirect savelist 2>&1 > /dev/null)})
  compadd "$@" -a -- _savelist_list
}

_ml_describe()
{
  _savelist_list=(${$(module -q -t --no_redirect savelist 2>&1 > /dev/null)})
  compadd "$@" -a -- _savelist_list
}

# Completion function for help
_module_help()
{
  _module_available_modules
  compadd "$@" -a -- _available_modules
}

# Completion function for load|add
_module_load()
{
  _module_available_modules
  compadd "$@" -a -- _available_modules
}

# Completion function for spider
_module_spider()
{
  _module_spider_list
  compadd "$@" -a -- _spider_list
}

# Completion function for unload|rm
_module_unload()
{
  _module_loaded_modules
  compadd "$@" -a -- _loaded_modules
}

# Completion function for switch|swap
_module_swap()
{
  # Actually first argument could be a loaded module

  if (( CURRENT == 2 )); then
    _module_loaded_modules
    compadd "$@" -a -- _loaded_modules
  else
    _module_available_modules
    compadd "$@" -a -- _available_modules
  fi
}

# Completion function for display|show
_module_show()
{
  _module_available_modules
  compadd "$@" -a -- _available_modules
}

# Completion function for avail
### No completion (yet?)

# Completion function for use
_module_use()
{
  _arguments -s \
    '(-a --append)'{--append,-a}'[append the directories instead of prepending]' \
    '*:directory:_files -/'
}

# Completion function for unuse
_module_unuse()
{
  compadd "$@" -- ${=MODULEPATH//:/ }
}

# Completion function for whatis
_module_whatis()
{
  _module_available_modules
  compadd "$@" -a -- _available_modules
}

_module "$@"

# Local Variables:
# mode: shell-script
# indent-tabs-mode: nil
# End:
