#!/usr/bin/env bash
set -Eeuo pipefail

if (($# < 2)); then
  echo "usage: $0 <cp9-session> <remote-command...>" >&2
  exit 64
fi

session=$1
shift
repo_root="${CP9_REPO_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)}"
state_dir="${CP9_STATE_DIR:-$repo_root/.cp9/$session}"
host_json="$state_dir/host.json"

[[ -f "$host_json" ]] || {
  echo "cp9 host state is missing: $host_json" >&2
  exit 1
}

host_field() {
  python3 - "$host_json" "$1" <<'PY'
import json
import sys

value = json.load(open(sys.argv[1], encoding="utf-8"))
for key in sys.argv[2].split("."):
    value = value.get(key, "") if isinstance(value, dict) else ""
print(str(value or "").strip())
PY
}

[[ "$(host_field mode)" == "remote-ec2-k3d" ]] || {
  echo "remote cp9 host state is required: $host_json" >&2
  exit 1
}

ssh_user=$(host_field remote.ssh_user)
ssh_key=$(host_field remote.ssh_private_key_path)
ssh_host=$(host_field remote.public_dns)
[[ -n "$ssh_host" ]] || ssh_host=$(host_field remote.public_ip)
transport=$(host_field remote.ssh_transport)
[[ -n "$transport" ]] || transport=direct
connect_timeout=10
proxy_command=

case "$transport" in
  direct) ;;
  ssm)
    connect_timeout=30
    profile=$(host_field remote.profile)
    [[ -n "$profile" ]] || profile=$(host_field profile)
    region=$(host_field remote.region)
    [[ -n "$region" ]] || region=$(host_field region)
    instance_id=$(host_field remote.instance_id)
    [[ "$profile" =~ ^[A-Za-z0-9_+=,.@-]+$ && "$region" =~ ^[a-z0-9-]+$ && "$instance_id" =~ ^i-[0-9a-f]+$ ]] || {
      echo "remote SSM transport identity is invalid in $host_json" >&2
      exit 1
    }
    proxy_command="ProxyCommand=aws ssm start-session --profile $profile --region $region --target $instance_id --document-name AWS-StartSSHSession --parameters portNumber=%p"
    ;;
  *)
    echo "unsupported remote SSH transport: $transport" >&2
    exit 1
    ;;
esac

for value in "$ssh_user" "$ssh_key" "$ssh_host"; do
  [[ -n "$value" ]] || {
    echo "remote cp9 host state is incomplete: $host_json" >&2
    exit 1
  }
done

ssh_args=(
  -o BatchMode=yes
  -o StrictHostKeyChecking=no
  -o UserKnownHostsFile=/dev/null
  -o LogLevel=ERROR
  -o "ConnectTimeout=$connect_timeout"
  -o ConnectionAttempts=3
  -o ServerAliveInterval=30
  -o ServerAliveCountMax=3
  -o ControlMaster=no
  -o ControlPath=none
)
[[ -z "$proxy_command" ]] || ssh_args+=(-o "$proxy_command")
printf -v remote_command '%q ' "$@"
exec ssh "${ssh_args[@]}" -i "$ssh_key" "$ssh_user@$ssh_host" "$remote_command"