#!/bin/bash


# Generates a Dockerfile for a development environment.
#  - Maps this script's user to the docker image, with sudo priveleges
#  - Installs a basic command line environment
#  - Installs a gui environment for nvidia machines
#
# Warning: currently no checking if it is in fact an nvidia machine

##############################################################################
# Colours
##############################################################################

BOLD="\e[1m"

CYAN="\e[36m"
GREEN="\e[32m"
RED="\e[31m"
YELLOW="\e[33m"

RESET="\e[0m"

padded_message ()
{
  line="........................................"
  printf "%s %s${2}\n" ${1} "${line:${#1}}"
}

pretty_header ()
{
  echo -e "${BOLD}${1}${RESET}"
}

pretty_print ()
{
  echo -e "${GREEN}${1}${RESET}"
}

pretty_warning ()
{
  echo -e "${YELLOW}${1}${RESET}"
}

pretty_error ()
{
  echo -e "${RED}${1}${RESET}"
}

##############################################################################
# Methods
##############################################################################

install_package ()
{
  PACKAGE_NAME=$1
  dpkg -s ${PACKAGE_NAME} > /dev/null
  if [ $? -ne 0 ]; then
    sudo apt-get -q -y install ${PACKAGE_NAME} > /dev/null
  else
    pretty_print "  $(padded_message ${PACKAGE_NAME} "found")"
    return 0
  fi
  if [ $? -ne 0 ]; then
    pretty_error "  $(padded_message ${PACKAGE_NAME} "failed")"
    return 1
  fi
  pretty_warning "  $(padded_message ${PACKAGE_NAME} "installed")"
  return 0
}

install_nvidia_docker2 ()
{
  PACKAGE_NAME=nvidia-docker2
  dpkg -s ${PACKAGE_NAME} > /dev/null
  if [ $? -ne 0 ]; then
    HOST_DISTRIBUTION=$(. /etc/os-release; echo $ID$VERSION_ID)
    curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
    curl -s -L https://nvidia.github.io/nvidia-docker/${HOST_DISTRIBUTION}/nvidia-docker.list | \
                sudo tee /etc/apt/sources.list.d/nvidia-docker.list
    sudo apt-get update
    sudo apt-get -q -y install ${PACKAGE_NAME} > /dev/null
    if [ $? -ne 0 ]; then
        pretty_error "  $(padded_message ${PACKAGE_NAME} "failed")"
        return 1
    fi
    sudo pkill -SIGHUP dockerd
    sudo service docker restart
    pretty_warning "  $(padded_message ${PACKAGE_NAME} "installed")"
  else
    pretty_print "  $(padded_message ${PACKAGE_NAME} "found")"
    return 0
  fi
  return 0
}

##############################################################################
# Command Line Parsing
##############################################################################

show_help ()
{
  pretty_print "    Usage:"
  pretty_print "        $0 <options>"
  pretty_print ""
  pretty_print "    Required Options:"
  pretty_print "        --distro=<distro> : one of bionic, focal"
  pretty_print ""
  pretty_print "    Options:"
  pretty_print "        --help            : show this help message"
  pretty_print ""
  pretty_print "    Note: this utilises the current shell user's id/name"
  pretty_print "    to create an analogous user within the docker image."
  pretty_print ""
  exit 0
}

for i in "$@"
do
case $i in
    --help)
    SHOW_HELP=1
    shift
    ;;
    -d=*|--distro=*)
    DISTRO="${i#*=}"
    shift
    ;;
    *)
       # nothing to do
    ;;
esac
done

if [ ! -z "$SHOW_HELP" ]; then
  show_help
  exit 0
fi

if [ -z "${DISTRO}" ]; then
  pretty_error "The argument '--distro' is a required option, refer to --help for more details."
  exit 1
fi

if [ ${DISTRO} = "bionic" ]; then
  RELEASE="18.04"
elif [ ${DISTRO} = "focal" ]; then
  RELEASE="20.04"
else
  pretty_error "Unsupported Release ['bionic', 'focal']"
  exit 1
fi

##############################################################################
# System Dependencies
##############################################################################

