#!/usr/bin/env bash

set -euo pipefail

# Provision an EC2 box, run the Oak benchmarks on it, download results/,

# then destroy everything it created (instance, security group, key pair).

#

# Usage:

#   scripts/aws_bench.sh [options] [-- <extra bench.py args>]

#

# Options:

#   --profile {micro,smoke,standard,large}   Benchmark profile (default: smoke)

#   --instance-type TYPE   EC2 instance type (default: c6id.2xlarge — has

#                          local NVMe, used as TMPDIR so fixture I/O hits

#                          instance-store SSD, not EBS)

#   --region REGION        AWS region (default: AWS CLI configured region)

#   --disk-gb N            Root EBS volume size in GB (default: 60)

#   --timeout-mins N       Max minutes to wait for the benchmark (default: 360)

#   --keep                 Do not destroy the instance afterwards (prints SSH cmd)

#   -- ...                 Everything after -- is passed through to bench.py

#

# Examples:

#   scripts/aws_bench.sh                                  # smoke run

#   scripts/aws_bench.sh --profile standard

#   scripts/aws_bench.sh --profile large -- --skip-diff

#

# Requires: aws CLI (authenticated — run `aws login` first), ssh, rsync, curl.

# Results land in results/aws/<run-id>/ locally.

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"

BENCH_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

PROFILE=smoke

INSTANCE_TYPE=c6id.2xlarge

REGION="$(aws configure get region 2>/dev/null || true)"

DISK_GB=60

TIMEOUT_MINS=360

KEEP=0

BENCH_EXTRA_ARGS=()

while [ $# -gt 0 ]; do

  case "$1" in

    --profile)        PROFILE="$2"; shift 2 ;;

    --instance-type)  INSTANCE_TYPE="$2"; shift 2 ;;

    --region)         REGION="$2"; shift 2 ;;

    --disk-gb)        DISK_GB="$2"; shift 2 ;;

    --timeout-mins)   TIMEOUT_MINS="$2"; shift 2 ;;

    --keep)           KEEP=1; shift ;;

    --)               shift; BENCH_EXTRA_ARGS=("$@"); break ;;

    -h|--help)        sed -n '4,30p' "$0"; exit 0 ;;

    *) echo "unknown option: $1" >&2; exit 2 ;;

  esac

done

[ -n "$REGION" ] || { echo "no AWS region configured; pass --region" >&2; exit 2; }

RUN_ID="oak-bench-$(date +%Y%m%d-%H%M%S)"

STATE_DIR="$(mktemp -d "${TMPDIR:-/tmp}/${RUN_ID}.XXXX")"

PEM="$STATE_DIR/key.pem"

SSH_OPTS=(-i "$PEM"

  -o StrictHostKeyChecking=accept-new

  -o UserKnownHostsFile="$STATE_DIR/known_hosts"

  -o ConnectTimeout=10

  -o ServerAliveInterval=15

  -o ServerAliveCountMax=8)

SSH_USER=ubuntu

log() { printf '[aws-bench] %s\n' "$*" >&2; }

aws_ec2() { aws ec2 --region "$REGION" "$@"; }

INSTANCE_ID=""

SG_ID=""

KEY_CREATED=0

cleanup() {

  local rc=$?

  set +e

  if [ "$KEEP" = 1 ] && [ -n "$INSTANCE_ID" ]; then

    log "--keep set: leaving $INSTANCE_ID running"

    log "ssh: ssh -i $PEM $SSH_USER@${PUBLIC_IP:-<ip>}"

    log "destroy later with:"

    log "  aws ec2 --region $REGION terminate-instances --instance-ids $INSTANCE_ID"

    log "  aws ec2 --region $REGION delete-security-group --group-id $SG_ID"

    log "  aws ec2 --region $REGION delete-key-pair --key-name $RUN_ID"

    exit "$rc"

  fi

  if [ -n "$INSTANCE_ID" ]; then

    log "terminating $INSTANCE_ID"

    aws_ec2 terminate-instances --instance-ids "$INSTANCE_ID" >/dev/null

    aws_ec2 wait instance-terminated --instance-ids "$INSTANCE_ID"

  fi

  if [ -n "$SG_ID" ]; then

    # SG deletion can lag the instance's ENI teardown; retry briefly.

    for _ in 1 2 3 4 5 6; do

      aws_ec2 delete-security-group --group-id "$SG_ID" 2>/dev/null && break

      sleep 10

    done

  fi

  [ "$KEY_CREATED" = 1 ] && aws_ec2 delete-key-pair --key-name "$RUN_ID" >/dev/null

  rm -rf "$STATE_DIR"

  exit "$rc"

}

