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

CACHE_DIR=${VBENCH_CACHE_DIR:-/root/.cache/vbench}
HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com}

download() {
  local url=$1 dest=$2 min_bytes=${3:-10240}
  local current=0
  if [[ -f "$dest" ]]; then
    current=$(stat -c '%s' "$dest")
  fi
  if ((current >= min_bytes)); then
    printf 'cached %s (%s bytes)\n' "$dest" "$current"
    return 0
  fi
  mkdir -p "$(dirname "$dest")"
  printf 'downloading %s\n' "$url"
  curl -fL --retry 5 --retry-delay 3 --connect-timeout 30 \
    -o "${dest}.part" "$url"
  current=$(stat -c '%s' "${dest}.part")
  ((current >= min_bytes)) || {
    printf 'download too small: %s (%s bytes)\n' "$url" "$current" >&2
    return 1
  }
  mv "${dest}.part" "$dest"
  printf 'ready %s (%s bytes)\n' "$dest" "$current"
}

download \
  "$HF_ENDPOINT/lalala125/AMT/resolve/main/amt-s.pth" \
  "$CACHE_DIR/amt_model/amt-s.pth"
download \
  "$HF_ENDPOINT/OpenGVLab/VBench_Used_Models/resolve/main/l16_ptk710_ftk710_ftk400_f16_res224.pth" \
  "$CACHE_DIR/umt_model/l16_ptk710_ftk710_ftk400_f16_res224.pth"
download \
  "$HF_ENDPOINT/OpenGVLab/VBench_Used_Models/resolve/main/grit_b_densecap_objectdet.pth" \
  "$CACHE_DIR/grit_model/grit_b_densecap_objectdet.pth"
download \
  "$HF_ENDPOINT/spaces/xinyu1205/recognize-anything/resolve/main/tag2text_swin_14m.pth" \
  "$CACHE_DIR/caption_model/tag2text_swin_14m.pth"
download \
  "$HF_ENDPOINT/OpenGVLab/VBench_Used_Models/resolve/main/ViClip-InternVid-10M-FLT.pth" \
  "$CACHE_DIR/ViCLIP/ViClip-InternVid-10M-FLT.pth"
download \
  "$HF_ENDPOINT/chaofengc/IQA-PyTorch-Weights/resolve/main/musiq_spaq_ckpt-358bb6af.pth" \
  "$CACHE_DIR/pyiqa_model/musiq_spaq_ckpt-358bb6af.pth"
download \
  "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt" \
  "/root/.cache/clip/ViT-B-32.pt"

printf 'VBench scoring weights prefetched successfully\n'