#!/bin/bash
set -e

trap 'echo "An error occurred. Please contact the manager." >&2' ERR

echo "================== dnew ==================="
echo "=== contact: xd21@mails.tsinghua.edu.cn ==="

test_not_empty()
{
    if [[ -z "$1" ]]; then
        echo $2 >&2
        exit
    fi
}

ask_for_yN()
{
    if [[ "$DFS_QUIET" == "1" ]]; then
        echo 0
    else
        read -p "$1 [yN]: " yn
        case $yn in
            [Yy]* ) echo 1;;
            * ) echo 0;;
        esac
    fi
}

container_name=$(whoami)_$(date +%y%m%d-%H%M%S)

base_command="sudo docker run -d --net=host -v /home2:/home2 \\
    --cap-add=SYS_NICE --cap-add=IPC_LOCK \\
    --security-opt seccomp=unconfined --ulimit memlock=-1:-1 \\
    --restart=unless-stopped \\
    --name $container_name"


if [[ "$EUID" == "0" ]]; then
    echo "Please run as your own user (i.e., DO NOT sudo)."
    exit 1
fi

ret=$(ask_for_yN "Use GPU?")
if [[ "$ret" == "1" ]]; then
    base_command="$base_command \\
    --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all"
fi
ret=$(ask_for_yN "Use RDMA?")
if [[ "$ret" == "1" ]]; then
    base_command="$base_command \\
    --device=/dev/infiniband/rdma_cm --device=/dev/infiniband/issm0 \\
    --device=/dev/infiniband/umad0 --device=/dev/infiniband/uverbs0"
fi

read -p "Image name (default: git.nasp.fit/nasp/nasp-ubuntu): "
if [[ -z "$REPLY" ]]; then
    image_name="git.nasp.fit/nasp/nasp-ubuntu"
else
    image_name="$REPLY"
fi
base_command="$base_command \\
    -it $image_name"

read -p "Start up command (default: /etc/startup.sh): "
if [[ -z "$REPLY" ]]; then
    cmd="/etc/startup.sh"
else
    cmd="$REPLY"
fi
base_command="$base_command $cmd"

echo "Will run:"
echo "====="
echo "$base_command"
echo "====="
ret=$(ask_for_yN "Start the container?")

if [[ "$ret" == "1" ]]; then
    eval "$base_command"
    echo "Container $container_name started. You can use the following command to get in:"
    echo "sudo docker exec -it $container_name bash"
else
    echo "Aborted."
fi