#compdef docker

# Docker autocompletion for oh-my-zsh
# Requires: Docker installed
# Author: Azaan (@aeonazaan)
# Updates: Bob Maerten (@bobmaerten) for Docker v0.9 & v1.3
#          Paul van den Berg (@bergvandenp) for Docker v1.3+


# ----- Helper functions
# Output a selectable list of all running docker containers
__docker_containers() {
    declare -a cont_cmd
    cont_cmd=($(docker ps | awk 'NR>1{print $NF":[CON("$1")"$2"("$3")]"}'))
    _describe 'containers' cont_cmd
}

# output a selectable list of all docker images
__docker_images() {
    declare -a img_cmd
    img_cmd=($(docker images | awk 'NR>1{print $1}'))
    _describe 'images' img_cmd
}

# output a selectable list of all docker commands
__docker_commands() {
    declare -a cmd_list
    NB=$(docker help 2>&1 | grep -n Commands | cut -f1 -d':')
    NB=$((NB+1))
    cmd_list=($(docker help 2>&1 | tail -n +${NB} | head -n -2 | awk 'NR>1{print $1}'))
    _describe 'commands' cmd_list
}

# ----- Commands
# Seperate function for each command, makes extension easier later
# ---------------------------
__attach() {
    _arguments \
        '--no-stdin[Do not attach stdin]' \
        '--sig-proxy[Proxify all received signal to the process (even in non-tty mode)]'
    __docker_containers
}

__build() {
    _arguments \
        '--force-rm[Always remove intermediate containers, even after unsuccessful builds]' \
        '--no-cache[Do not use cache when building the image]' \
        '(-q,--quiet=)'{-q,--quiet=}'[Suppress the verbose output generated by the containers]' \
        '--rm[Remove intermediate containers after a successful build]' \
        '(-t,--tag=)'{-t,--tag=}'[Repository name (and optionally a tag) to be applied to the resulting image in case of success]' \
        '*:files:_files'
}

__commit() {
    _arguments \
        '(-a,--author=)'{-a,--author=}'[Author (eg. "John Hannibal Smith <hannibal@a-team.com>"]' \
        '(-m,--message=)'{-m,--message=}'[Commit message]' \
        '(-p,--pause=)'{-p,--pause=}'[Pause container during commit]'
    __docker_containers
}

__cp() {
    __docker_containers
}
__create() {
    _arguments \
        '(-a,--attach=)'{-a,--attach=}'[Attach to STDIN, STDOUT or STDERR.]' \
        '--add-host=[Add a custom host-to-IP mapping (host:ip)]' \
        '(-c,--cpu-shares=)'{-c,--cpu-shares=}'[CPU shares (relative weight)]' \
        '--cap-add=[Add Linux capabilities]' \
        '--cap-drop=[Drop Linux capabilities]' \
        '--cidfile=[Write the container ID to the file]' \
        '--cpuset=[CPUs in which to allow execution (0-3, 0,1)]' \
        '--device=[Add a host device to the container (e.g. --device=/dev/sdc:/dev/xvdc)]' \
        '--dns=[Set custom DNS servers]' \
        '--dns-search=[Set custom DNS search domains]' \
        '(-e,--env=)'{-e,--env=}'[Set environment variables]' \
        '--entrypoint=[Overwrite the default ENTRYPOINT of the image]' \
        '--env-file=[Read in a line delimited file of environment variables]' \
        '--expose=[Expose a port from the container without publishing it to your host]' \
        '(-h,--hostname=)'{-h,--hostname=}'[Container host name]' \
        '(-i,--interactive=)'{-i,--interactive=}'[Keep STDIN open even if not attached]' \
        '--link=[Add link to another container in the form of name:alias]' \
        '--lxc-conf=[(lxc exec-driver only) Add custom lxc options --lxc-conf="lxc.cgroup.cpuset.cpus = 0,1"]' \
        '(-m,--memory=)'{-m,--memory=}'[Memory limit (format: <number><optional unit>, where unit = b, k, m or g)]' \
        '--name=[Assign a name to the container]' \
        '--net=[Set the Network mode for the container]' \
        '(-P,--publish-all=)'{-P,--publish-all=}'[Publish all exposed ports to the host interfaces]' \
        '(-p,--publish=)'{-p,--publish=}'[Publish a container''s port to the host]' \
        '--privileged=[Give extended privileges to this container]' \
        '--restart=[Restart policy to apply when a container exits]' \
        '--security-opt=[Security Options]' \
        '(-t,--tty=)'{-t,--tty=}'[Allocate a pseudo-TTY]' \
        '(-u,--user=)'{-u,--user=}'[Username or UID]' \
        '(-v,--volume=)'{-v,--volume=}'[Bind mount a volume (e.g., from the host: -v /host:/container, from Docker: -v /container)]' \
        '--volumes-from=[Mount volumes from the specified container(s)]' \
        '(-w,--workdir=)'{-w,--workdir=}'[Working directory inside the container]'
    __docker_images
}
__diff() {
    __docker_containers
}

