#!/bin/zsh
# VIDGOD phase 2 — ProPainter (vg-remove), Cutie (vg-cutie), RIFE (vg-interp):
# clones, weights, deps, patches. Idempotent: re-run to resume.
# Heartbeat: ~/.jobs/vidgod-phase2.status
export PATH=/opt/homebrew/bin:/usr/bin:/bin:/usr/sbin:/sbin
VG="$(cd "$(dirname "$0")/.." && pwd)"
JOB=vidgod-phase2
hb(){ mkdir -p ~/.jobs; echo "$(date '+%F %T') | $1 | $2" > ~/.jobs/$JOB.status; }
die(){ hb "FAILED" "$1"; echo "FATAL: $1" >&2; exit 1; }
get(){ # get URL DEST
  [ -s "$2" ] && return 0
  curl -sfL --retry 3 -o "$2.part" "$1" || return 1
  mv "$2.part" "$2"
}
RPY="$VG/venvs/roto/bin/python"
[ -x "$RPY" ] || die "run setup_venvs.sh first (no roto venv)"

hb "1/8" "clone ProPainter"
[ -d "$VG/tools/ProPainter" ] || git clone --depth 1 https://github.com/sczhou/ProPainter "$VG/tools/ProPainter" || die "propainter clone"

hb "2/8" "ProPainter weights (4 files, ~500MB)"
mkdir -p "$VG/tools/ProPainter/weights"
B=https://github.com/sczhou/ProPainter/releases/download/v0.1.0
get $B/ProPainter.pth "$VG/tools/ProPainter/weights/ProPainter.pth" || die "ProPainter.pth"
get $B/recurrent_flow_completion.pth "$VG/tools/ProPainter/weights/recurrent_flow_completion.pth" || die "rfc.pth"
get $B/raft-things.pth "$VG/tools/ProPainter/weights/raft-things.pth" || die "raft.pth"
get $B/i3d_rgb_imagenet.pt "$VG/tools/ProPainter/weights/i3d_rgb_imagenet.pt" || echo "WARN: i3d optional weight failed"

hb "3/8" "patch ProPainter video reader (torchvision >= 0.23 removed read_video)"
if ! grep -q "CAP_PROP_FPS" "$VG/tools/ProPainter/inference_propainter.py"; then
  git -C "$VG/tools/ProPainter" apply "$VG/patches/propainter-cv2-reader.patch" || die "propainter patch"
  echo "patch applied"
fi

hb "4/8" "clone Cutie"
[ -d "$VG/tools/Cutie" ] || git clone --depth 1 https://github.com/hkchengrex/Cutie "$VG/tools/Cutie" || die "cutie clone"

hb "5/8" "Cutie weights + device patch"
mkdir -p "$VG/tools/Cutie/weights"
get https://github.com/hkchengrex/Cutie/releases/download/v1.0/cutie-base-mega.pth "$VG/tools/Cutie/weights/cutie-base-mega.pth" || die "cutie weights"
if ! grep -q "mps" "$VG/tools/Cutie/cutie/utils/get_default_model.py"; then
  git -C "$VG/tools/Cutie" apply "$VG/patches/cutie-device.patch" || die "cutie patch"
  echo "patch applied"
fi

hb "6/8" "RIFE ncnn macos build (universal binary + models, ~450MB zip)"
mkdir -p "$VG/tools/rife"
if [ ! -x "$VG/tools/rife/rife-ncnn-vulkan" ]; then
  get https://github.com/nihui/rife-ncnn-vulkan/releases/download/20221029/rife-ncnn-vulkan-20221029-macos.zip /tmp/rife.zip || die "rife zip"
  ditto -x -k /tmp/rife.zip "$VG/tools/rife-extract" || die "rife unzip"
  inner=$(find "$VG/tools/rife-extract" -name "rife-ncnn-vulkan" -type f | head -1)
  [ -n "$inner" ] || die "rife binary not in zip"
  cp -R "$(dirname "$inner")"/* "$VG/tools/rife/"
  chmod +x "$VG/tools/rife/rife-ncnn-vulkan"
  rm -rf "$VG/tools/rife-extract" /tmp/rife.zip
fi

hb "7/8" "python deps (ProPainter + Cutie) into roto venv"
uv pip install --python "$RPY" einops scipy av addict pyyaml future matplotlib \
  PySide6 easydict gitpython gdown hickle pycocotools charset-normalizer || die "deps"
uv pip install --python "$RPY" --no-deps -e "$VG/tools/Cutie" || die "cutie -e"

hb "8/8" "smoke: imports + rife arch"
"$RPY" -c "import cutie, matplotlib, einops, scipy; print('phase2 imports ok')" || die "import smoke"
file "$VG/tools/rife/rife-ncnn-vulkan" | grep -q arm64 || echo "WARN: rife binary has no arm64 slice"

hb "DONE" "phase2 tools + weights + patches ready"
echo "PHASE2 FETCH DONE — vg-remove / vg-interp / vg-cutie ready"