#!/bin/bash -l
# Very important to start with 'bash -l' - to escape SageMaker Studio notebook environment

set -e

dir=$(dirname "$0")
source "$dir"/sm-helper-functions

if [[ "$1" == "configure" ]]; then
    # Everything that requires root and Internet goes to 'configure'
    if [[ -f /opt/sagemaker-ssh-helper/.ssh-ide-configured ]]; then
        echo "sm-ssh-ide: Already configured. Remove /opt/sagemaker-ssh-helper/.ssh-ide-configured to force reconfiguration."
        exit 0
    fi

    OPTIONS=$2
    mkdir -p /opt/sagemaker-ssh-helper/
    echo "$OPTIONS" > /opt/sagemaker-ssh-helper/.sm-ssh-ide-options

    _install_helper_scripts

    cat >/etc/profile.d/sm-ssh-ide.sh <<EOF
export XAUTHORITY="/tmp/.Xauthority"
export ICEAUTHORITY="/tmp/.ICEauthority"
touch "/tmp/.Xauthority"
touch "/tmp/.ICEauthority"
chmod 600 "/tmp/.Xauthority"
chmod 600 "/tmp/.ICEauthority"
EOF
    source /etc/profile.d/sm-ssh-ide.sh

    sm-setup-ssh configure

    if [[ "$OPTIONS" != "--ssh-only" ]]; then

      echo "sm-ssh-ide: Installing VNC and GUI apps"

      if _is_centos; then

        # Based on https://repost.aws/knowledge-center/ec2-linux-2-install-gui
        amazon-linux-extras install mate-desktop1.x
        echo PREFERRED=/usr/bin/mate-session > /etc/sysconfig/desktop

        yum install -y tigervnc-server
        mkdir -p /etc/tigervnc/
        echo localhost > /etc/tigervnc/vncserver-config-mandatory

        amazon-linux-extras install epel
        yum install -y chromium

      else

        apt-get -qq -y install mate-desktop-environment tigervnc-standalone-server

        echo "/usr/bin/mate-session" > ~/.xsession
        chmod +x ~/.xsession

      fi

    else
      echo "sm-ssh-ide: Skipping VNC and GUI apps install"
    fi

    touch /opt/sagemaker-ssh-helper/.ssh-ide-configured

    echo "sm-ssh-ide: Configuration complete."

elif [[ "$1" == "get-metadata" ]]; then

  echo "App name: $(_print_sm_app_name)"
  echo "User profile name: $(_print_sm_user_profile_name)"
  echo "Domain: $(_print_sm_domain_id)"
  echo ""
  echo "Connect from local machine (with GUI and Jupyter): sm-local-ssh-ide connect $(_print_sm_app_name)"
  echo "SSH only: sm-local-ssh-ide connect $(_print_sm_app_name) --ssh-only"

elif [[ "$1" == "set-jb-license-server" ]]; then

    JB_LICENSE_SERVER_HOST="$2"

    if [ -f ~/.sm-jb-license-server ]; then
        echo "sm-ssh-ide: PyCharm license server host is already configured in ~/.sm-jb-license-server, skipping override"
    else
        echo "sm-ssh-ide: Saving PyCharm License server host into ~/.sm-jb-license-server"
        echo "$JB_LICENSE_SERVER_HOST" > ~/.sm-jb-license-server
    fi

    JB_LICENSE_SERVER_HOST="$(cat ~/.sm-jb-license-server)"
    # TODO: check that it's not started with 'http'

    if grep -q "$JB_LICENSE_SERVER_HOST" /etc/hosts; then
        echo "sm-ssh-ide: Skipping the update of /etc/hosts with PyCharm license server (already there)"
    else
        echo "sm-ssh-ide: Updating /etc/hosts with PyCharm license server from ~/.sm-jb-license-server"
        echo "127.0.0.1  $JB_LICENSE_SERVER_HOST" >> /etc/hosts
    fi

