#!/usr/bin/env bash
# Self-destructing, budget-guarded AWS validation run for SimurgArray.
#
#   bash scripts/cloud/trial.sh setup  <alert-email>   # one-time: budget alert, results bucket, instance role
#   bash scripts/cloud/trial.sh launch                  # start one spot run (terminates itself when done)
#   bash scripts/cloud/trial.sh status                  # instance state + last lines of the remote log
#   bash scripts/cloud/trial.sh results                 # download log/summary of the latest run
#   bash scripts/cloud/trial.sh kill                    # emergency: terminate every trial instance now
#   bash scripts/cloud/trial.sh audit                   # list anything still costing money
#
# Cost guards: spot price cap, hard time cap (shutdown = terminate), EXIT trap shutdown, CloudWatch idle alarm
# that terminates from outside, disk deleted with the instance, monthly budget e-mail alert.
set -euo pipefail
cd "$(dirname "$0")/../.."
REGION=${REGION:-us-east-1}
TYPE=${TYPE:-m7g.4xlarge}
MAX_PRICE=${MAX_PRICE:-0.30}        # USD/hour spot cap (on-demand is 0.6528)
MAX_MINUTES=${MAX_MINUTES:-360}     # hard cap: 6 hours
DISK_GB=${DISK_GB:-250}
DISK_MBPS=${DISK_MBPS:-125}         # gp3 throughput (125 = free baseline; each extra MB/s ~ $0.00005/h)
DISK_IOPS=${DISK_IOPS:-3000}        # gp3 IOPS (3000 = free baseline)
BUDGET_USD=${BUDGET_USD:-10}
TASK=${TASK:-p3}                 # scripts/cloud/tasks/<TASK>.sh
TAG=simurg-trial
ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
BUCKET=simurg-trial-$ACCOUNT
ROLE=simurg-trial-role
AWS="aws --region $REGION"

cmd=${1:-help}
case "$cmd" in
setup)
  EMAIL=${2:?usage: trial.sh setup <alert-email>}
  echo "-> budget alert: \$$BUDGET_USD/month, e-mail $EMAIL at 50%, 100% actual and 100% forecast"
  cat > /tmp/simurg-budget.json <<EOF
{"BudgetName":"$TAG-monthly","BudgetLimit":{"Amount":"$BUDGET_USD","Unit":"USD"},"TimeUnit":"MONTHLY","BudgetType":"COST"}
EOF
  notif() { echo "{\"Notification\":{\"NotificationType\":\"$1\",\"ComparisonOperator\":\"GREATER_THAN\",\"Threshold\":$2,\"ThresholdType\":\"PERCENTAGE\"},\"Subscribers\":[{\"SubscriptionType\":\"EMAIL\",\"Address\":\"$EMAIL\"}]}"; }
  echo "[$(notif ACTUAL 50),$(notif ACTUAL 100),$(notif FORECASTED 100)]" > /tmp/simurg-notif.json
  aws budgets create-budget --account-id "$ACCOUNT" --budget file:///tmp/simurg-budget.json \
    --notifications-with-subscribers file:///tmp/simurg-notif.json 2>/dev/null || echo "   (budget already exists)"
  echo "-> results bucket s3://$BUCKET (objects expire after 30 days)"
  $AWS s3api create-bucket --bucket "$BUCKET" >/dev/null 2>&1 || true
  $AWS s3api put-public-access-block --bucket "$BUCKET" --public-access-block-configuration \
    BlockPublicAcls=true,IgnorePublicAcls=true,BlockPublicPolicy=true,RestrictPublicBuckets=true
  $AWS s3api put-bucket-lifecycle-configuration --bucket "$BUCKET" --lifecycle-configuration \
    '{"Rules":[{"ID":"expire","Status":"Enabled","Filter":{},"Expiration":{"Days":30}}]}'
  echo "-> instance role $ROLE (read/write only this bucket)"
  aws iam create-role --role-name $ROLE --assume-role-policy-document \
    '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"ec2.amazonaws.com"},"Action":"sts:AssumeRole"}]}' \
    >/dev/null 2>&1 || true
  aws iam put-role-policy --role-name $ROLE --policy-name bucket-only --policy-document \
    "{\"Version\":\"2012-10-17\",\"Statement\":[{\"Effect\":\"Allow\",\"Action\":[\"s3:GetObject\",\"s3:PutObject\"],\"Resource\":\"arn:aws:s3:::$BUCKET/*\"}]}"
  aws iam create-instance-profile --instance-profile-name $ROLE >/dev/null 2>&1 || true
  aws iam add-role-to-instance-profile --instance-profile-name $ROLE --role-name $ROLE 2>/dev/null || true
  echo "setup done (budget alerts need you to confirm nothing; AWS e-mails $EMAIL when thresholds are crossed)"
  ;;
