#!/bin/bash

# NOTE: This is a lower level interface to the SSH Helper.
# Run `sm-ssh -h` for the high-level interface.

# Commands:
# [--domain-id <domain-id>] [--user-profile-name <user-profile-name>] connect-app <app_name> [--ssh-only] [<extra_ssh_args>]
# [--domain-id <domain-id>] [--user-profile-name <user-profile-name>] proxy-host <fqdn>
# run-command <command> <args...>
# set-domain-id <domain_id>
# set-user-profile-name <user_profile_name>
# set-jb-license-server <jb-license-server-hostname-without-http>

# SageMaker Studio app name for Kernel Gateways is usually the same as the hostname,
#   e. g. 'sagemaker-data-science-ml-m5-large-1234567890abcdef0' (without quotes).
# For JupyterServer app it's 'default'.

# Tip: to open SageMaker Studio UI in Firefox from command line on macOS, use the following command:
# open -a Firefox $(AWS_PROFILE=terry aws sagemaker create-presigned-domain-url --domain-id d-lnwlaexample --user-profile-name terry-whitlock --query AuthorizedUrl --output text)


# replace with your JetBrains License Server host, or leave it as is if you don't use one
JB_LICENSE_SERVER_HOST="jetbrains-license-server.example.com"


set -e
set -o pipefail

_on_error() {
  echo 'sm-local-ssh-ide: Terminated.'
}
trap '_on_error' ERR

self=$(realpath "${BASH_SOURCE[0]}")
dir=$(dirname "$self")
source "$dir"/sm-helper-functions 2>/dev/null || source sm-helper-functions

echo "sm-local-ssh-ide: Starting in $dir"

DOMAIN_ID=""
USER_PROFILE_NAME=""

if [[ "$1" == "--domain-id" ]]; then
    DOMAIN_ID=$2
    shift
    shift
fi

if [[ "$1" == "--user-profile-name" ]]; then
    USER_PROFILE_NAME=$2
    shift
    shift
fi

COMMAND=$1

if [[ "$COMMAND" != "set-user-profile-name" && "$COMMAND" != "set-domain-id" \
      && "$DOMAIN_ID" == "" && "$USER_PROFILE_NAME" == "" ]]; then
  if [ -f ~/.sm-studio-domain-id ]; then
    DOMAIN_ID="$(cat ~/.sm-studio-domain-id)"
  else
    echo "sm-local-ssh-ide: WARNING: SageMaker Studio domain ID is not set."\
      "Run 'sm-local-ssh-ide set-domain-id' to override."
  fi
  if [ -f ~/.sm-studio-user-profile-name ]; then
    USER_PROFILE_NAME="$(cat ~/.sm-studio-user-profile-name)"
  else
    echo "sm-local-ssh-ide: WARNING: SageMaker Studio user profile name is not set."\
      "Run 'sm-local-ssh-ide set-user-profile-name' to override."
  fi
fi

if [[ "$COMMAND" == "proxy-host" ]]; then
  SM_SSH_HOST_NAME="$2"
  SM_RESOURCE_TYPE="studio"

  _check_ssh_proxy_host_name "$SM_SSH_HOST_NAME" "$SM_RESOURCE_TYPE"
  _export_ssh_key_env_var "$SM_SSH_HOST_NAME"

  INSTANCE_ID=$(_generate_key_and_print_instance_id "$SM_SSH_HOST_NAME" "$DOMAIN_ID" "$USER_PROFILE_NAME")

  sm-local-start-ssh --proxy-setup-only "${INSTANCE_ID}"

elif [[ "$COMMAND" == "connect-app" ]]; then

  SM_STUDIO_KGW_NAME="$2"
  OPTIONS="$3"

  # shellcheck disable=SC2091
  INSTANCE_ID=$($(_python) <<EOF
import sagemaker; from sagemaker_ssh_helper.ide import SSHIDE;
import logging; logging.basicConfig(level=logging.INFO);
SSHIDE("$DOMAIN_ID", "$USER_PROFILE_NAME").print_kernel_instance_id("$SM_STUDIO_KGW_NAME", timeout_in_sec=300)
EOF
  )

  if [[ "$OPTIONS" == "--ssh-only" ]]; then
    echo "sm-local-ssh-ide: Connecting only SSH to local port 10022 (got the flag --ssh-only)"
    shift
    shift
    shift
    sm-local-start-ssh "$INSTANCE_ID" \
        -L localhost:10022:localhost:22 \
        $*
  else
    if [ -f ~/.sm-jb-license-server ]; then
        JB_LICENSE_SERVER_HOST="$(cat ~/.sm-jb-license-server)"
    fi
    echo "sm-local-ssh-ide: Connecting SSH, VNC and Jupyter to local ports 10022, 5901 and 8889 (add --ssh-only flag to override)"
    shift
    shift
    sm-local-start-ssh "$INSTANCE_ID" \
        -L localhost:10022:localhost:22 \
        -L localhost:5901:localhost:5901 \
        -L localhost:8889:localhost:8889 \
        -R 127.0.0.1:443:"$JB_LICENSE_SERVER_HOST":443 \
        $*
  fi

elif [[ "$COMMAND" == "set-jb-license-server" ]]; then
    JB_LICENSE_SERVER_HOST="$2"

    echo "sm-local-ssh-ide: Saving PyCharm License server host into ~/.sm-jb-license-server"
    echo "$JB_LICENSE_SERVER_HOST" > ~/.sm-jb-license-server

elif [[ "$COMMAND" == "set-domain-id" ]]; then
    DOMAIN_ID="$2"
    if [[ "$DOMAIN_ID" == "" ]]; then
      echo "sm-local-ssh-ide: ERROR: <domain-id> argument is expected"
      exit 1
    fi
    echo "sm-local-ssh-ide: Saving SageMaker Studio domain ID into ~/.sm-studio-domain-id"
    echo "$DOMAIN_ID" > ~/.sm-studio-domain-id

elif [[ "$COMMAND" == "set-user-profile-name" ]]; then
    USER_PROFILE_NAME="$2"
    if [[ "$USER_PROFILE_NAME" == "" ]]; then
      echo "sm-local-ssh-ide: ERROR: <user-profile-name> argument is expected"
      exit 1
    fi
    echo "sm-local-ssh-ide: Saving SageMaker Studio user profile name into ~/.sm-studio-user-profile-name"
    echo "$USER_PROFILE_NAME" > ~/.sm-studio-user-profile-name

elif [[ "$COMMAND" == "run-command" ]]; then

  shift
  ARGS=$*

  # shellcheck disable=SC2086
  ssh -4 -i ~/.ssh/sagemaker-ssh-gw -p 10022 root@localhost \
    -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    $ARGS

else
    echo "ERROR: Unknown command: '$COMMAND'"
    exit 1
fi