#!/bin/bash
# Prepare the Trigger-Off poisoned dataset from Mind2Web.
#
# Steps:
#   1. Download / convert Mind2Web from HuggingFace.
#   2. Filter and balance trajectories by site.
#   3. Build the poisoned dataset (immunity + attack + clean samples).
#
# Environment variables (all have defaults):
#   DATA_DIR   - root data directory (default: ./data)
#   CONFIG     - path to training config (default: configs/training_config.yaml)
#   SPLIT      - Mind2Web HF split to download (default: train)
#   MAX_PER_SITE - max trajectories per site after balancing (default: 500)

set -e
cd "$(dirname "$0")/.."

DATA_DIR=${DATA_DIR:-./data}
CONFIG=${CONFIG:-configs/training_config.yaml}
SPLIT=${SPLIT:-train}
MAX_PER_SITE=${MAX_PER_SITE:-500}

echo "============================================================"
echo "  Trigger-Off: Data Preparation"
echo "============================================================"
echo "  DATA_DIR     = $DATA_DIR"
echo "  CONFIG       = $CONFIG"
echo "  SPLIT        = $SPLIT"
echo "  MAX_PER_SITE = $MAX_PER_SITE"
echo "------------------------------------------------------------"

# ---- Step 1: Download / convert Multimodal-Mind2Web ----
echo ""
echo "[1/3] Downloading osunlp/Multimodal-Mind2Web (split=$SPLIT) …"
python -m src.data_pipeline.mind2web_loader \
    --output_dir "$DATA_DIR/mind2web_raw" \
    --split "$SPLIT" \
    --save_images

# ---- Step 2: Filter and balance by site ----
echo ""
echo "[2/3] Filtering and balancing by site …"
python -m src.data_pipeline.site_filter \
    --input_dir  "$DATA_DIR/mind2web_raw" \
    --output_dir "$DATA_DIR/processed" \
    --max_per_site "$MAX_PER_SITE"

# ---- Step 3: Build poisoned dataset ----
echo ""
echo "[3/3] Building Trigger-Off poisoned dataset …"
python -m src.data_pipeline.dataset_builder \
    --input_dir  "$DATA_DIR/processed" \
    --output_dir "$DATA_DIR/poisoned" \
    --config     "$CONFIG"

echo ""
echo "============================================================"
echo "  Data preparation complete."
echo "  Train JSON : $DATA_DIR/poisoned/train.json"
echo "  Test  JSON : $DATA_DIR/poisoned/test.json"
echo "============================================================"