#!/bin/bash

# SageMaker Studio Kernel Gateway name is usually the same as the hostname,
# e. g. datascience-1-0-ml-g4dn-xlarge-afdb4b3051726e2ee18a399903

SM_STUDIO_KGW_NAME="$1"

echo "Querying SSM instance ID for kernel gateway $SM_STUDIO_KGW_NAME"

INSTANCE_ID=$(python <<EOF
import sagemaker; from sagemaker_ssh_helper.log import SSHLog;
import logging; logging.basicConfig(level=logging.INFO);
print(SSHLog().get_studio_kgw_ssm_instance_ids("$SM_STUDIO_KGW_NAME", retry=30)[0])
EOF
)

# 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"

if [ -f ~/.sm-jb-license-server ]; then
    # TODO: extract duplicating logic from sm-ssh-ide
    echo "sm-local-ssh-ide: ~/.sm-jb-license-server file with PyCharm license server host is already configured, skipping override"
    JB_LICENSE_SERVER_HOST="$(cat ~/.sm-jb-license-server)"
else
    echo "sm-local-ssh-ide: Saving PyCharm License server host into ~/.sm-jb-license-server"
    echo "$JB_LICENSE_SERVER_HOST" > ~/.sm-jb-license-server
fi


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
