#!/bin/bash

# Commands:
# connect <processing_job_name> [<extra_ssh_args>]
# stop-waiting
# run-command <command> <args...>

COMMAND=$1

if [[ "$COMMAND" == "connect" ]]; then
  JOB_NAME=$2

  INSTANCE_ID=$(python <<EOF
import sagemaker; from sagemaker_ssh_helper.manager import SSMManager;
import logging; logging.basicConfig(level=logging.INFO);
print(SSMManager().get_processing_instance_ids("$JOB_NAME", timeout_in_sec=300)[0])
EOF
  )

  shift
  shift
  sm-local-start-ssh "$INSTANCE_ID" \
      -R localhost:12345:localhost:12345 \
      -L localhost:14022:localhost:22 \
      $*

elif [[ "$COMMAND" == "stop-waiting" ]]; then

  $0 run-command sm-wait stop

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

  shift
  ARGS=$*

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

else
    echo "Unknown command: $COMMAND"
fi