#!/usr/bin/env bash

DOCKER_DS_DONT_PULL=${DOCKER_DS_DONT_PULL:-pull}
DOCKER_DS_DIFFS=${DOCKER_DS_DIFFS:-no_diffs}

DOCKER_ENV=""
NOTEBOOK_IMAGE=resero/docker-ds

. dock-sync

# Parse command line arguments in any order
gflag=''
regflag='true'
while getopts 'g' flag; do    # if a character is followed by a colon, that argument is expected to have an argument.
  case "${flag}" in
    g) gflag='-g';;
    *) error "Unexpected option ${flag}" ;;
  esac
done

DOCKER_TAG='latest'
GPU_OPTS=''
if [ -n "$gflag" ]; then
    DOCKER_TAG='gpu'
    GPU_OPTS='--runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all'
fi

if [ "$DOCKER_DS_DIFFS" != "no_diffs" ]; then
    DOCKER_ENV="-e DOCKER_DS_DIFFS=1"
fi

AWS_ENV_NAMES=('AWS_ACCESS_KEY_ID' 'AWS_SECRET_ACCESS_KEY' 'AWS_DEFAULT_REGION' 'AWS_REGION' 'AWS_SESSION_TOKEN' 'AWS_SECURITY_TOKEN')
AWS_ENV=''
for aws_env in ${AWS_ENV_NAMES[@]}; do
    if [ ! -z "${!aws_env}" ]; then
        AWS_ENV+=" -e $aws_env=\${$aws_env}"
    fi
done

if [ -d "./docker/notebook" ]; then
    build-image -f notebook
    if [ -z "$DOCKER_HOST" ]
    then
        run-image $gflag notebook
    else
        sync-up
        run-image $gflag notebook
        sync-down
    fi
else
    # Script to run resero-labs/docker-ds notebok against the current directory
    if [ "$DOCKER_DS_DONT_PULL" = "pull" ]; then
        docker pull $NOTEBOOK_IMAGE:$DOCKER_TAG
    fi

    date_stamp=$(date "+%Y_%m_%d_%H.%M.%S")
    if [ -z "$DOCKER_HOST" ]
    then
        if [[ ! -d "./.ipynb_checkponts" ]]; then
            mkdir ./.ipynb_checkpoints
        fi
        if [[ -d "/data" ]]; then
            if [[ -d "$HOME/.aws" ]]; then
                volume_mounts="--mount type=bind,source=$(pwd),target=/home/jovyan/project -v /data:/data --mount type=bind,source=$HOME/.aws,target=/home/jovyan/.aws"
            else
                volume_mounts="--mount type=bind,source=$(pwd),target=/home/jovyan/project -v /data:/data"
            fi
        else
            if [[ -d "$HOME/.aws" ]]; then
                volume_mounts="--mount type=bind,source=$(pwd),target=/home/jovyan/project --mount type=bind,source=$HOME/.aws,target=/home/jovyan/.aws"
            else
                volume_mounts="--mount type=bind,source=$(pwd),target=/home/jovyan/project"
            fi
        fi
        cmd='docker run '$GPU_OPTS' --init --name '$USER'_notebook_'$date_stamp' '$DOCKER_ENV' -e NOTEBOOK_MODE=lab '$AWS_ENV' --rm -it '$volume_mounts' -p 8888:8888 '$NOTEBOOK_IMAGE':'$DOCKER_TAG
        echo $cmd
        eval $cmd
    else
        sync-up
        cmd='docker run '$GPU_OPTS' --init --name '$USER'_notebook_'$date_stamp' '$DOCKER_ENV' -e NOTEBOOK_MODE=lab '$AWS_ENV' --rm -it --mount type=bind,source=/data/workspaces/'$USER'/code/'${PWD##*/}',target=/home/jovyan/project -v /data:/data -p 8888:8888 '$NOTEBOOK_IMAGE':'$DOCKER_TAG
        echo $cmd
        eval $cmd
        sync-down
    fi
fi
