#!/usr/bin/env bash
# The annotator dies on a 503 UNAVAILABLE: keypool.is_rate_limit only matches 429/quota-style
# errors, so a transient "model experiencing high demand" escapes the handler and kills the
# process mid-range. It resumes from its own results JSON (episodes already present are skipped),
# so re-entering with backoff is safe and cheap. Keep ONE model id — the results filename embeds
# it, and switching mid-run starts a fresh file.
set -uo pipefail
cd /workspace/fd-studio/eval/ci_mse
export GEMINI_API_KEYS="AQ.Ab8RN6JI3VRv4hhZjDTbdpzxICp5A_lv-v7sEi0mWJ-XcYXZPQ,AQ.Ab8RN6J4-1le25DDURnrcp4ySEB-FOTRsgYsZhFNlw62NOOAGA,AQ.Ab8RN6LmjvzzwUZ72Eq8roPdeHbbtKvysUTEf-XrL4eXx_Kawg,AQ.Ab8RN6K1sICN6XDz3WepwiHi2ssloe9gWfOxgmMY0GL_GmHj_A"
export HF_HOME=/workspace/.hf
MODEL="${MODEL:-gemini-flash-lite-latest}"
OUT=/workspace/evalout/ci_annot
RESULTS="$OUT/test_zero_shot_${MODEL}.json"
PY=/workspace/smolvla_ft/.venv/bin/python

for attempt in $(seq 1 40); do
  n=$($PY -c "
import json,sys
try: print(len(json.load(open('$RESULTS'))))
except Exception: print(0)" 2>/dev/null)
  echo "=== attempt $attempt: $n/6 episodes annotated"
  [[ "$n" -ge 6 ]] && { echo "ANNOTATION COMPLETE"; break; }
  $PY vlm_annotator/gemini_annotator.py \
      --demo_path /workspace/mm_split/val --demo_type lerobot \
      --prompt_config vlm_annotator/prompts/PlaceBlueCubeInOrangeBox.json \
      --episode_idx "[0, 5]" --model "$MODEL" --backend gemini --downsample 3 \
      --output_dir "$OUT" --video_out "$OUT/episode.mp4" 2>&1 | tail -4
  sleep $(( attempt < 5 ? 15 : 45 ))
done