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

# 必要目录
mkdir -p code_depth/checkpoints \
         code_edit/stage1/checkpoint-4800 \
         code_edit/stage2/checkpoint-20000

# 通用下载函数：优先 curl，退回 wget；内建重试和断点续传
fetch() {
  local url="$1"
  local out="$2"
  # 已存在就跳过
  if [ -s "$out" ]; then
    echo "✔ Exists: $out"
    return 0
  fi
  echo "↓ Fetch: $url -> $out"
  if command -v curl >/dev/null 2>&1; then
    # --retry 对网络/5xx/超时都重试；-C - 断点续传；-f 让 4xx/5xx 变为非0退出
    curl -fL --retry 5 --retry-all-errors --connect-timeout 20 -C - -o "$out" "${url}?download=1"
  else
    # wget 也加 tries 与 continue
    wget --tries=5 -c -O "$out" "${url}?download=1"
  fi
}

# 1) VDA 权重
fetch "https://huggingface.co/depth-anything/Video-Depth-Anything-Small/resolve/main/video_depth_anything_vits.pth" \
      "code_depth/checkpoints/video_depth_anything_vits.pth"

fetch "https://huggingface.co/depth-anything/Video-Depth-Anything-Large/resolve/main/video_depth_anything_vitl.pth" \
      "code_depth/checkpoints/video_depth_anything_vitl.pth"

# 2) 你的 stage1 / stage2 两个 safetensors
fetch "https://huggingface.co/buxiangzhiren/GeoRemover/resolve/main/stage1/checkpoint-4800/pytorch_lora_weights.safetensors" \
      "code_edit/stage1/checkpoint-4800/pytorch_lora_weights.safetensors"

fetch "https://huggingface.co/buxiangzhiren/GeoRemover/resolve/main/stage2/checkpoint-20000/pytorch_lora_weights.safetensors" \
      "code_edit/stage2/checkpoint-20000/pytorch_lora_weights.safetensors"

# 最终校验：缺哪个报名字
missing=()
need=(
  "code_depth/checkpoints/video_depth_anything_vits.pth"
  "code_depth/checkpoints/video_depth_anything_vitl.pth"
  "code_edit/stage1/checkpoint-4800/pytorch_lora_weights.safetensors"
  "code_edit/stage2/checkpoint-20000/pytorch_lora_weights.safetensors"
)
for f in "${need[@]}"; do
  [ -s "$f" ] || missing+=("$f")
done
if [ ${#missing[@]} -ne 0 ]; then
  echo "❌ Missing after download:" >&2
  printf ' - %s\n' "${missing[@]}" >&2
  exit 1
fi

echo "✅ All assets ready."