#!/bin/bash

source $(dirname $0)/shared-run

getDiagnosticSessionRedirection()
{
   if [[ -n "$JUPYTER_DIAGNOSTICS_ENABLED" ]]; then
      echo "|& tee -a $JUPYTER_DIAGNOSTICS_FILE"
   else
      echo ""
   fi
}

getDiagnosticsPrintenv()
{
   if [[ -n "$JUPYTER_DIAGNOSTICS_ENABLED" ]]; then
      echo "echo -e \"[$(date)] ---- launching jupyter process ----\n[$(date)] ---- final env ----\n[$(date)] $(printenv)\" $(getDiagnosticSessionRedirection) &&"
   else
      echo ""
   fi
}

emitDiagnostic()
{
   if [[ -n "$JUPYTER_DIAGNOSTICS_ENABLED" ]]; then
      if ! echo "[$(date)] $1" >> $JUPYTER_DIAGNOSTICS_FILE ; then
         echo "Error: $$ trying to append to $JUPYTER_DIAGNOSTICS_FILE - turning off session diagnostics"
         JUPYTER_DIAGNOSTICS_ENABLED=
      fi
   fi
}

getDiagnosticSessionStraceEnabled()
{
   if [[ -z "$JUPYTER_DIAGNOSTICS_STRACE_ENABLED" ]]; then
      echo "strace not enabled for this diagnostics report"
   elif command -v strace &> /dev/null; then
      echo "strace enabled for this diagnostics report"
      echo "strace report location: $JUPYTER_STRACE_FILE"
      echo ""
   else
      echo "strace is not installed; cannot generate strace diagnostics"
   fi
}

getDiagnosticSessionStrace()
{
   if [[ -z "$JUPYTER_DIAGNOSTICS_STRACE_ENABLED" ]]; then
      echo ""
   elif command -v strace &> /dev/null; then
      echo "strace -fttTy -o $JUPYTER_STRACE_FILE"
   else
      echo ""
   fi
}

runHooks "start"

# determine the config directory
if [ "$RSTUDIO_CONFIG_DIR" != "" ] ; then
   configDirs="$RSTUDIO_CONFIG_DIR"
else
   configDirs="$XDG_CONFIG_DIRS:/etc"
fi

configDirs="$XDG_CONFIG_DIRS:/etc"
IFS=':' read -ra configDirList <<< "$configDirs"
for configDir in "${configDirList[@]}"
do
   if test -e "$configDir/rstudio/jupyter.conf"
   then
      configFile="$configDir/rstudio/jupyter.conf"
      break
   fi
done

# load jupyter-profile script
if test -e "$configDir/rstudio/jupyter-profile"
then
   . "$configDir/rstudio/jupyter-profile"
fi

# session arguments
# escape backslashes in arguments to pass safely to bash below
JUPYTER_ARGS="${@//\\/\\\\}"

# eval variables like $USER in XDG_DATA_HOME and XDG_STATE_HOME
XDG_DATA_HOME=$(eval echo $XDG_DATA_HOME)
XDG_STATE_HOME=$(eval echo $XDG_STATE_HOME)

USER_DATA_HOME=${XDG_DATA_HOME:-${HOME}/.local/share}

# Allow individual users to enable session diagnostics to their local log directory
USER_DIAG_CONFIG="$HOME/.config/pwb/jupyter-diagnostics"
USER_DIAG_LOG_DIR="$USER_DATA_HOME/pwb/log"
if [[ -f "$USER_DIAG_CONFIG" ]] ; then
   USER_DIAG_MODE=$(cat "$USER_DIAG_CONFIG")
   if [[ "$USER_DIAG_MODE" =~ debug|strace ]] ; then
      export JUPYTER_DIAGNOSTICS_ENABLED=true
      mkdir -p "$USER_DIAG_LOG_DIR"
      export JUPYTER_DIAGNOSTICS_FILE="$USER_DIAG_LOG_DIR/jupyter-diagnostics.log"
      JUPYTER_STRACE_FILE="${JUPYTER_DIAGNOSTICS_FILE}.strace"
      if [[ "$USER_DIAG_MODE" = "strace" ]] ; then
         export JUPYTER_DIAGNOSTICS_STRACE_ENABLED=true
      fi

      LOG_PATTERN="Session.debug"
      LOG_ARG="--Session.debug=True"

      if [[ ! "$JUPYTER_ARGS" =~ .*"$LOG_PATTERN".* ]] ; then 
         JUPYTER_ARGS="$JUPYTER_ARGS $LOG_ARG"
      fi
      echo "Enabling $USER_DIAG_MODE diagnostics from config file $USER_DIAG_CONFIG" 
      echo "Diag files: $JUPYTER_DIAGNOSTICS_FILE + .strace"
   elif [[ "$USER_DIAG_MODE" != "off" ]] ; then
      echo "Ignoring invalid session-diagnostics value: $USER_DIAG_MODE for file: $USER_DIAG_CONFIG - use one of strace, debug or off"
   fi