trap cleanup EXIT

# --- Preflight -------------------------------------------------------------

aws sts get-caller-identity >/dev/null 2>&1 || {

  echo "AWS credentials missing/expired — run: aws login" >&2; exit 2; }

log "run id: $RUN_ID  region: $REGION  type: $INSTANCE_TYPE  profile: $PROFILE"

AMI_ID="$(aws --region "$REGION" ssm get-parameter \

  --name /aws/service/canonical/ubuntu/server/24.04/stable/current/amd64/hvm/ebs-gp3/ami-id \

  --query Parameter.Value --output text)"

log "ubuntu 24.04 ami: $AMI_ID"

MY_IP="$(curl -fsS https://checkip.amazonaws.com)"

# --- Provision -------------------------------------------------------------

aws_ec2 create-key-pair --key-name "$RUN_ID" --key-type ed25519 \

  --query KeyMaterial --output text > "$PEM"

chmod 600 "$PEM"

KEY_CREATED=1

VPC_ID="$(aws_ec2 describe-vpcs --filters Name=isDefault,Values=true \

  --query 'Vpcs[0].VpcId' --output text)"

[ "$VPC_ID" != "None" ] || { echo "no default VPC in $REGION" >&2; exit 2; }

SG_ID="$(aws_ec2 create-security-group --group-name "$RUN_ID" \

  --description "ephemeral oak benchmark runner" --vpc-id "$VPC_ID" \

  --query GroupId --output text)"

aws_ec2 authorize-security-group-ingress --group-id "$SG_ID" \

  --protocol tcp --port 22 --cidr "${MY_IP}/32" >/dev/null

INSTANCE_ID="$(aws_ec2 run-instances \

  --image-id "$AMI_ID" \

  --instance-type "$INSTANCE_TYPE" \

  --key-name "$RUN_ID" \

  --security-group-ids "$SG_ID" \

  --block-device-mappings "[{\"DeviceName\":\"/dev/sda1\",\"Ebs\":{\"VolumeSize\":$DISK_GB,\"VolumeType\":\"gp3\",\"DeleteOnTermination\":true}}]" \

  --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=$RUN_ID},{Key=purpose,Value=oak-benchmarks}]" \

  --query 'Instances[0].InstanceId' --output text)"

log "instance: $INSTANCE_ID (waiting for running state)"

aws_ec2 wait instance-running --instance-ids "$INSTANCE_ID"

PUBLIC_IP="$(aws_ec2 describe-instances --instance-ids "$INSTANCE_ID" \

  --query 'Reservations[0].Instances[0].PublicIpAddress' --output text)"

[ "$PUBLIC_IP" != "None" ] || { echo "instance has no public IP" >&2; exit 1; }

log "public ip: $PUBLIC_IP (waiting for SSH)"

for i in $(seq 1 30); do

  ssh "${SSH_OPTS[@]}" "$SSH_USER@$PUBLIC_IP" true 2>/dev/null && break

  [ "$i" = 30 ] && { echo "SSH never came up" >&2; exit 1; }

  sleep 10

done

# --- Upload + setup --------------------------------------------------------

log "syncing benchmarks repo to instance"

rsync -az -e "ssh ${SSH_OPTS[*]}" \

  --exclude results/ --exclude workdirs/ --exclude worktrees/ --exclude runs/ \

  --exclude fixtures/ --exclude __pycache__/ --exclude '.git/' --exclude '.oak/' \

  "$BENCH_ROOT/" "$SSH_USER@$PUBLIC_IP:benchmarks/"

log "installing dependencies (git, python3, oak)"

ssh "${SSH_OPTS[@]}" "$SSH_USER@$PUBLIC_IP" 'bash -s' <<'SETUP'

set -euo pipefail

sudo apt-get update -qq

sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \

  git python3 rsync curl ca-certificates >/dev/null

# Use instance-store NVMe (if this type has one) as the benchmark temp dir,

# so fixture generation hits local SSD rather than the EBS root volume.

ROOT_DISK="$(lsblk -no PKNAME "$(findmnt -no SOURCE /)" | head -1)"

