#!/bin/bash -u

# This script is used to run a docker container with graphics support.
# All arguments to this script except "-c <container_name>" will be appended to a docker run command.
# If a container name is specified, and this container already exists, the container is re-entered,
# which easily allows entering the same persistent container from multiple terminals.
# See documentation for detailed examples: https://moveit.ros.org/install/docker/

# Example commands:
# ./gui-docker --rm -it moveit/moveit:master-source /bin/bash     # Run a (randomly named) container that is removed on exit
# ./gui-docker -v ~/ros_ws:/root/ros_ws --rm -it moveit/moveit:master-source /bin/bash   # Same, but also link host volume ~/ros_ws to /root/ros_ws in the container
# ./gui-docker -c container_name                                  # Start (or continue) an interactive bash in a moveit/moveit:master-source container
# ./gui-docker                                                    # Same, but use the default container name "default_moveit_container"

function check_nvidia2() {
    # If we don't have an NVIDIA graphics card, bail out
    lspci | grep -qi "vga .*nvidia" || return 1
    # If we don't have the nvidia runtime, bail out
    if ! docker -D info | grep -qi "runtimes.* nvidia" ; then
       echo "nvidia-docker v2 not installed (see https://github.com/NVIDIA/nvidia-docker/wiki)"
       return 2
    fi
    echo "found nvidia-docker v2"
    DOCKER_PARAMS="\
      --runtime=nvidia \
      --env=NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES:-all} \
      --env=NVIDIA_DRIVER_CAPABILITIES=${NVIDIA_DRIVER_CAPABILITIES:+$NVIDIA_DRIVER_CAPABILITIES,}all"
    return 0
}

function check_nvidia1() {
    # If we don't have an NVIDIA graphics card, bail out
    lspci | grep -qi "vga .*nvidia" || return 1
    # Check whether nvidia-docker is available
    if ! which nvidia-docker > /dev/null ; then
       echo "nvidia-docker v1 not installed either"
       return 2
    fi
    # Check that nvidia-modprobe is installed
    if ! which nvidia-modprobe > /dev/null ; then
       echo "nvidia-docker-plugin requires nvidia-modprobe. Please install it!"
       return 3
    fi
    # Retrieve device parameters from nvidia-docker-plugin
    if ! DOCKER_PARAMS=$(curl -s http://localhost:3476/docker/cli) ; then
       echo "nvidia-docker-plugin not responding on http://localhost:3476/docker/cli"
       echo "See https://github.com/NVIDIA/nvidia-docker/wiki/Installation-(version-1.0)"
       return 3
    fi
    echo "found nvidia-docker v1"
    DOCKER_EXECUTABLE=nvidia-docker
}

function check_dri() {
    # If there is no /dev/dri, bail out
    test -d /dev/dri || return 1
    DOCKER_PARAMS="--device=/dev/dri --group-add video"
}

function transfer_x11_permissions() {
    # store X11 access rights in temp file to be passed into docker container
    XAUTH=/tmp/.docker.xauth
    touch $XAUTH
    xauth nlist $DISPLAY | sed -e 's/^..../ffff/' | xauth -f $XAUTH nmerge -
}

function count_positional_args() {
    while true ; do
       case "${1:-}" in
          # Skip common options with a subsequent positional argument
          # This list is not exhaustive! Augment as you see fit.
          -v|--volume) shift ;;
          -w) shift ;;
          -e) shift ;;
          # Skip all other options
          -*) ;;
          *) break ;;
       esac
       shift
    done
    # Return remaining number of arguments
    echo $#
}

if [ $# -eq 0 ] ; then
   # If no options are specified at all, use the name "default_moveit_container"
   CONTAINER_NAME=default_moveit_container
else
  # Check for option -c or --container in first position
  case "$1" in
    -c|--container)
      shift
      # If next argument is not an option, use it as the container name
      if [[ "${1:-}" != -* ]] ; then
         CONTAINER_NAME="${1:-}"
         shift
      fi
      # Set default container name if still undefined
      CONTAINER_NAME="${CONTAINER_NAME:-default_moveit_container}"
      ;;
  esac
fi

transfer_x11_permissions

# Probe for nvidia-docker (version 2 or 1)
check_nvidia2 || check_nvidia1 || check_dri || echo "No supported graphics card found"

DOCKER_EXECUTABLE=${DOCKER_EXECUTABLE:-docker}

# If CONTAINER_NAME was specified and this container already exists, continue it
if [ -n "${CONTAINER_NAME:-}" ] ; then
   if [ -z "$($DOCKER_EXECUTABLE ps -aq --filter name=^$CONTAINER_NAME\$)" ] ; then
      # container not yet existing: add an option to name the container when running docker below
      NAME_OPTION="--name=$CONTAINER_NAME"
      if [ "$(count_positional_args $@)" == "0" ] ; then
         # If no further (positional) arguments were provided, start a bash in the default image (for dummy users)
         DUMMY_DEFAULTS="-it moveit/moveit:master-source bash"
      fi
   else
      if [ -z "$($DOCKER_EXECUTABLE ps -q --filter name=^$CONTAINER_NAME\$)" ] ; then
         echo -n "Start existing, but stopped container: "
         docker start $CONTAINER_NAME
      fi
      echo "Entering container: $CONTAINER_NAME"
      if [ $# -eq 0 ] ; then
         docker exec -it $CONTAINER_NAME bash
      else
         docker exec $CONTAINER_NAME $@
      fi
      rm $XAUTH
      exit 0
   fi
fi

${DOCKER_EXECUTABLE:-docker} run \
    --env="DISPLAY=$DISPLAY" \
    --env="QT_X11_NO_MITSHM=1" \
    --env="XAUTHORITY=$XAUTH" \
    --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" \
    --volume="$XAUTH:$XAUTH" \
    ${NAME_OPTION:-} \
    ${DOCKER_PARAMS:-} \
    $@ ${DUMMY_DEFAULTS:-}

# cleanup
rm $XAUTH