fi

if [[ -n "${RS_SESSION_SSL_CERT}" ]]; then
   if [[ -z "${RS_SESSION_SSL_CERT_KEY}" ]]; then
      echo "Error - public key without private key"
      exit 10
   fi

   CERT_DIR="/tmp/pwb-${USER}"
   mkdir -m 700 -p "${CERT_DIR}"

   CERT_ID=$(date +%s.%N)
   CERT_FILE="${CERT_DIR}/cert-${CERT_ID}.pem"
   KEY_FILE="${CERT_DIR}/key-${CERT_ID}.pem"
   rm -f ${CERT_FILE} ${KEY_FILE}
   (umask 077;
    echo "${RS_SESSION_SSL_CERT}" > "${CERT_FILE}";
    echo "${RS_SESSION_SSL_CERT_KEY}" > "${KEY_FILE}")
   JUPYTER_ARGS="${JUPYTER_ARGS} --certfile ${CERT_FILE} --keyfile ${KEY_FILE}"
elif [[ -n "${RS_SESSION_SSL_CERT_PATH}" ]] ; then
   # eval variables like $HOME and $USER in RS_SESSION_CERT_PATH
   eval CERT_PATH=`echo ${RS_SESSION_SSL_CERT_PATH}`
   eval KEY_PATH=`echo ${RS_SESSION_SSL_CERT_KEY_PATH}`
   JUPYTER_ARGS="${JUPYTER_ARGS} --certfile ${CERT_PATH} --keyfile ${KEY_PATH}"
fi

if ! command -v ${JUPYTER_EXEC_COMMAND} &> /dev/null
then
    echo "jupyter command: '${JUPYTER_EXEC_COMMAND}' not found. Aborting."
    exit 1
fi

manage_jlab_script=$(dirname $(readlink -f $0))/../extras/scripts/manage-jlab-ext.sh
if [ -f "$manage_jlab_script" ]; then
   source ${manage_jlab_script}
   setupEnvironment ${JUPYTER_EXEC_COMMAND}

   MAJOR_VERSION=$(echo "${JNOTEBOOK_VERSION}" | cut -d. -f1)
   if [[ "${JUPYTER_ARGS}" == "lab"* || "${MAJOR_VERSION}" -ge 7 ]]; then
      # check if the expected version of the jupyterlab extension has been installed
      echo "Checking if the extension ${JLAB_EXT_VERSION_STRING} is installed..."
   
      if [ "$(echo $(${JUPYTER_EXEC_COMMAND} labextension list 2>&1) | grep -c ${JLAB_EXT_NAME}[^0-9]*${JLAB_EXT_VERSION})" -lt 1 ]; then
         # locate pip from jupyter's path
         if [ "${JUPYTER_EXEC_COMMAND}" == "jupyter" ]; then
            # using the PATH
            JUPYTER_PATH=$(which jupyter)
         else
            JUPYTER_PATH=${JUPYTER_EXEC_COMMAND}
         fi
         # evaluate symlinks
         JUPYTER_PATH=$(readlink -f ${JUPYTER_PATH})

         pip_path=$(dirname ${JUPYTER_PATH})/pip
         if [ -f "${pip_path}" ]; then
            echo "Installing the Workbench JupyterLab extension..."
            ${pip_path} install ${JLAB_EXT_VERSION_STRING} || echo "ERROR: pip failed to install ${JLAB_EXT_VERSION_STRING}" 
         else
            echo "Could not install the Workbench JupyterLab extension - could not locate pip from Jupyter path. \
            To manually install the extension, run 'python -m pip install ${JLAB_EXT_VERSION_STRING}' \
            using the same version of python that '${JUPYTER_EXEC_COMMAND}' relies on"
         fi
      else
         echo "The extension ${JLAB_EXT_VERSION_STRING} has already been installed"
      fi
   fi
else
   echo "Skipping Jupyter maintenance because $manage_jlab_script was not found."
fi

emitDiagnostic "Running session: ${JUPYTER_EXEC_COMMAND} ${JUPYTER_ARGS}"

echo "Running session: ${JUPYTER_EXEC_COMMAND} ${JUPYTER_ARGS}"

# SIGTERM is used to stop the session - so the hooks need to be run as a trap
trap 'runHooks "stop"' SIGTERM

# run the session
/bin/bash --login $JUPYTER_PROFILE_OPTIONS -c "$(getDiagnosticsPrintenv) $(getDiagnosticSessionStrace) $JUPYTER_EXEC_COMMAND $JUPYTER_ARGS"

# store exit code of command
EXIT_CODE=$?

if [[ -n "${CERT_DIR}" ]]; then
   rm -f "${CERT_FILE}"
   rm -f "${KEY_FILE}"
fi

# exit with previously-stored exit code, after cleanup
exit $EXIT_CODE