__events() {
    _arguments \
        '--since=[Show previously created events and then stream.]' \
        '--until=[Stream events until this timestamp]'
}

__exec() {
    _arguments \
        '(-d,--detach=)'{-d,--detach=}'[Detached mode: run command in the background]' \
        '(-i,--interactive=)'{-i,--interactive=}'[Keep STDIN open even if not attached]' \
        '(-t,--tty=)'{-t,--tty=}'[Allocate a pseudo-TTY]'
    __docker_containers
}

__export() {
   __docker_containers
}

__help() {
    __docker_commands
}


__history() {
    _arguments \
        '--no-trunc[Don''t truncate output]' \
        '(-q,--quiet)'{-q,--quiet}'[Only show numeric IDs]'
    __docker_images
}

__images() {
    _arguments \
        '(-a,--all)'{-a,--all}'[Show all images (by default filter out the intermediate images used to build)]' \
        '--no-trunc[Don''t truncate output]' \
        '(-q,--quiet=)'{-q,--quiet=}'[Only show numeric IDs]' \
        '(-t,--tree=)'{-t,--tree=}'[Output graph in tree format]'
    __docker_images
}

__import() {
    _arguments '*:files:_files'
}

__info() {
    # no arguments
}

__inspect() {
    __docker_images
    _arguments \
        '(-f,--format=)'{-f,--format=}'[Format the output using the given go template.]'
    __docker_containers
}

__kill() {
    _arguments \
        '(-s,--signal=)'{-s,--signal=}'[Signal to send to the container]'
    __docker_containers
}

__load() {
    _arguments \
        '(-i,--input=)'{-i,--input=}'[Read from a tar archive file, instead of STDIN]' \
        '*:files:_files'
}

__login() {
    _arguments \
        '(-e,--email=)'{-e,-email=}'[Email]' \
        '(-p,--password=)'{-p,-password=}'[Password]' \
        '(-u,--username=)'{-u,-username=}'[Username]'
}

__logout() {
    # no autocomplete
}

__logs() {
    _arguments \
        '(-f,--follow=)'{-f,--follow=}'[Follow log output]' \
        '(-t,--timestamps=)'{-t,--timestamps=}'[Show timestamps]' \
        '--tail[Output the specified number of lines at the end of logs (defaults to all logs)]'
    __docker_containers
}

__port() {
    __docker_containers
}

__top() {
    __docker_containers
}

__pause() {
    __docker_containers
}

__ps() {
    _arguments \
        '(-a,--all=)'{-a,--all=}'[Show all containers. Only running containers are shown by default.]' \
        '--before=[Show only container created before Id or Name, include non-running ones.]' \
        '(-f,--filter=)'{-f,--filter=}'[Provide filter values.]' \
        '(-l,--latest=)'{-l,--latest=}'[Show only the latest created container, include non-running ones.]' \
        '-n=[Show n last created containers, include non-running ones.]' \
        '--no-trunc[Don''t truncate output]' \
        '(-q,--quiet=)'{-q,--quiet=}'[Only display numeric IDs]' \
        '(-s,--size=)'{-s,--size=}'[Display sizes]' \
        '--since=[Show only containers created since Id or Name, include non-running ones.]'
}

__pull() {
    _arguments \
        '(-a,--all-tags=)'{-a,--all-tags=}'[Download all tagged images in the repository]'
}

__push() {
    # no arguments
}

