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

# focuschain-ab-long extension: waits for the main experiment's DONE-control-2,
# then runs treatment-L and control-L (24 structurally-long tasks, 1 run/arm)
# on the same runner with identical config. Budget/deadline guarded.

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-long"
LOGDIR="$ROOT/logs/focuschain-ab-long"
DEADLINE_EPOCH=$(date -u -d '2026-07-04 09:10:00' +%s)  # hard stop 09:32Z
PACK_LOCAL="$ROOT/trials/shared/osworld-long-pack-candidates/long-pack-24.json"
PACK_REMOTE="/srv/controller-data/long-pack-24.json"

mkdir -p "$RESULTS" "$LOGDIR"
exec >>"$LOGDIR/orchestrator-long.log" 2>&1
log() { echo "[$(date -u +%FT%TZ)] $*"; }

log "waiting for DONE-control-2"
while [[ ! -f "$ROOT/results/focuschain-ab/DONE-control-2" ]]; do
  sleep 60
  if (( $(date -u +%s) > DEADLINE_EPOCH )); then log "deadline before main finished; abort"; exit 0; fi
  systemctl is-active fc2-finalize >/dev/null 2>&1 || { [[ -f "$ROOT/results/focuschain-ab/DONE-control-2" ]] || { log "finalizer died without DONE-control-2; abort"; exit 0; }; }
done
log "main experiment complete — starting extension"

rsync -a -e "ssh -i $KEY" "$PACK_LOCAL" "ubuntu@$RUNNER_IP:$PACK_REMOTE" </dev/null

for entry in treatment-L:focus_chain control-L:keep_all; do
  BLOCK="${entry%%:*}"; POLICY="${entry##*:}"
  NOW=$(date -u +%s)
  if (( NOW > DEADLINE_EPOCH - 3600 )); then log "insufficient time for $BLOCK; stopping"; break; fi

  # budget guard: sum upstream costs across both experiments so far
  SPEND=$(python3 - <<'PY'
import json, glob
tot=0.0
for traj in glob.glob('/srv/controller-data/results/focuschain-ab*/**/traj.jsonl', recursive=True):
    seen=set()
    for line in open(traj, errors='replace'):
        try: row=json.loads(line)
        except: continue
        u=(row.get('action') or {}).get('usage') or {}
        sn=row.get('step_num')
        if sn in seen or not u: continue
        seen.add(sn)
        tot+=float((u.get('cost_details') or {}).get('upstream_inference_cost') or 0)
        cm=u.get('context_mgmt_usage') or {}
        tot+=float((cm.get('cost_details') or {}).get('upstream_inference_cost') or 0)
print(f"{tot:.2f}")
PY
)
  log "spend so far: \$$SPEND (cap 120)"
  if python3 -c "exit(0 if float('$SPEND') > 100 else 1)"; then log "budget guard tripped (> \$100); stopping"; break; fi

  log "=== starting $BLOCK (policy=$POLICY) ==="
  $SSH "sudo systemd-run --unit=fcl-$BLOCK --collect --property=User=ubuntu \
    --property=WorkingDirectory=/srv/controller-data \
    --setenv=POLICY=$POLICY --setenv=BLOCK=$BLOCK --setenv=NUM_ENVS=12 \
    --setenv=EXPERIMENT=focuschain-ab-long --setenv=TEST_META_PATH=$PACK_REMOTE \
    /srv/controller-data/run_block.sh"
  [[ $? -ne 0 ]] && { log "FAILED to start fcl-$BLOCK"; break; }

  BLOCK_START=$(date -u +%s)
  while true; do
    sleep 120
    STATE=$($SSH "systemctl is-active fcl-$BLOCK" 2>/dev/null </dev/null)
    NOW=$(date -u +%s)
    [[ "$STATE" != "active" ]] && { log "$BLOCK unit state: $STATE"; break; }
    if (( NOW - BLOCK_START > 5400 )) || (( NOW > DEADLINE_EPOCH )); then
      log "$BLOCK time cap/deadline — stopping unit"
      $SSH "sudo systemctl stop fcl-$BLOCK" </dev/null
      break
    fi
  done

  rsync -a -e "ssh -i $KEY" "ubuntu@$RUNNER_IP:$ROOT/results/focuschain-ab-long/$BLOCK/" "$RESULTS/$BLOCK/" </dev/null
  rsync -a -e "ssh -i $KEY" "ubuntu@$RUNNER_IP:$ROOT/logs/focuschain-ab-long/$BLOCK.log" "$LOGDIR/" </dev/null
  N=$(find "$RESULTS/$BLOCK" -name result.txt 2>/dev/null | wc -l)
  echo "block=$BLOCK policy=$POLICY finished_utc=$(date -u +%FT%TZ) results=$N/24" > "$RESULTS/DONE-$BLOCK"
  log "wrote DONE-$BLOCK ($N/24 results)"
done
log "extension orchestrator finished"