for dev in /dev/nvme*n1; do

  [ -b "$dev" ] || continue

  [ "$(basename "$dev")" = "$ROOT_DISK" ] && continue

  if [ -z "$(lsblk -no MOUNTPOINTS "$dev" | tr -d '[:space:]')" ]; then

    sudo mkfs.ext4 -q -F "$dev"

    sudo mkdir -p /mnt/bench

    sudo mount "$dev" /mnt/bench

    sudo chown "$USER" /mnt/bench

    echo "[setup] instance-store NVMe mounted at /mnt/bench"

    break

  fi

done

curl -fsSL oak.space/install | sh

export PATH="$HOME/.oak/bin:$HOME/.local/bin:/usr/local/bin:$PATH"

echo "[setup] git:    $(git --version)"

echo "[setup] python: $(python3 --version)"

echo "[setup] oak:    $(command -v oak) ($(oak --version 2>&1 | head -1))"

SETUP

# --- Run benchmark (detached, so an SSH drop doesn't kill a long run) -------

EXTRA_ARGS_STR=""

[ ${#BENCH_EXTRA_ARGS[@]} -gt 0 ] && EXTRA_ARGS_STR="$(printf ' %q' "${BENCH_EXTRA_ARGS[@]}")"

log "starting bench.py --profile $PROFILE${EXTRA_ARGS_STR}"

ssh "${SSH_OPTS[@]}" "$SSH_USER@$PUBLIC_IP" 'bash -s' <<RUN

set -euo pipefail

cat > run_bench.sh <<'INNER'

#!/usr/bin/env bash

set -uo pipefail

export PATH="\$HOME/.oak/bin:\$HOME/.local/bin:/usr/local/bin:\$PATH"

if mountpoint -q /mnt/bench; then

  export TMPDIR=/mnt/bench/tmp

  mkdir -p "\$TMPDIR"

fi

cd "\$HOME/benchmarks"

python3 scripts/bench.py --profile $PROFILE${EXTRA_ARGS_STR}

echo \$? > "\$HOME/bench.exit"

INNER

chmod +x run_bench.sh

rm -f bench.exit bench.log

setsid nohup ./run_bench.sh > bench.log 2>&1 < /dev/null &

RUN

# Poll until the exit file appears, streaming new log output as it arrives.

DEADLINE=$(( $(date +%s) + TIMEOUT_MINS * 60 ))

LOG_OFFSET=0

BENCH_EXIT=""

while :; do

  sleep 20

  CHUNK="$(ssh "${SSH_OPTS[@]}" "$SSH_USER@$PUBLIC_IP" \

    "tail -c +$((LOG_OFFSET + 1)) bench.log 2>/dev/null; \

     [ -f bench.exit ] && printf '\n__BENCH_EXIT__%s' \"\$(cat bench.exit)\"" \

    2>/dev/null)" || { log "ssh poll failed; retrying"; continue; }

  if [[ "$CHUNK" == *__BENCH_EXIT__* ]]; then

    BENCH_EXIT="${CHUNK##*__BENCH_EXIT__}"

    CHUNK="${CHUNK%$'\n'__BENCH_EXIT__*}"

  fi

  if [ -n "$CHUNK" ]; then

    printf '%s\n' "$CHUNK"

    LOG_OFFSET=$((LOG_OFFSET + $(printf '%s' "$CHUNK" | wc -c) + 1))

  fi

  [ -n "$BENCH_EXIT" ] && break

  if [ "$(date +%s)" -ge "$DEADLINE" ]; then

    echo "benchmark exceeded ${TIMEOUT_MINS}m timeout" >&2

    exit 1

  fi

done

log "bench.py exited with code $BENCH_EXIT"

# --- Download results ------------------------------------------------------

DEST="$BENCH_ROOT/results/aws/$RUN_ID"

mkdir -p "$DEST"

rsync -az -e "ssh ${SSH_OPTS[*]}" \

  "$SSH_USER@$PUBLIC_IP:benchmarks/results/" "$DEST/" 2>/dev/null || true

ssh "${SSH_OPTS[@]}" "$SSH_USER@$PUBLIC_IP" "cat bench.log" > "$DEST/bench.log" || true

log "results downloaded to results/aws/$RUN_ID/"

if [ -f "$DEST/latest.summary.md" ]; then

  echo

  cat "$DEST/latest.summary.md"

fi

exit "$BENCH_EXIT"