#!/bin/bash
# 并行分块下载 hf-mirror 上的大文件（绕开单连接限速），完成后原子落位
set -u
M=/home/zyw/ComfyUI/models
BASE=https://hf-mirror.com
TMP=/home/zyw/yue2-par
CHUNKS=${CHUNKS:-4}
JOBS=${JOBS:-2}
mkdir -p "$TMP"

sig_url() {  # 取签名直链
  curl -sI -m 30 "$1" | grep -i '^location:' | tail -1 | tr -d '\r' | sed 's/^[Ll]ocation: //'
}
size_of() {
  curl -sIL -m 30 "$1" | grep -i '^content-length' | tail -1 | tr -d '\r' | awk '{print $2}'
}

dl_one() {  # repo rev path destdir
  local repo="$1" rev="$2" rel="$3" dest="$4"
  local out="$dest/$rel" url="$BASE/$repo/resolve/$rev/$rel"
  mkdir -p "$dest/$(dirname "$rel")"
  if [ -s "$out" ]; then
    local e; e=$(size_of "$url")
    [ -n "$e" ] && [ "$(stat -c%s "$out")" = "$e" ] && { echo "SKIP $rel"; return 0; }
  fi
  local sz; sz=$(size_of "$url")
  [ -z "$sz" ] && { echo "FAIL(no size) $rel"; return 1; }
  local work="$TMP/$(echo "$repo-$rel" | tr '/' '_')"
  rm -rf "$work"; mkdir -p "$work"
  local part=$(( sz / CHUNKS + 1 ))
  local pids=()
  for i in $(seq 0 $((CHUNKS-1))); do
    local s=$(( i * part )); local e=$(( s + part - 1 )); [ $e -ge $sz ] && e=$(( sz - 1 ))
    ( curl -L --retry 8 --retry-delay 3 --retry-all-errors -sS -r "$s-$e" -o "$work/p$i" "$url" ) &
    pids+=($!)
  done
  local ok=1
  for p in "${pids[@]}"; do wait "$p" || ok=0; done
  cat "$work"/p* > "$work/full" 2>/dev/null
  local got; got=$(stat -c%s "$work/full" 2>/dev/null || echo 0)
  if [ "$got" != "$sz" ]; then echo "FAIL $rel (得到 $got / 期望 $sz)"; return 1; fi
  mv "$work/full" "$out"; rm -rf "$work"
  echo "OK $rel $(stat -c%s "$out")"
}

# 并行跑指定数量的大文件
run_group() {  # 每个参数 "repo|rev|rel|dest"
  local pids=() item
  for item in "$@"; do
    IFS='|' read -r repo rev rel dest <<< "$item"
    dl_one "$repo" "$rev" "$rel" "$dest" >> "$TMP/log.txt" 2>&1 &
    pids+=($!)
  done
  for p in "${pids[@]}"; do wait "$p"; done
}

echo "=== 并行下载开始 $(date +%T) ==="
run_group \
  "m-a-p/SheetSage2|main|model.safetensors|$M/SheetSage2" \
  "m-a-p/YuE2-Vae|main|model.safetensors|$M/YuE2/YuE2-Vae"
run_group \
  "m-a-p/MERT-v2-FullSong|d8ba1c745e733b3908ce6ad16ebeb17ac7600a42|model.safetensors|$M/MERT2/MERT-v2-FullSong"
echo "=== 结束 $(date +%T) ==="
cat "$TMP/log.txt"
du -sh $M/YuE2/YuE2-3B $M/YuE2/YuE2-Vae $M/SheetSage2 $M/MERT2/MERT-v2-FullSong