#!/usr/bin/env bash
# Build preprocessing data used by CHNB experiments.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
. "$PROJECT_ROOT/scripts/chnb_env.sh"

RUN_LOCAL=1
RUN_GEMINI=0
DRY_RUN=0
SERVER_STEP="all"
GEMINI_STEP="all"
UNIFORM_N=(64 128 256)
CHUNK_MODES=(shot_grouped)
FPS="0.5"
ASSUME_FPS="24.0"
SERVER_EXTRA=()
GEMINI_EXTRA=()

usage() {
    cat <<'EOF'
Usage:
  bash scripts/preprocess_all.sh [options]

Default:
  Runs local preprocessing only: frames, CLIP embeddings, and clips.

Options:
  --all                     Run local preprocessing and Gemini captioning.
  --with-gemini             Also run Gemini/Vertex captioning.
  --gemini-only             Run only Gemini/Vertex preprocessing.
  --local-only              Run only local preprocessing.
  --server-step STEP        Step for run_server.py (default: all).
  --gemini-step STEP        Step for run_gemini.py (default: all).
  --uniform-n N...          Uniform frame budgets (default: 64 128 256).
  --chunk-modes MODE...     Clip chunk modes (default: shot_grouped).
  --fps FPS                 FPS frame extraction rate (default: 0.5).
  --assume-fps FPS          FPS used for shot-boundary frame indices (default: 24.0).
  --all-videos              Pass --all-videos to local preprocessing.
  --bucket NAME             Pass --bucket to Gemini preprocessing.
  --project ID              Pass --project to Gemini preprocessing.
  --vids CSV                Pass --vids to Gemini preprocessing.
  --dry-run                 Print commands without running them.
  -h, --help                Show this help.

Environment:
  CHNB_DATA_ROOT, CHNB_VIDEO_DIR, CHNB_FRAMES_DIR, CHNB_CLIP_EMBEDDINGS_DIR,
  CHNB_CLIPS_DIR, CHNB_CAPTIONS_DIR, CHNB_SHOTS_DIR, CHNB_GCS_BUCKET, and
  GOOGLE_CLOUD_PROJECT are honored by the underlying runners.
EOF
}

while [ "$#" -gt 0 ]; do
    case "$1" in
        --all)
            RUN_LOCAL=1
            RUN_GEMINI=1
            shift
            ;;
        --with-gemini)
            RUN_GEMINI=1
            shift
            ;;
        --gemini-only)
            RUN_LOCAL=0
            RUN_GEMINI=1
            shift
            ;;
        --local-only)
            RUN_LOCAL=1
            RUN_GEMINI=0
            shift
            ;;
        --server-step)
            SERVER_STEP="$2"
            shift 2
            ;;
        --gemini-step)
            GEMINI_STEP="$2"
            shift 2
            ;;
        --uniform-n)
            shift
            UNIFORM_N=()
            while [ "$#" -gt 0 ] && [[ "$1" != --* ]]; do
                UNIFORM_N+=("$1")
                shift
            done
            if [ "${#UNIFORM_N[@]}" -eq 0 ]; then
                echo "--uniform-n requires at least one value" >&2
                exit 2
            fi
            ;;
        --chunk-modes)
            shift
            CHUNK_MODES=()
            while [ "$#" -gt 0 ] && [[ "$1" != --* ]]; do
                CHUNK_MODES+=("$1")
                shift
            done
            if [ "${#CHUNK_MODES[@]}" -eq 0 ]; then
                echo "--chunk-modes requires at least one value" >&2
                exit 2
            fi
            ;;
        --fps)
            FPS="$2"
            shift 2
            ;;
        --assume-fps)
            ASSUME_FPS="$2"
            shift 2
            ;;
        --all-videos)
            SERVER_EXTRA+=("--all-videos")
            shift
            ;;
        --bucket)
            GEMINI_EXTRA+=("--bucket" "$2")
            shift 2
            ;;
        --project)
            GEMINI_EXTRA+=("--project" "$2")
            shift 2
            ;;
        --vids)
            GEMINI_EXTRA+=("--vids" "$2")
            shift 2
            ;;
        --dry-run)
            DRY_RUN=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

run_cmd() {
    local label="$1"
    shift
    echo
    echo "=== $label ==="
    printf '$'
    printf ' %q' "$@"
    echo
    if [ "$DRY_RUN" = "1" ]; then
        return 0
    fi
    "$@"
}

if [ "$RUN_LOCAL" = "1" ]; then
    cmd=(
        bash "$PROJECT_ROOT/experiments/preprocess/run_server.sh"
        --step "$SERVER_STEP"
        --uniform-n "${UNIFORM_N[@]}"
        --fps "$FPS"
        --chunk-modes "${CHUNK_MODES[@]}"
        --assume-fps "$ASSUME_FPS"
    )
    if [ "${#SERVER_EXTRA[@]}" -gt 0 ]; then
        cmd+=("${SERVER_EXTRA[@]}")
    fi
    run_cmd "local preprocessing" "${cmd[@]}"
fi

if [ "$RUN_GEMINI" = "1" ]; then
    cmd=(
        bash "$PROJECT_ROOT/experiments/preprocess/run_gemini.sh"
        --step "$GEMINI_STEP"
        --uniform-n "${UNIFORM_N[@]}"
        --fps "$FPS"
        --chunk-modes "${CHUNK_MODES[@]}"
        --assume-fps "$ASSUME_FPS"
    )
    if [ "${#GEMINI_EXTRA[@]}" -gt 0 ]; then
        cmd+=("${GEMINI_EXTRA[@]}")
    fi
    run_cmd "Gemini preprocessing" "${cmd[@]}"
fi

echo
echo "Preprocessing orchestration complete."