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

usage() {
    cat <<'EOF'
Usage: daylily-ssh-into-headnode [options]

Options:
  --profile NAME   AWS profile to use (sets AWS_PROFILE for this invocation)
  --region REGION  AWS region (skip interactive region selection)
  --cluster NAME   Cluster name (skip interactive cluster selection)
  --pem FILE       Path to PEM key file (skip interactive PEM selection)
  --user USER      SSH username (default: ubuntu)
  -h, --help       Show this help message and exit

When flags are omitted the script falls back to interactive selection.
EOF
}

need_cmd() {
    command -v "$1" >/dev/null 2>&1 || {
        echo "Error: missing required command: $1" >&2
        exit 3
    }
}

flag_profile=""
flag_region=""
flag_cluster=""
flag_pem=""
flag_user="ubuntu"

while [[ $# -gt 0 ]]; do
    case "$1" in
        --profile) flag_profile="$2"; shift 2 ;;
        --region)  flag_region="$2"; shift 2 ;;
        --cluster) flag_cluster="$2"; shift 2 ;;
        --pem)     flag_pem="$2"; shift 2 ;;
        --user)    flag_user="$2"; shift 2 ;;
        -h|--help) usage; exit 0 ;;
        *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
    esac
done

if [[ -n "$flag_profile" ]]; then
    export AWS_PROFILE="$flag_profile"
fi

if [[ -z "${AWS_PROFILE:-}" ]]; then
    echo "Error: AWS_PROFILE is not set. Provide --profile or export AWS_PROFILE." >&2
    exit 2
fi

need_cmd aws
need_cmd pcluster
need_cmd python3

DEFAULT_REGION="us-west-2"

region="$flag_region"
if [[ -z "$region" ]]; then
    regions=$(aws ec2 describe-regions \
        --profile "$AWS_PROFILE" \
        --region "$DEFAULT_REGION" \
        --query "Regions[].RegionName" \
        --output text 2>/dev/null | tr '\t' '\n' || true)
    if [[ -z "$regions" ]]; then
        echo "Error: unable to retrieve AWS regions. Check AWS_PROFILE=$AWS_PROFILE." >&2
        exit 3
    fi

    region_array=()
    while IFS= read -r r; do region_array+=("$r"); done <<<"$regions"
    echo "Select an AWS region:"
    select region in "${region_array[@]}"; do
        [[ -n "$region" ]] && break
        echo "Invalid selection. Please try again."
    done
fi

cluster="$flag_cluster"
if [[ -z "$cluster" ]]; then
    clusters_json=$(pcluster list-clusters --region "$region" 2>/dev/null || true)
    cluster_names=$(python3 - <<'PY' <<<"$clusters_json"
import json, sys
try:
    data = json.load(sys.stdin)
except Exception:
    sys.exit(0)
clusters = data.get("clusters", []) or []
for c in clusters:
    name = c.get("clusterName")
    if name:
        print(name)
PY
)
    if [[ -z "$cluster_names" ]]; then
        echo "Error: no clusters found (or cannot list clusters) in region $region." >&2
        exit 3
    fi

    cluster_array=()
    while IFS= read -r n; do cluster_array+=("$n"); done <<<"$cluster_names"

    if [[ ${#cluster_array[@]} -eq 1 ]]; then
        cluster="${cluster_array[0]}"
        echo "Only one cluster found: $cluster. Auto-selecting it."
    else
        echo "Select a cluster name:"
        select cluster in "${cluster_array[@]}"; do
            [[ -n "$cluster" ]] && break
            echo "Invalid selection. Please try again."
        done
    fi
fi

head_ip=$(pcluster describe-cluster -n "$cluster" --region "$region" 2>/dev/null \
    | python3 - <<'PY'
import json, sys
try:
    d = json.load(sys.stdin)
except Exception:
    print("")
    sys.exit(0)
print((d.get("headNode") or {}).get("publicIpAddress") or "")
PY
)
if [[ -z "$head_ip" ]]; then
    echo "Error: could not determine head node public IP for cluster $cluster in $region." >&2
    exit 3
fi

pem="$flag_pem"
if [[ -z "$pem" ]]; then
    shopt -s nullglob
    pem_files=(~/.ssh/*.pem)
    shopt -u nullglob
    if [[ ${#pem_files[@]} -eq 0 ]]; then
        echo "Error: no PEM files found in ~/.ssh. Provide --pem." >&2
        exit 2
    fi
    if [[ ${#pem_files[@]} -eq 1 ]]; then
        pem="${pem_files[0]}"
        echo "Only one PEM found: $pem. Auto-selecting it."
    else
        echo "Select a PEM file:"
        select pem in "${pem_files[@]}"; do
            [[ -n "$pem" ]] && break
            echo "Invalid selection. Please try again."
        done
    fi
fi

if [[ ! -f "$pem" ]]; then
    echo "Error: PEM file does not exist: $pem" >&2
    exit 2
fi

echo "Connecting to ${flag_user}@${head_ip} (cluster=$cluster region=$region profile=$AWS_PROFILE)"
exec ssh -i "$pem" \
    -o StrictHostKeyChecking=no \
    -o UserKnownHostsFile=/dev/null \
    "${flag_user}@${head_ip}"