pretty_header "System Dependencies"
install_package docker.io || exit 1
install_nvidia_docker2 || exit 1

##############################################################################
# Variables
##############################################################################

DOCKERFILE_DIR=`mktemp -d -p "/tmp"`
DOCKERFILE=${DOCKERFILE_DIR}/Dockerfile
# For vanilla ubuntu images, ubuntu:${RELEASE}
UBUNTU_IMAGE=nvidia/opengl:1.0-glvnd-runtime-ubuntu${RELEASE}
DOCKER_IMAGE=groot:${DISTRO}
USER_ID=$(id -u)
SYSTEM_DEPS="apt-utils bash-completion build-essential curl debian-keyring debian-archive-keyring git locales lsb-release openssh-server python3-dev software-properties-common sudo vim wget x11-apps"

##############################################################################
# Debug
##############################################################################

pretty_print ""
pretty_print "Variables"
pretty_print ""
pretty_print "  DOCKERFILE............${DOCKERFILE}"
pretty_print "  UBUNTU_IMAGE..........${UBUNTU_IMAGE}"
pretty_print "  DOCKER_IMAGE..........${DOCKER_IMAGE}"
pretty_print "  USER..................${USER}"
pretty_print "  USER_ID...............${USER_ID}"
pretty_print "  SYSTEM_DEPS...........${SYSTEM_DEPS}"
pretty_print ""
pretty_print "---------------------------------------------------------------"
pretty_print ""

##############################################################################
# Dockerfile
##############################################################################

pretty_header "Dockerfile"

# This is a quoted heredoc ('EOL') that will escape special characters so the
# bash variables don't get substituted here.
#  https://stackoverflow.com/questions/45047380/write-text-to-file-literally-including-special-characters
#
# - DEBIAN_FRONTEND chooses the method to configure debs when installing, prefer not to set it permanenently via ENV
#
# export DEBIAN_FRONTEND=noninteractive && \
# ENV DEBIAN_FRONTEND noninteractive

cat << EOL > ${DOCKERFILE}

# Ubuntu images are listed at https://hub.docker.com/_/ubuntu/, can also use 'latest'
FROM ${UBUNTU_IMAGE}
LABEL maintainer="Daniel Stonier"

# Build Arguments
ARG WORKSPACE_DIR="/home/${USER}/workspace"

# Install system packages
RUN apt-get update
# Install this first - will result in only one spurious warning that apt-utils is not installed
RUN apt-get install -y --no-install-recommends apt-utils
RUN DEBIAN_FRONTEND=noninteractive apt-get install -y ${SYSTEM_DEPS}

# System environment
ENV TERM linux
RUN locale-gen en_US.UTF-8
ENV LANG en_US.UTF-8

# NVIDIA runtime environment
ENV NVIDIA_VISIBLE_DEVICES \\${NVIDIA_VISIBLE_DEVICES:-all}
ENV NVIDIA_DRIVER_CAPABILITIES \\
   \\${NVIDIA_DRIVER_CAPABILITIES:+\\$NVIDIA_DRIVER_CAPABILITIES,}graphics

# Setup user
RUN adduser --quiet --disabled-password \
    --shell /bin/bash --home /home/${USER}  \
    --gecos "Daniel Stonier" ${USER} \
    --uid ${USER_ID}
RUN echo "${USER} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers
RUN usermod -a -G sudo ${USER}

RUN echo "I'm grooty, you should be too"

USER ${USER}
# USER root

# Not sure if this is the best option, just what my non-docker user has, needed by some shell commands.
ENV TERM xterm-256color

SHELL ["/bin/bash", "--login", "-i"]

WORKDIR \${WORKSPACE_DIR}

# Used instead of CMD here so it can't be overridden on the command line
ENTRYPOINT ["/bin/bash", "--login", "-i"]
EOL

##############################################################################
# Build Image
##############################################################################

pretty_header "Docker Image"

#   --rm: remove intermediate objects
#   -f: filename
#   -t: a label for the image (name:tag)

docker image build \
    --rm -t ${DOCKER_IMAGE} \
    ${DOCKERFILE_DIR}
