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

# focuschain-ab orchestrator (runs on the controller, long-lived).
# Serially runs 4 blocks on the runner, syncs results back after each block,
# and drops DONE-<block> markers for the downstream consumer agent.

RUNNER_IP="172.31.44.229"
KEY="/home/ubuntu/.ssh/controller-worker-aws-20260702"
SSH="ssh -i $KEY -o ConnectTimeout=15 -o ServerAliveInterval=30 ubuntu@$RUNNER_IP"
ROOT="/srv/controller-data"
RESULTS="$ROOT/results/focuschain-ab"
LOGDIR="$ROOT/logs/focuschain-ab"
DEADLINE_EPOCH=$(date -u -d '2026-07-04 08:45:00' +%s)  # leave margin before 09:32Z hard stop

mkdir -p "$RESULTS" "$LOGDIR"
exec >>"$LOGDIR/orchestrator.log" 2>&1

log() { echo "[$(date -u +%FT%TZ)] $*"; }

BLOCKS=(control-1:keep_all treatment-1:focus_chain treatment-2:focus_chain control-2:keep_all)

snapshot_fleet() {
  aws ec2 describe-instances --region us-west-2 \
    --query 'Reservations[].Instances[?State.Name==`running`].[InstanceId,InstanceType,Tags[?Key==`Name`].Value|[0]]' \
    --output text | sort
}

for entry in "${BLOCKS[@]}"; do
  BLOCK="${entry%%:*}"; POLICY="${entry##*:}"
  if [[ -f "$RESULTS/DONE-$BLOCK" ]]; then log "$BLOCK already done, skipping"; continue; fi
  NOW=$(date -u +%s)
  if (( NOW > DEADLINE_EPOCH )); then log "DEADLINE reached before $BLOCK — stopping"; break; fi

  log "=== starting block $BLOCK (policy=$POLICY) ==="
  log "fleet before block:"; snapshot_fleet

  $SSH "sudo systemd-run --unit=fc-$BLOCK --collect --property=User=ubuntu \
    --property=WorkingDirectory=/srv/controller-data \
    --setenv=POLICY=$POLICY --setenv=BLOCK=$BLOCK --setenv=NUM_ENVS=12 \
    /srv/controller-data/run_block.sh"
  if [[ $? -ne 0 ]]; then log "FAILED to start unit fc-$BLOCK — aborting"; break; fi

  # Wait for the block to finish (poll every 2 min); guard against runaway blocks.
  BLOCK_START=$(date -u +%s)
  while true; do
    sleep 120
    STATE=$($SSH "systemctl is-active fc-$BLOCK" 2>/dev/null </dev/null)
    NOW=$(date -u +%s)
    if [[ "$STATE" != "active" ]]; then log "block $BLOCK unit state: $STATE"; break; fi
    if (( NOW - BLOCK_START > 4*3600 )); then
      log "block $BLOCK exceeded 4h — stopping unit"
      $SSH "sudo systemctl stop fc-$BLOCK" </dev/null
      break
    fi
    if (( NOW > DEADLINE_EPOCH )); then
      log "DEADLINE hit mid-block $BLOCK — stopping unit"
      $SSH "sudo systemctl stop fc-$BLOCK" </dev/null
      break
    fi
  done

  # Sync results + log back to the controller.
  rsync -a -e "ssh -i $KEY" "ubuntu@$RUNNER_IP:$ROOT/results/focuschain-ab/$BLOCK/" "$RESULTS/$BLOCK/" </dev/null
  rsync -a -e "ssh -i $KEY" "ubuntu@$RUNNER_IP:$ROOT/logs/focuschain-ab/$BLOCK.log" "$LOGDIR/" </dev/null
  N_RESULTS=$(find "$RESULTS/$BLOCK" -name result.txt 2>/dev/null | wc -l)
  log "block $BLOCK synced: $N_RESULTS/40 result.txt files"

  {
    echo "block=$BLOCK policy=$POLICY finished_utc=$(date -u +%FT%TZ) results=$N_RESULTS/40"
  } > "$RESULTS/DONE-$BLOCK"
  log "wrote DONE-$BLOCK"

  # Give harness-initiated worker teardown a moment, then snapshot fleet.
  sleep 60
  log "fleet after block:"; snapshot_fleet
done

log "orchestrator finished"
