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

PROFILE="${1:-today}"
WORK_ROOT="${NAYARA_WORK_ROOT:-/teamspace/studios/this_studio/nayara-work}"

case "$PROFILE" in
  preview)
    TOTAL_TOKENS=50000000
    MAX_DOCS=40000
    ;;
  today)
    TOTAL_TOKENS=300000000
    MAX_DOCS=150000
    ;;
  overnight)
    # A realistic single-H100 target for an evening-to-04:00 run. The packer
    # stops at each source's weighted token target, so disk use stays bounded.
    TOTAL_TOKENS="${NAYARA_OVERNIGHT_TOKENS:-2000000000}"
    MAX_DOCS="${NAYARA_OVERNIGHT_MAX_DOCS:-1200000}"
    ;;
  *)
    echo "usage: bash scripts/lightning_prepare.sh [preview|today|overnight]" >&2
    exit 2
    ;;
esac

mkdir -p "$WORK_ROOT"
python -m nayara.prepare_data profile \
  --sources data/sources.example.json \
  --output "$WORK_ROOT/sources-${PROFILE}.json" \
  --total-tokens "$TOTAL_TOKENS"

python -m nayara.prepare_data fetch \
  --sources "$WORK_ROOT/sources-${PROFILE}.json" \
  --raw-dir "$WORK_ROOT/raw-${PROFILE}" \
  --max-docs "$MAX_DOCS"

python -m nayara.prepare_data tokenizer \
  --raw-dir "$WORK_ROOT/raw-${PROFILE}" \
  --output "$WORK_ROOT/tokenizer-${PROFILE}.json"

python -m nayara.prepare_data pack \
  --raw-dir "$WORK_ROOT/raw-${PROFILE}" \
  --tokenizer "$WORK_ROOT/tokenizer-${PROFILE}.json" \
  --output "$WORK_ROOT/packed-${PROFILE}" \
  --sources "$WORK_ROOT/sources-${PROFILE}.json"

python - <<PY
import json
from pathlib import Path
manifest = json.loads(Path("$WORK_ROOT/packed-$PROFILE/manifest.json").read_text())
print("packed tokens:", {name: item["tokens"] for name, item in manifest["sources"].items()})
print("total:", sum(item["tokens"] for item in manifest["sources"].values()))
PY