__restart() {
    _arguments \
        '(-t,--time=)'{-t,--time=}'[Number of seconds to try to stop for before killing the container. Once killed it will then be restarted. Default is 10 seconds.]'
    __docker_containers
}

__rm() {
    _arguments \
        '(-f,--force=)'{-f,--force=}'[Force removal of running container]' \
        '(-l,--link=)'{-l,--link=}'[Remove the specified link and not the underlying container]' \
        '(-v,--volumes=)'{-v,--volumes=}'[Remove the volumes associated to the container]'
    __docker_containers
}

__rmi() {
    _arguments \
        '(-f,--force=)'{-f,--force=}'[Force removal of the image]' \
        '--no-prune[Do not delete untagged parent]'
    __docker_images
}

__run() {
    _arguments \
        '(-a,--attach=)'{-a,--attach=}'[Attach to STDIN, STDOUT or STDERR.]' \
        '--add-host=[Add a custom host-to-IP mapping (host:ip)]' \
        '(-c,--cpu-shares=)'{-c,--cpu-shares=}'[CPU shares (relative weight)]' \
        '--cap-add=[Add Linux capabilities]' \
        '--cap-drop=[Drop Linux capabilities]' \
        '--cidfile=[Write the container ID to the file]' \
        '--cpuset=[CPUs in which to allow execution (0-3, 0,1)]' \
        '(-d,--detach=)'{-d,--detach=}'[Detached mode: run the container in the background and print the new container ID]' \
        '--device=[Add a host device to the container (e.g. --device=/dev/sdc:/dev/xvdc)]' \
        '--dns=[Set custom DNS servers]' \
        '--dns-search=[Set custom DNS search domains]' \
        '(-e,--env=)'{-e,--env=}'[Set environment variables]' \
        '--entrypoint=[Overwrite the default ENTRYPOINT of the image]' \
        '--env-file=[Read in a line delimited file of environment variables]' \
        '--expose=[Expose a port from the container without publishing it to your host]' \
        '(-h,--hostname=)'{-h,--hostname=}'[Container host name]' \
        '(-i,--interactive=)'{-i,--interactive=}'[Keep STDIN open even if not attached]' \
        '--link=[Add link to another container in the form of name:alias]' \
        '--lxc-conf=[(lxc exec-driver only) Add custom lxc options --lxc-conf="lxc.cgroup.cpuset.cpus = 0,1"]' \
        '(-m,--memory=)'{-m,--memory=}'[Memory limit (format: <number><optional unit>, where unit = b, k, m or g)]' \
        '--name=[Assign a name to the container]' \
        '--net=[Set the Network mode for the container]' \
        '(-P,--publish-all=)'{-P,--publish-all=}'[Publish all exposed ports to the host interfaces]' \
        '(-p,--publish=)'{-p,--publish=}'[Publish a container''s port to the host]' \
        '--privileged=[Give extended privileges to this container]' \
        '--restart=[Restart policy to apply when a container exits]' \
        '--rm=[Automatically remove the container when it exits (incompatible with -d)]' \
        '--security-opt=)'{--security-opt=}'[Security Options]' \
        '--sig-proxy=[Proxy received signals to the process (even in non-TTY mode). SIGCHLD, SIGSTOP, and SIGKILL are not proxied.]' \
        '(-t,--tty=)'{-t,--tty=}'[Allocate a pseudo-TTY]' \
        '(-u,--user=)'{-u,--user=}'[Username or UID]' \
        '(-v,--volume=)'{-v,--volume=}'[Bind mount a volume (e.g., from the host: -v /host:/container, from Docker: -v /container)]' \
        '--volumes-from=[Mount volumes from the specified container(s)]' \
        '(-w,--workdir=)'{-w,--workdir=}'[Working directory inside the container]'
    __docker_images
}

__search() {
    _arguments \
        '--automated[Only show automated builds]' \
        '--no-trunc[Don''t truncate output]' \
        '(-s,--stars=)'{-s,--stars=}'[Only displays with at least x stars]'
}

__save() {
    _arguments \
        '(-o,--output=)'{-o,--output=}'[Write to a file, instead of STDOUT]'
    __docker_images
}

__start() {
    _arguments \
        '(-a,--attach=)'{-a,--attach=}'[Attach container''s STDOUT and STDERR and forward all signals to the process]' \
        '(-i,--interactive=)'{-i,--interactive=}'[Attach container''s STDIN]'
    __docker_containers
}

