#compdef docker

# Docker autocompletion for oh-my-zsh
# Requires: Docker installed
# Author: Azaan (@aeonazaan)
# Updates: Bob Maerten (@bobmaerten) for Docker v0.9+
#          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")]"}'))
    if [[  'X$cont_cmd' != 'X' ]]
        _describe 'containers' cont_cmd
}

# Output a selectable list of all containers, even not running
__docker_all_containers() {
    declare -a cont_cmd
    cont_cmd=($(docker ps -a | awk 'NR>1{print $NF":[CON("$1")"$2"("$3")]"}'))
    if [[  'X$cont_cmd' != 'X' ]]
        _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}'| sed 's/:/\\:/g'))
    _describe 'images' img_cmd
}

# ----- 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 \
        '--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 (e.g. "John Hannibal Smith <hannibal@a-team.com>")]' \
        '(-c,--change=)'{-c,--change=}'[Apply Dockerfile instruction to the created image]' \
        '(-m,--message=)'{-m,--message=}'[Commit message]' \
        '(-p,--pause=)'{-p,--pause=}'[Pause container during commit]' \
}

__cp() {
    __docker_containers
}

__create() {
    _arguments \
        '(-P,--publish-all=)'{-P,--publish-all=}'[Publish all exposed ports to the host interfaces]' \
        '(-a,--attach=)'{-a,--attach=}'[Attach to STDIN, STDOUT or STDERR]' \
        '--add-host=[Add a custom host-to-IP mapping]' \
        '--cap-add=[Add Linux capabilities]' \
        '--cap-drop=[Drop Linux capabilities]' \
        '--cpuset-cpus=[CPUs in which to allow execution (0-3, 0,1)]' \
        '(-c,--cpu-shares=)'{-c,--cpu-shares=}'[CPU shares (relative weight)]' \
        '--cidfile=[Write the container ID to the file]' \
        '--device=[Add a host device to the container]' \
        '--dns=[Set custom dns servers]' \
        '--dns-search=[Set custom DNS search domains]' \
        '(-e,--env=)'{-e,--env=}'[Set environment variables]' \
        '--env-file=[Read in a file of environment variables]' \
        '--entrypoint=[Overwrite the default entrypoint of the image]' \
        '--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]' \
        '--ipc=[IPC namespace to use]' \
        '(-l,--label=)'{-l,--label=}'[Set meta data on a container]' \
        '--link=[Add link to another container (name:alias)]' \
        '--log-driver=[Logging driver for the container]' \
        '--lxc-conf=[Add custom LXC options]' \
        '--mac-address=[Container MAC address (e.g. 92:d0:c6:0a:29:33)]' \
        '(-m,--memory=)'{-m,--memory=}'[Memory limit (format: <number><optional unit>, where unit = b, k, m or g)]' \
        '--net=[Set the Network mode for the container]' \
        '--name=[Assign a name to the container]' \
        '--pid=[PID namespace to use]' \
        '(-p,--publish=)'{-p,--publish=}'[Publish a container''s port to the host (format: ip:hostPort:containerPort/protocol)]' \
        '--privileged=[Give extended privileges to this container]' \
        '--restart=[Restart policy to apply when a container exits]' \
        '--security-opt=[Security Options]' \
        '--sig-proxy=[Proxify all received signal to the process (even in non-tty mode)]' \
        '(-t,--tty=)'{-t,--tty=}'[Allocate a pseudo-tty]' \
        '(-u,--user=)'{-u,--user=}'[Username or UID]' \
        '--ulimit=[Ulimit options]' \
        '(-v,--volume=)'{-v,--volume=}'[Bind mount a volume (e.g. -v /host:/container or -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.]'
}

__export() {
   __docker_containers
}

__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]' \
        '(-v,--viz=)'{-v,--viz=}'[Output graph in graphviz format]'
    __docker_images
}

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

__info() {
    # no arguments
}

__inspect() {
    __docker_images
    __docker_all_containers
}

__kill() {
    _arguments \
        '(-s,--signal=)'{-s,--signal=}'[KILL Signal]'
    __docker_containers
}

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

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

__logs() {
    _arguments \
        '(-f,--follow)'{-f,-follow}'[Follow log output]'
    __docker_containers
}

