# Description:
# This script will help the user to setup rootless docker[1] on CMLab Server.
# We follow our CMLab's rootless docker guideline[2] to build the environment:
# 
# We will add the user's username into subuid & subgid so that the docker service
# can start the containers. Then we run the script provided by docker to install
# rootless docker's service.
# 
# After checking whether the docker service can start, we should allow Docker
# can access NVIDIA Container Runtime so that our Docker container can use GPUs.
#
# Then, we should change the directory of Docker to /tmp2 so that Docker can save
# images.
#
# [1] rootless docker's documentation: https://docs.docker.com/engine/security/rootless/
# [2] CMLab's rootless docker guideline: https://hackmd.io/@snsd0805/SkWMyB1ER
# 
# If changed this scripts, you should copy it to every CML Server's /usr/bin/ directory


# utils functions

function print_info {
	local msg="$1"
	echo -e -n "\033[44m[CML INFO]\033[0m "
	echo "$msg"
}

function print_err {
	local msg="$1"
	echo -e -n "\033[41m[CML INFO]\033[0m "
	echo "$msg"
}

function run_command_check {
	local status="$1"
	local success_prefix="$2"
	local error_prefix="$3"
	if [ $status -eq 0 ]; then
		print_info "${success_prefix}"
	else
		print_err "${error_prefix}"
		exit 1
	fi
}

# start
print_info "Start setting up Docker service"

# insert the username into /etc/subuid & /etc/subuid
# setup-uid's source code is on https://gitea.snsd0805.com/snsd0805/CMLAB-Docker-Builder/src/branch/master/tool/setup_uid.c
setup-uid
run_command_check $? "Inserted $(whoami) into /etc/subuid & /etc/subgid" "failed to set subuid & subgid. Please contact us: unix_manager@cmlab.csie.ntu.edu.tw"

# Run install script provided by docker
if [ ! -d "/tmp2/$(whoami)/docker/bin" ]; then
	curl -fsSL https://get.docker.com/rootless | DOCKER_BIN=/tmp2/$(whoami)/docker/bin sh
	run_command_check $? "Docker installed." "failed to install rootless docker. Please contact us: unix_manager@cmlab.csie.ntu.edu.tw"
	export PATH="/tmp2/$(whoami)/docker/bin:$PATH"
	export DOCKER_HOST="unix:///run/user/$(id -u)/docker.sock"
	echo "export PATH=/tmp2/$(whoami)/docker/bin:\$PATH" >> ~/.bashrc
	echo "export DOCKER_HOST=unix:///run/user/$(id -u)/docker.sock" >> ~/.bashrc
	print_info "We have add some environment variables in your ~/.bashrc"
else
	print_info "Docker has been installed. We skip this step."	
fi

# Link NVIDIA Container Runtime to Docker & change docker image's path
if [ ! -d /tmp2/$(whoami)/docker-data ]; then
	mkdir /tmp2/$(whoami)/docker-data
fi
cat << EOF > $HOME/.config/docker/daemon.json
{
    "runtimes": {
        "nvidia": {
            "args": [],
            "path": "nvidia-container-runtime"
        }
    },
    "data-root": "/tmp2/$(whoami)/docker-data"
}
EOF
run_command_check $? "Linked NVIDIA Container Runtime & changed the data path." "failed to link NVIDIA Container Runtime to Docker. Please contact us: unix_manager@cmlab.csie.ntu.edu.tw"

# restart docker
print_info "Restarting Docker & test..."
systemctl --user restart docker && docker pull ubuntu && docker run --rm --runtime=nvidia --gpus '"device=0,1"' ubuntu:22.04 nvidia-smi

print_info "These are 2 GPUs on this server."
print_info "Complete all install steps."
print_info "If you have any problems, please contact us!"

echo 
echo 'CMLab Unix Manager'
echo 'CMLab, National Taiwan University'
echo 'Email: unix_manager@cmlab.csie.ntu.edu.tw'
echo '            __  __     _             '
echo '           / / / /__  (_)_ __        '
echo '   __  ___/ /_/ / _ \/ /\ \ /        '
echo '  /  |/  /\____/_//_/_//_\_\____ ____'
echo ' / /|_/ / _ `/ _ \/ _ `/ _ `/ -_) __/'
echo '/_/  /_/\_,_/_//_/\_,_/\_, /\__/_/   '
echo '                      /___/   	   '