elif [[ "$1" == "set-vnc-password" ]]; then
    OPTIONS=$(cat /opt/sagemaker-ssh-helper/.sm-ssh-ide-options || echo "")

    if [[ "$OPTIONS" == "--ssh-only" ]]; then
      echo "sm-ssh-ide: Setting VNC password for SSH-only configuration has no effect, exiting."
      exit 0
    fi

    VNC_PASSWORD="$2"

    if [ -f ~/.vnc/passwd ]; then
        echo "sm-ssh-ide: VNC password is already set in ~/.vnc/passwd, skipping override"
    else
        echo "sm-ssh-ide: Encrypting and saving VNC password to ~/.vnc/passwd"
        mkdir -p ~/.vnc
        echo "$VNC_PASSWORD" | vncpasswd -f > ~/.vnc/passwd
        chmod 600 ~/.vnc/passwd
    fi

elif [[ "$1" == "set-local-user-id" ]]; then

    LOCAL_USER_ID="$2"

    if [ -f ~/.sm-ssh-owner ]; then
        echo "sm-ssh-ide: Local user ID is already set in ~/.sm-ssh-owner, skipping override"
        LOCAL_USER_ID="$(cat ~/.sm-ssh-owner)"
    else
        echo "sm-ssh-ide: Saving local user ID into ~/.sm-ssh-owner"
        echo "$LOCAL_USER_ID" > ~/.sm-ssh-owner
    fi

elif [[ "$1" == "init-ssm" ]]; then
    LOCAL_USER_ID="$(cat ~/.sm-ssh-owner)"
    echo "sm-ssh-ide: Will use local user ID: $LOCAL_USER_ID"

    user_profile_json=$(aws sagemaker describe-user-profile \
      --domain "$(_print_sm_domain_id)" \
      --user-profile-name "$(_print_sm_user_profile_name)" \
      --output json \
      | tr -d "\n")
    execution_role=$(echo "$user_profile_json" | grep "ExecutionRole" \
      | sed -e 's/^.*"ExecutionRole": \"\([^"]*\)\".*$/\1/')

    SSH_SSM_ROLE=$(echo "$execution_role" | sed -e 's/^.*:role\/\(.*\)/\1/')
    echo "sm-ssh-ide: Will use SSM role: $SSH_SSM_ROLE"

    export SSH_SSM_ROLE
    export SSH_OWNER_TAG=$LOCAL_USER_ID

    # Init SSM and pipe the output to CloudWatch (stdout of the pid 1).
    sm-init-ssm | tee /proc/1/fd/1

elif [[ "$1" == "start" ]]; then
    OPTIONS=$(cat /opt/sagemaker-ssh-helper/.sm-ssh-ide-options || echo "")

    if [[ -f /tmp/.ssh-ide-local-lock ]]; then
        echo "sm-ssh-ide: Already running on this instance? Call stop first." && exit 1
    fi

    touch /tmp/.ssh-ide-local-lock

    echo "sm-ssh-ide: Saving env variables for remote SSH interpreter"
    sm-save-env

    echo "sm-ssh-ide: Starting SSH service"
    if _is_centos; then
      /usr/sbin/sshd
    else
      service ssh start || (echo "sm-ssh-ide: ERROR: Failed to start ssh service" && exit 255)
    fi

    if [[ "$OPTIONS" != "--ssh-only" ]]; then

      echo "sm-ssh-ide: Starting VNC server"

      export USER=root
      vncserver :1

      SM_STUDIO_PYTHON=$(_print_sm_studio_python)
      echo "sm-ssh-ide: Using SageMaker Studio Python: $SM_STUDIO_PYTHON"

      "$SM_STUDIO_PYTHON" -m jupyter notebook --version >/dev/null 2>&1 || (
        # Note: you should have access to Internet or have notebook module pre-installed in the container
        echo "sm-ssh-ide: Installing Jupyter notebook"
        "$SM_STUDIO_PYTHON" -m pip install notebook
        # Note: reinstall of some libraries is needed for compatibility with SageMaker Studio:
        "$SM_STUDIO_PYTHON" -m pip install --force-reinstall traitlets tornado
      )

      echo "sm-ssh-ide: Starting Jupyter notebook server"

      "$SM_STUDIO_PYTHON" -m jupyter notebook --no-browser --port=8889 --ip=127.0.0.1 --allow-root \
        >/tmp/jupyter-notebook.log 2>&1 &

      sleep 2

      tail /tmp/jupyter-notebook.log

    else
      echo "sm-ssh-ide: Skipping VNC and Jupyter start"
    fi

    echo "sm-ssh-ide: Started all services"