__port() {
    __docker_containers
}

__top() {
    __docker_containers
}

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

__pull() {
    _arguments \
        '(-t,--tag=)'{-t,--tag=}'[Download tagged image in 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=10]'
    __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_all_containers
}

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

__run() {
    _arguments \
        '(-P,--publish-all=)'{-P,--publish-all=}'[Publish all exposed ports to the host interfaces]' \
        '(-a,--attach=)'{-a,--attach=}'[Attach to STDIN, STDOUT or STDERR]' \
        '--add-host=[Add a custom host-to-IP mapping]' \
        '--cap-add=[Add Linux capabilities]' \
        '--cap-drop=[Drop Linux capabilities]' \
        '--cpuset-cpus=[CPUs in which to allow execution (0-3, 0,1)]' \
        '(-c,--cpu-shares=)'{-c,--cpu-shares=}'[CPU shares (relative weight)]' \
        '--cidfile=[Write the container ID to the file]' \
        '(-d,--detach=)'{-d,--detach=}'[Run container in the background, print new container id]' \
        '--device=[Add a host device to the container]' \
        '--dns=[Set custom dns servers]' \
        '--dns-search=[Set custom DNS search domains]' \
        '(-e,--env=)'{-e,--env=}'[Set environment variables]' \
        '--env-file=[Read in a file of environment variables]' \
        '--entrypoint=[Overwrite the default entrypoint of the image]' \
        '--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]' \
        '--ipc=[IPC namespace to use]' \
        '(-l,--label=)'{-l,--label=}'[Set meta data on a container]' \
        '--link=[Add link to another container (name:alias)]' \
        '--log-driver=[Logging driver for the container]' \
        '--lxc-conf=[Add custom LXC options]' \
        '--mac-address=[Container MAC address (e.g. 92:d0:c6:0a:29:33)]' \
        '(-m,--memory=)'{-m,--memory=}'[Memory limit (format: <number><optional unit>, where unit = b, k, m or g)]' \
        '--net=[Set the Network mode for the container]' \
        '--name=[Assign a name to the container]' \
        '--pid=[PID namespace to use]' \
        '(-p,--publish=)'{-p,--publish=}'[Publish a container''s port to the host (format: ip:hostPort:containerPort/protocol)]' \
        '--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 Options]' \
        '--sig-proxy=[Proxify all received signal to the process (even in non-tty mode)]' \
        '(-t,--tty=)'{-t,--tty=}'[Allocate a pseudo-tty]' \
        '(-u,--user=)'{-u,--user=}'[Username or UID]' \
        '--ulimit=[Ulimit options]' \
        '(-v,--volume=)'{-v,--volume=}'[Bind mount a volume (e.g. -v /host:/container or -v /container)]' \
        '--volumes-from=[Mount volumes from the specified container(s)]' \
        '(-w,--workdir=)'{-w,--workdir=}'[Working directory inside the container]'
    __docker_images
}

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

__save() {
    __docker_images
}

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

__stats() {
    __docker_containers
}

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

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

__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 a container from a Dockerfile"
    "commit":"Create a new image from a container's changes"
    "cp":"Copy files/folders from the containers filesystem to the host path"
    "create":"Create new container without running it"
    "diff":"Inspect changes on a container's filesystem"
    "events":"Get real time events from the server"
    "export":"Stream the contents of a container as a tar archive"
    "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 Login to the docker registry server"
    "logs":"Fetch the logs of a container"
    "port":"Lookup the public-facing port which is NAT-ed to PRIVATE_PORT"
    "ps":"List containers"
    "pull":"Pull an image or a repository from the docker registry server"
    "push":"Push an image or a repository to the 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 in the docker index"
    "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"
    "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 ;;
    cp)
        __cp ;;
    create)
        __create ;;
    diff)
        __diff ;;
    events)
        __events ;;
    export)
        __export ;;
    history)
        __history ;;
    images)
        __images ;;
    import)
        __import ;;
    info)
        __info ;;
    inspect)
        __inspect ;;
    kill)
        __kill ;;
    load)
        __load ;;
    login)
        __login ;;
    logs)
        __logs ;;
    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 ;;
    version)
        __version ;;
    wait)
        __wait ;;
    exec)
        __exec ;;
esac