#!/bin/zsh
# VIDGOD bootstrap — venvs, repos, checkpoints, patches. Idempotent: re-run to resume.
# Run from anywhere: paths derive from this script's location.
# Heartbeat: ~/.jobs/vidgod-setup.status
export PATH=/opt/homebrew/bin:/usr/bin:/bin:/usr/sbin:/sbin
VG="$(cd "$(dirname "$0")/.." && pwd)"
JOB=vidgod-setup

hb(){ mkdir -p ~/.jobs; echo "$(date '+%F %T') | $1 | $2" > ~/.jobs/$JOB.status; }
die(){ hb "FAILED" "$1"; echo "FATAL: $1" >&2; exit 1; }

command -v uv >/dev/null || die "uv not installed (brew install uv)"
command -v ffmpeg >/dev/null || die "ffmpeg not installed (brew install ffmpeg)"

hb "1/10" "creating dirs"
mkdir -p "$VG"/{venvs,tools,models,library/thumbs,work,tests} || die "mkdir"

# ---- roto venv (torch/MPS + SAM2 + MatAnyone) ----
hb "2/10" "roto venv (py3.12)"
[ -x "$VG/venvs/roto/bin/python" ] || uv venv "$VG/venvs/roto" --python 3.12 || die "uv venv roto"
RPY="$VG/venvs/roto/bin/python"

hb "3/10" "torch+torchvision into roto (big download)"
uv pip install --python "$RPY" torch torchvision numpy opencv-python pillow tqdm huggingface_hub imageio imageio-ffmpeg hatchling editables || die "torch install"

hb "4/10" "clone + install SAM2"
if [ ! -d "$VG/tools/sam2" ]; then
  git clone --depth 1 https://github.com/facebookresearch/sam2 "$VG/tools/sam2" || die "sam2 clone"
fi
SAM2_BUILD_CUDA=0 uv pip install --python "$RPY" --no-build-isolation -e "$VG/tools/sam2" || die "sam2 install"

hb "5/10" "clone + install MatAnyone"
if [ ! -d "$VG/tools/MatAnyone" ]; then
  git clone --depth 1 https://github.com/pq-yang/MatAnyone "$VG/tools/MatAnyone" || die "matanyone clone"
fi
if [ -f "$VG/tools/MatAnyone/requirements.txt" ]; then
  uv pip install --python "$RPY" -r "$VG/tools/MatAnyone/requirements.txt" || echo "WARN: some MatAnyone reqs failed, continuing"
fi
if [ -f "$VG/tools/MatAnyone/pyproject.toml" ] || [ -f "$VG/tools/MatAnyone/setup.py" ]; then
  uv pip install --python "$RPY" --no-deps --no-build-isolation -e "$VG/tools/MatAnyone" || echo "WARN: matanyone -e install failed (will use sys.path)"
fi

hb "6/10" "patch MatAnyone video reader (torchvision >= 0.23 removed read_video)"
if ! grep -q "CAP_PROP_FPS" "$VG/tools/MatAnyone/matanyone/utils/inference_utils.py"; then
  git -C "$VG/tools/MatAnyone" apply "$VG/patches/matanyone-cv2-reader.patch" || die "matanyone patch failed to apply"
  echo "patch applied"
else
  echo "patch already applied"
fi

# ---- index venv (scene detect + whisper + beats) ----
hb "7/10" "index venv + packages"
[ -x "$VG/venvs/index/bin/python" ] || uv venv "$VG/venvs/index" --python 3.12 || die "uv venv index"
IPY="$VG/venvs/index/bin/python"
uv pip install --python "$IPY" mlx-whisper "scenedetect[opencv]" librosa soundfile numpy pillow tqdm || die "index pkgs"

# ---- checkpoints ----
hb "8/10" "SAM2.1 large checkpoint (~856MB)"
CKPT="$VG/models/sam2.1_hiera_large.pt"
if [ ! -s "$CKPT" ]; then
  curl -fL --retry 3 -o "$CKPT.part" https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt || die "sam2 ckpt download"
  mv "$CKPT.part" "$CKPT"
fi

hb "9/10" "MatAnyone checkpoint (~135MB, github release)"
MA="$VG/models/matanyone.pth"
if [ ! -s "$MA" ]; then
  curl -fL --retry 3 -o "$MA.part" https://github.com/pq-yang/MatAnyone/releases/download/v1.0.0/matanyone.pth || die "matanyone ckpt download"
  mv "$MA.part" "$MA"
fi
# its inference script hard-codes pretrained_models/ relative to the repo dir
mkdir -p "$VG/tools/MatAnyone/pretrained_models"
ln -sf "$MA" "$VG/tools/MatAnyone/pretrained_models/matanyone.pth"

hb "10/10" "smoke: torch/MPS + imports"
"$RPY" - <<'PYEOF' || die "roto smoke"
import torch, sam2
print("torch", torch.__version__, "mps:", torch.backends.mps.is_available())
PYEOF
"$IPY" - <<'PYEOF' || die "index smoke"
import mlx_whisper, scenedetect, librosa
print("index venv ok")
PYEOF

hb "DONE" "venvs+repos+ckpts+patches ready"
echo "SETUP DONE — try: $VG/setup/smoke_test.sh"