#!/usr/bin/env bash
# One-click environment setup for this repo.
#
#   ./setup.sh                 # install deps + download all required weights
#   ./setup.sh --skip-weights  # only set up the Python environment
#   ./setup.sh --weights-only  # only download weights (env already set up)
#
# Requires: Python 3.10+, ~90GB free disk for weights, and (for gated HF repos)
# an authenticated `huggingface-cli login` — set HF_TOKEN in your shell first,
# or run `huggingface-cli login` before this script if downloads 401.
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$REPO_ROOT"

DO_ENV=1
DO_WEIGHTS=1
for arg in "$@"; do
  case "$arg" in
    --skip-weights) DO_WEIGHTS=0 ;;
    --weights-only) DO_ENV=0 ;;
    *) echo "Unknown argument: $arg" >&2; exit 1 ;;
  esac
done

if [[ "$DO_ENV" -eq 1 ]]; then
  echo "=== [1/2] Python environment (uv sync) ==="
  if ! command -v uv >/dev/null 2>&1; then
    echo "uv not found -> installing via the official installer"
    curl -LsSf https://astral.sh/uv/install.sh | sh
    export PATH="$HOME/.local/bin:$PATH"
  fi
  uv sync
  echo "Environment ready. Activate it with: source .venv/bin/activate"
fi

if [[ "$DO_WEIGHTS" -eq 1 ]]; then
  echo "=== [2/2] Downloading model weights (~90GB total) ==="
  if ! command -v huggingface-cli >/dev/null 2>&1; then
    echo "huggingface-cli not found -> installing huggingface_hub CLI"
    uv run pip install -U "huggingface_hub[cli]" >/dev/null || pip install -U "huggingface_hub[cli]"
  fi
  HF="huggingface-cli"
  command -v uv >/dev/null 2>&1 && HF="uv run huggingface-cli"

  mkdir -p weights/ltx-2.3 weights/gemma-3-12b-it-qat-q4_0-unquantized

  echo "--- LTX-2.3 dev checkpoint (base model for training + custom-LoRA inference) ---"
  $HF download Lightricks/LTX-2.3 ltx-2.3-22b-dev.safetensors \
    --local-dir weights/ltx-2.3

  echo "--- LTX-2.3 spatial upscaler (2x stage-2 upsampling) ---"
  $HF download Lightricks/LTX-2.3 ltx-2.3-spatial-upscaler-x2-1.1.safetensors \
    --local-dir weights/ltx-2.3

  echo "--- Gemma 3 12B text encoder ---"
  $HF download google/gemma-3-12b-it-qat-q4_0-unquantized \
    --local-dir weights/gemma-3-12b-it-qat-q4_0-unquantized

  echo "Weights downloaded to: $REPO_ROOT/weights/"
  echo "NOTE: this repo does NOT ship pretrained IC-LoRA checkpoints (they were"
  echo "trained on licensed footage). Train your own with scripts/train_ic_lora.sh"
  echo "before running scripts/infer_v2v.py --structure-lora <your checkpoint>."
fi

echo "=== Done ==="