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

# focuschain-ab: run one A/B block on the runner instance.
# Usage: POLICY=keep_all|focus_chain BLOCK=control-1 [NUM_ENVS=12] ./run_block.sh
# Both arms share every flag except --context_policy.

ROOT="/srv/controller-data"
OSWORLD_DIR="$ROOT/repos/OSWorld-focuschain"
VENV="$ROOT/cache/osworld-venv"

POLICY="${POLICY:?keep_all (control) or focus_chain (treatment)}"
BLOCK="${BLOCK:?e.g. control-1}"
NUM_ENVS="${NUM_ENVS:-12}"
MAX_STEPS="${MAX_STEPS:-50}"
MODEL="${MODEL:-us.anthropic.claude-sonnet-4-6}"
TEST_META_PATH="${TEST_META_PATH:-$ROOT/trials/shared/osworld-40-pack/pack.json}"
EXPERIMENT="${EXPERIMENT:-focuschain-ab}"
RESULT_DIR="$ROOT/results/$EXPERIMENT/$BLOCK"
LOG_DIR="$ROOT/logs/$EXPERIMENT"

mkdir -p "$LOG_DIR" "$RESULT_DIR"

set -a
source "$ROOT/.env"
set +a
# ANTHROPIC_API_KEY must never reach the OpenRouter client (401 footgun).
unset ANTHROPIC_API_KEY

export AWS_REGION="${AWS_REGION:-us-west-2}"
export AWS_DEFAULT_REGION="$AWS_REGION"
export AWS_SUBNET_ID="${AWS_SUBNET_ID:-subnet-036cf5036ea57dc61}"
export AWS_SECURITY_GROUP_ID="${AWS_SECURITY_GROUP_ID:-sg-092b15357dee7f1b5}"
export TOKENIZERS_PARALLELISM=false

cd "$OSWORLD_DIR"

exec "$VENV/bin/python" scripts/python/run_multienv_claude.py \
  --provider_name aws \
  --provider openrouter \
  --region "$AWS_REGION" \
  --client_password osworld-public-evaluation \
  --action_space claude_computer_use \
  --observation_type screenshot \
  --screen_width 1920 \
  --screen_height 1080 \
  --image_target_width 960 \
  --image_target_height 540 \
  --model "$MODEL" \
  --no_thinking \
  --max_steps "$MAX_STEPS" \
  --max_tokens 8000 \
  --num_envs "$NUM_ENVS" \
  --context_policy "$POLICY" \
  --test_all_meta_path "$TEST_META_PATH" \
  --result_dir "$RESULT_DIR" \
  --log_level INFO \
  2>&1 | tee -a "$LOG_DIR/$BLOCK.log"
