#!/bin/bash -l

jupyter_output=$(mktemp ~/.jupyter_output.XXXX)

clean_up() {
        job_id=$(awk '/srun: job/ {print $3;exit}' $jupyter_output)
        scancel $job_id
	rm $jupyter_output
        echo "Terminated SLURM job $job_id"
}

trap clean_up EXIT

set -m

OPT_HEADLESS=
HAS_TIME=

process_args() {
        local arg args i
        args=( "$@" )
        for arg in "$@"; do
                case "${args[i]}" in
                    (-t|--time|--time=*)
                        HAS_TIME=yes
                        ;;
                    (-h|--headless)
                        OPT_HEADLESS=yes
                        continue
                        ;;
                esac
                SLURM_ARGS=( "${SLURM_ARGS[@]}" "$arg" )
        done
        # FIXME: Can do more thorough arg checking later.
}

# Allow several choices:
# - anaconda3
# - anaconda2
# - python/2.7
# - python/3.x

enable_lmod
if ! module is-loaded slurm; then module load slurm; fi

if [ -n "$ANACONDA3_VER" ]; then
        echo "Launching Jupyter for anaconda3 module version $ANACONDA3_VER"
elif [ -n "$ANACONDA2_VER" ]; then
        echo "Launching Jupyter for anaconda2 module version $ANACONDA2_VER"
elif [ -n "$PYTHON_VER" ]; then
        echo "Launching Jupyter for python module version $PYTHON_VER"
        echo " * loading jupyter module"
        if ! module is-loaded numpy; then module load numpy; fi
        if ! module is-loaded ipython; then module load ipython; fi
        if ! module is-loaded jupyter; then module load jupyter; fi
elif [ x"$PYTHON_VER" = x ]; then
        PYTHON_VER_DEFAULT=3.6
        echo "Warning: Loading default Python module version $PYTHON_VER_DEFAULT"
        echo "Warning: Please load python and any needed python modules before launching juypter!"
        echo "Example: "
        echo "  enable_lmod"
        echo "  module load python/3.6   # 2.7/3.7 is also avaiable"
        sleep 5s
        #enable_lmod
        module load "python/$PYTHON_VER_DEFAULT"
        echo " * loading jupyter module"
        module load jupyter
fi

if [ -z "$(which jupyter-notebook)" ]; then
        echo "Error: cannot find jupyter-notebook program" >&2
        echo "Please contact site administrator with the following diagnosis:" >&2
        echo
        module avail
        echo
        echo "PATH=$PATH"
        exit 1
fi

# FIXME::
SLURM_ARGS=( )
process_args "$@"

echo 
if [ x"$HAS_TIME" != xyes ]; then
        SLURM_ARGS=( -t 240 "${SLURM_ARGS[@]}" )
        echo "   jupyter server will be TERMINATED after 4 hour"
fi
echo "   please do not close this terminal window, or jupyter will be terminated immediately"
echo 

module list

echo " * launching jupyter notebook server"
(srun -n 1 -c 2 -J juypter "${SLURM_ARGS[@]}" jupyter-notebook --no-browser --port=$((8000 + $RANDOM % 1000)) --ip=0.0.0.0 2>&1 | tee $jupyter_output > /dev/null) &

echo " * waiting for Jupyter server to be ready ..."
for ((i=0; i < 600; ++i)); do
        tokens=$(awk -F: '/  http:/ {print $3}'        $jupyter_output)
        if [ -n "$tokens" ]; then
                break
        fi
        sleep 1s
done
if [ -z "$tokens" ]; then
        echo "Error: Cannot find the Jupyter server URL!"
        echo
        echo "Dump of Jupyter process output:"
        echo "-------------------------------"
        cat $jupyter_output
        echo
        echo "Exiting due to error..."
        exit 2
fi

# Don't use echo to avoid divulging the token (!!)
port=${tokens%/*}
token=${tokens#*token=}

# Look for the hostname where the server is running:
job_id=$(awk     '/srun: job/ {print $3;exit}' $jupyter_output)
host=$(squeue -j $job_id -O NodeList -h | sed 's/ //g')
url="http://$host:$tokens"

echo " * SLURM job ID: $job_id"

if [[ -n "$DISPLAY" && "$OPT_HEADLESS" != yes ]]; then
        echo " * launching browser"
        (firefox $url 2>&1 )>/dev/null &
        echo "   if you close your browser, please the link use below to connect back"
        echo "      $url"
else
        echo
        echo "NOTE: Jupyter is running in a headless mode!"
        echo "To connect, please keep an open connection from your local computer using:"
        echo
        echo "    ssh -L 8888:$host:$port $USER@turing.hpc.odu.edu"
        echo
        echo "and point your browser to:"
        echo
        # https://stackoverflow.com/a/15229498/655885
        cat <<EOF
    http://localhost:8888/?token=${token}
EOF
        echo
        echo "(If ssh fails because local port 8888 is taken, use another number.)"
fi

wait