#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────
# scripts/setup_wav2lip.sh
#
# One-shot installer for the custom video role-play module. Run on
# the Ubuntu server once after deploy, from the repo root:
#
#   bash scripts/setup_wav2lip.sh
#
# What this script does:
#   1. Installs system packages: ffmpeg + libgl/libsm (cv2 needs them).
#   2. Clones the Wav2Lip repository into vendor/Wav2Lip/.
#   3. Installs Wav2Lip's Python dependencies into the active venv.
#   4. Downloads the pretrained `wav2lip_gan.pth` model into
#      vendor/Wav2Lip/checkpoints/.
#   5. Verifies your 4 avatar MP4s are present under
#      static/assets/avatars/.
#   6. Pre-creates the cache directory used by services/video_engine.
#
# After this script finishes, the FastAPI app picks Wav2Lip up
# automatically — no code change needed. Sanity-check from your
# browser at:
#   http://<host>:8000/api/_health/video
# ─────────────────────────────────────────────────────────────────
set -euo pipefail

# Resolve the repo root (parent of this script's directory).
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
ROOT_DIR="$( cd "$SCRIPT_DIR/.." && pwd )"
cd "$ROOT_DIR"

echo "▶ repo root: $ROOT_DIR"

VENDOR_DIR="$ROOT_DIR/vendor"
WAV2LIP_DIR="$VENDOR_DIR/Wav2Lip"
CHECKPOINT_DIR="$WAV2LIP_DIR/checkpoints"
CHECKPOINT_PATH="$CHECKPOINT_DIR/wav2lip_gan.pth"
AVATAR_DIR="$ROOT_DIR/static/assets/avatars"
CACHE_DIR="$ROOT_DIR/data/cache/video"
TMP_DIR="$ROOT_DIR/data/cache/tmp"

mkdir -p "$VENDOR_DIR" "$AVATAR_DIR" "$CACHE_DIR" "$TMP_DIR"

# ── 1. System packages ──────────────────────────────────────────
echo "▶ Installing system packages (ffmpeg, libGL, libSM, libXext)…"
if command -v apt-get >/dev/null 2>&1; then
  sudo apt-get update -qq
  sudo apt-get install -y -qq ffmpeg libgl1 libsm6 libxext6 git
else
  echo "  ⚠ apt-get not found — make sure ffmpeg + libgl1 are installed manually."
fi

# ── 2. Clone Wav2Lip ────────────────────────────────────────────
if [ -d "$WAV2LIP_DIR/.git" ]; then
  echo "▶ Wav2Lip already cloned at $WAV2LIP_DIR — skipping."
else
  echo "▶ Cloning Wav2Lip…"
  git clone --depth 1 https://github.com/Rudrabha/Wav2Lip.git "$WAV2LIP_DIR"
fi

# ── 3. Python deps ──────────────────────────────────────────────
# Wav2Lip's requirements.txt pins ancient versions; we install them
# inside whatever venv is currently active so the CPython resolver
# can pick compatible binary wheels for our Python version.
echo "▶ Installing Python dependencies…"
PIP_EXTRA=()
# CPU-only torch wheels from the official PyPI index. If you later
# add a GPU, remove this and install the CUDA-matched torch wheel.
PIP_EXTRA=("torch==2.2.2" "torchvision==0.17.2" "--index-url" "https://download.pytorch.org/whl/cpu")
python -m pip install --upgrade pip
python -m pip install "${PIP_EXTRA[@]}"
# Core deps Wav2Lip's inference.py imports at module top.
#
# IMPORTANT: setuptools is pinned explicitly. On Python 3.12 (Anaconda
# / venv created with python -m venv 3.12+), `pkg_resources` is no
# longer bundled — but librosa < 0.11 still imports it from inside
# setuptools. Without this, every Wav2Lip invocation crashes with
#   ModuleNotFoundError: No module named 'pkg_resources'
# and the server silently falls back to raw audio (no lip-sync).
python -m pip install \
  setuptools \
  numpy "opencv-python<5" librosa==0.10.1 numba scipy tqdm \
  resampy "audioread" \
  soundfile

# ── 4. Pretrained checkpoint ────────────────────────────────────
mkdir -p "$CHECKPOINT_DIR"
if [ -s "$CHECKPOINT_PATH" ]; then
  echo "▶ Checkpoint already present at $CHECKPOINT_PATH — skipping."
else
  echo "▶ Downloading wav2lip_gan.pth (≈ 430 MB)…"
  # The original Google Drive link rate-limits; we pull from the
  # community mirror that hosts the same file. If this URL ever
  # 404s, replace it with any current Wav2Lip release mirror.
  WAV2LIP_URL="${WAV2LIP_URL:-https://huggingface.co/numz/wav2lip_studio/resolve/main/Wav2lip/wav2lip_gan.pth}"
  curl -L --fail --retry 3 -o "$CHECKPOINT_PATH" "$WAV2LIP_URL"
  echo "  ✓ checkpoint saved to $CHECKPOINT_PATH"
fi

# ── 5. Avatar MP4s ──────────────────────────────────────────────
echo "▶ Checking avatar MP4s under $AVATAR_DIR…"
EXPECTED=(male_one.mp4 male_two.mp4 female_one.mp4 female_two.mp4)
MISSING=()
for f in "${EXPECTED[@]}"; do
  if [ ! -s "$AVATAR_DIR/$f" ]; then MISSING+=("$f"); fi
done
if [ ${#MISSING[@]} -gt 0 ]; then
  echo "  ⚠ Missing avatar files (place these under $AVATAR_DIR before using video mode):"
  for f in "${MISSING[@]}"; do echo "      - $f"; done
else
  echo "  ✓ all 4 avatar files found."
fi

# ── 6. Smoke test ──────────────────────────────────────────────
echo "▶ Verifying Wav2Lip can import its own modules…"
python - <<'PY' || echo "  ⚠ Wav2Lip import check failed — see error above."
import sys, importlib
sys.path.insert(0, "vendor/Wav2Lip")
for mod in ("audio", "models", "face_detection"):
    try:
        importlib.import_module(mod)
        print(f"   ✓ {mod}")
    except Exception as exc:
        print(f"   ✗ {mod}: {exc}")
PY

echo ""
echo "✅ Setup complete."
echo "   • Restart your FastAPI service: systemctl restart <your-service>  (or rerun uvicorn)"
echo "   • Health check:  curl http://localhost:8000/api/_health/video"
echo "   • The first video render takes 30–90 s on CPU; cached responses replay instantly."