elif [[ "$1" == "ssm-agent" ]]; then

    echo "sm-ssh-ide: Starting SSM agent"
    /usr/bin/amazon-ssm-agent

elif [[ "$1" == "status" ]]; then

    netstat -nptl | grep '5901\|8889'

elif [[ "$1" == "get-studio-python-path" ]]; then

    SM_STUDIO_PYTHON=$(_print_sm_studio_python)
    echo "$SM_STUDIO_PYTHON"

elif [[ "$1" == "get-studio-python-version" ]]; then

    SM_STUDIO_PYTHON=$(_print_sm_studio_python)
    $SM_STUDIO_PYTHON --version

elif [[ "$1" == "env-diagnostics" ]]; then

    echo "System Python location: $(which python || echo 'Not found')"
    echo "System Python version: $(python --version || echo 'Not found')"

    echo "Conda location: $(which conda || echo 'Not found')"
    echo "Conda environments: $(conda env list || echo 'Not found')"

    echo "SageMaker Studio Python location: $($0 get-studio-python-path || echo 'Not found')"
    echo "SageMaker Studio Python version: $($0 get-studio-python-version || echo 'Not found')"

    # Should be the same as `jupyter kernelspec list` executed from SageMaker Studio Python
    echo "Jupyter Kernels: $(find "$JUPYTER_PATH" -name 'kernel.json')"
    find "$JUPYTER_PATH" -name 'kernel.json' -print0 | xargs -0 cat

elif [[ "$1" == "stop" ]]; then
    OPTIONS=$(cat /opt/sagemaker-ssh-helper/.sm-ssh-ide-options || echo "")

    echo "sm-ssh-ide: Stopping SSM agent"
    pkill -ef amazon-ssm-agent || echo "sm-ssh-ide: SSM agent already stopped?"

    echo "sm-ssh-ide: Stopping SSH service"
    if _is_centos; then
      pkill -ef sshd || echo "sm-ssh-ide: SSH already stopped?"
    else
      service ssh stop
    fi

    if [[ "$OPTIONS" != "--ssh-only" ]]; then

      echo "sm-ssh-ide: Stopping VNC server"
      vncserver -kill :1 || echo "sm-ssh-ide: VNC server already stopped?"

      # Wait dbus to shop
      sleep 5

      pkill -e Xtigervnc || :
      pkill -e Xtightvnc || :

      echo "sm-ssh-ide: Stopping SSH and GPG agents"
      pkill -e ssh-agent || :
      pkill -e gpg-agent || :

      echo "sm-ssh-ide: Stopping Jupyter notebook server"
      pkill -fe jupyter-notebook || echo "sm-ssh-ide: Jupyter notebook already stopped?"

    else
      echo "sm-ssh-ide: Skipping VNC and Jupyter stop"
    fi

    # Wait all other services to stop
    sleep 5
    
    if [[ -f /tmp/.ssh-ide-local-lock ]]; then
        rm /tmp/.ssh-ide-local-lock
    else
        echo "sm-ssh-ide: Local lock is missing, was not not running on this instance?"
    fi

    echo "sm-ssh-ide: Stopped all services"

else

    echo "sm-ssh-ide: Unknown command: $1"

fi