__stats() {
    __docker_containers
}

__stop() {
    _arguments \
        '(-t,--time=)'{-t,--time=}'[Number of seconds to wait for the container to stop before killing it. Default is 10 seconds.]'
    __docker_containers
}

__tag() {
    _arguments \
        '(-f,--force=)'{-f,--force=}'[Force]'
    __docker_images
}

__unpause() {
    __docker_containers
}

__version() {
    # no arguments
}

__wait() {
    __docker_containers
}

__exec() {
    _arguments \
      '(-d,--detach=)'{-d,--detach=}'[Detached mode: run command in the background]' \
      '(-i,--interactive=)'{-i,--interactive=}'[Keep STDIN open even if not attached]' \
      '(-t,--tty=)'{-t,--tty=}'[Allocate a pseudo-TTY]'
    __docker_containers
}

# end commands ---------
# ----------------------

local -a _1st_arguments
_1st_arguments=(
    "attach":"Attach to a running container"
    "build":"Build an image from a Dockerfile"
    "commit":"Create a new image from a container's changes"
    "cp":"Copy files/folders from a container's filesystem to the host path"
    "create":"Create a new container"
    "diff":"Inspect changes on a container's filesystem"
    "events":"Get real time events from the server"
    "exec":"Run a command in an existing container"
    "export":"Stream the contents of a container as a tar archive"
    "help":"Display help for a given command"
    "history":"Show the history of an image"
    "images":"List images"
    "import":"Create a new filesystem image from the contents of a tarball"
    "info":"Display system-wide information"
    "inspect":"Return low-level information on a container"
    "kill":"Kill a running container"
    "load":"Load an image from a tar archive"
    "login":"Register or log in to a Docker registry server"
    "logout":"Log out from a Docker registry server"
    "logs":"Fetch the logs of a container"
    "port":"Lookup the public-facing port that is NAT-ed to PRIVATE_PORT"
    "pause":"Pause all processes within a container"
    "ps":"List containers"
    "pull":"Pull an image or a repository from a Docker registry server"
    "push":"Push an image or a repository to a Docker registry server"
    "restart":"Restart a running container"
    "rm":"Remove one or more containers"
    "rmi":"Remove one or more images"
    "run":"Run a command in a new container"
    "save":"Save an image to a tar archive"
    "search":"Search for an image on the Docker Hub"
    "start":"Start a stopped container"
    "stats":"Display a live stream of one or more containers' resource usage statistics"
    "stop":"Stop a running container"
    "tag":"Tag an image into a repository"
    "top":"Lookup the running processes of a container"
    "unpause":"Unpause a paused container"
    "version":"Show the Docker version information"
    "wait":"Block until a container stops, then print its exit code"
    "exec":"Run a task inside a running container"
)

_arguments '*:: :->command'

if (( CURRENT == 1 )); then
    _describe -t commands "docker command" _1st_arguments
    return
fi

local -a _command_args
case "$words[1]" in
    attach)
       __attach ;;
    build)
        __build ;;
    commit)
        __commit ;;
    create)
        __create ;;
    cp)
        __cp ;;
    diff)
        __diff ;;
    events)
        __events ;;
    exec)
        __exec ;;
    export)
        __export ;;
    help)
        __help ;;
    history)
        __history ;;
    images)
        __images ;;
    import)
        __import ;;
    info)
        __info ;;
    inspect)
        __inspect ;;
    kill)
        __kill ;;
    load)
        __load ;;
    login)
        __login ;;
    logout)
        __logout ;;
    logs)
        __logs ;;
    pause)
        __pause ;;
    port)
        __port ;;
    ps)
        __ps ;;
    pull)
        __pull ;;
    push)
        __push ;;
    restart)
        __restart ;;
    rm)
        __rm ;;
    rmi)
        __rmi ;;
    run)
        __run ;;
    save)
        __save ;;
    search)
        __search ;;
    stats)
        __stats ;;
    start)
        __start ;;
    stop)
        __stop ;;
    tag)
        __tag ;;
    top)
        __top ;;
    unpause)
        __unpause ;;
    version)
        __version ;;
    wait)
        __wait ;;
    exec)
        __exec ;;
esac