launch)
  running=$($AWS ec2 describe-instances --filters Name=tag:Project,Values=$TAG \
    Name=instance-state-name,Values=pending,running --query 'length(Reservations[].Instances[])')
  [[ "$running" == "0" ]] || { echo "a trial instance is already running; refusing to start another"; exit 1; }
  RUN=run-$(date -u +%Y%m%dT%H%M%SZ)
  echo "-> packing repository at $(git rev-parse --short HEAD) -> s3://$BUCKET/$RUN/"
  git archive --format=tar.gz -o /tmp/simurg-repo.tar.gz HEAD
  $AWS s3 cp /tmp/simurg-repo.tar.gz "s3://$BUCKET/$RUN/repo.tar.gz" --only-show-errors
  [[ -f scripts/cloud/tasks/$TASK.sh ]] || { echo "unknown TASK $TASK"; exit 2; }
  OMNI=${OMNI:-1}  # 0: the task does not need the 107 GB Omni2.5 archive
  sed -e "s/__BUCKET__/$BUCKET/" -e "s/__RUN__/$RUN/" -e "s/__MAX_MINUTES__/$MAX_MINUTES/" -e "s/__TASK__/$TASK/g" \
    -e "s/__OMNI__/$OMNI/" \
    scripts/cloud/userdata.sh > /tmp/simurg-userdata.sh
  AMI=$($AWS ssm get-parameter --name /aws/service/canonical/ubuntu/server/24.04/stable/current/arm64/hvm/ebs-gp3/ami-id \
    --query Parameter.Value --output text)
  echo "-> task: $TASK"
  echo "-> launching $TYPE spot (cap \$$MAX_PRICE/h, hard limit ${MAX_MINUTES} min, ${DISK_GB} GB disk deleted on exit)"
  ID=$($AWS ec2 run-instances --image-id "$AMI" --instance-type "$TYPE" --count 1 \
    --iam-instance-profile Name=$ROLE \
    --instance-initiated-shutdown-behavior terminate \
    --instance-market-options "MarketType=spot,SpotOptions={MaxPrice=$MAX_PRICE,SpotInstanceType=one-time,InstanceInterruptionBehavior=terminate}" \
    --block-device-mappings "DeviceName=/dev/sda1,Ebs={VolumeSize=$DISK_GB,VolumeType=gp3,Throughput=$DISK_MBPS,Iops=$DISK_IOPS,DeleteOnTermination=true}" \
    --metadata-options HttpTokens=required \
    --user-data file:///tmp/simurg-userdata.sh \
    --tag-specifications "ResourceType=instance,Tags=[{Key=Project,Value=$TAG},{Key=Run,Value=$RUN}]" \
      "ResourceType=volume,Tags=[{Key=Project,Value=$TAG},{Key=Run,Value=$RUN}]" \
    --query 'Instances[0].InstanceId' --output text)
  echo "$ID $RUN" > .trial-last
  echo "-> idle alarm: terminate if average CPU < 2% for 30 minutes"
  $AWS cloudwatch put-metric-alarm --alarm-name "$TAG-idle-$ID" --namespace AWS/EC2 --metric-name CPUUtilization \
    --dimensions Name=InstanceId,Value="$ID" --statistic Average --period 300 --evaluation-periods 6 \
    --threshold 2 --comparison-operator LessThanThreshold --treat-missing-data notBreaching \
    --alarm-actions "arn:aws:automate:$REGION:ec2:terminate"
  echo "launched $ID ($RUN). Check with: bash scripts/cloud/trial.sh status"
  ;;
status)
  read -r ID RUN < .trial-last
  $AWS ec2 describe-instances --instance-ids "$ID" \
    --query 'Reservations[0].Instances[0].[InstanceId,State.Name,InstanceType,LaunchTime]' --output text
  $AWS s3 cp "s3://$BUCKET/$RUN/log.txt" - 2>/dev/null | tail -n "${LINES_:-15}" || echo "(no log uploaded yet)"
  ;;
results)
  read -r ID RUN < .trial-last
  mkdir -p "results/$RUN"
  $AWS s3 cp "s3://$BUCKET/$RUN/" "results/$RUN/" --recursive --exclude repo.tar.gz --only-show-errors
  ls -la "results/$RUN"
  ;;
kill)
  ids=$($AWS ec2 describe-instances --filters Name=tag:Project,Values=$TAG \
    Name=instance-state-name,Values=pending,running,stopping,stopped --query 'Reservations[].Instances[].InstanceId' --output text)
  [[ -n "$ids" ]] && $AWS ec2 terminate-instances --instance-ids $ids --output text || echo "no trial instances"
  ;&
audit)
  echo "-- instances (any project) --"
  $AWS ec2 describe-instances --filters Name=instance-state-name,Values=pending,running,stopping,stopped \
    --query 'Reservations[].Instances[].[InstanceId,State.Name,InstanceType]' --output text
  echo "-- volumes --"
  $AWS ec2 describe-volumes --query 'Volumes[].[VolumeId,State,Size]' --output text
  echo "-- snapshots owned --"
  $AWS ec2 describe-snapshots --owner-ids self --query 'Snapshots[].[SnapshotId,VolumeSize]' --output text
  echo "-- elastic IPs --"
  $AWS ec2 describe-addresses --query 'Addresses[].PublicIp' --output text
  echo "-- trial alarms (deleting alarms of terminated instances) --"
  for a in $($AWS cloudwatch describe-alarms --alarm-name-prefix "$TAG-idle-" --query 'MetricAlarms[].AlarmName' --output text); do
    iid=${a#"$TAG-idle-"}
    st=$($AWS ec2 describe-instances --instance-ids "$iid" --query 'Reservations[0].Instances[0].State.Name' --output text 2>/dev/null || echo gone)
    if [[ "$st" == "terminated" || "$st" == "gone" ]]; then $AWS cloudwatch delete-alarms --alarm-names "$a"; echo "deleted $a"; else echo "$a ($st)"; fi
  done
  ;;
*)
  sed -n '2,15p' "$0"
  ;;
esac