#!/usr/bin/env bash
# 下载 P0 基准所需模型。走 hf-mirror.com（本机 huggingface.co 直连不通）。
# 可重复执行：已存在且大小正确的文件会跳过。
set -uo pipefail

ROOT="$(cd "$(dirname "$0")" && pwd)"
MIRROR="${HF_MIRROR:-https://hf-mirror.com}"
STT_REPO="csukuangfj/sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20"
TTS_REPO="csukuangfj/sherpa-onnx-vits-zh-ll"

STT_DIR="$ROOT/models/stt"
TTS_DIR="$ROOT/models/tts"
WAV_DIR="$ROOT/wavs"
mkdir -p "$STT_DIR" "$TTS_DIR" "$WAV_DIR"

fetch() { # fetch <repo> <remote-path> <local-path>
  local repo="$1" remote="$2" local="$3"
  mkdir -p "$(dirname "$local")"
  local url="$MIRROR/$repo/resolve/main/$remote"
  local expect
  expect=$(curl -sIL "$url" | grep -i '^content-length' | tail -1 | tr -d '\r' | awk '{print $2}')
  if [ -f "$local" ] && [ -n "$expect" ] && [ "$(stat -c%s "$local")" = "$expect" ]; then
    printf '  skip  %-45s %s\n' "$(basename "$remote")" "(已完整)"
    return 0
  fi
  printf '  get   %-45s %s\n' "$(basename "$remote")" "${expect:+$((expect/1048576))MB}"
  curl -sL --retry 3 --retry-delay 2 -o "$local.part" "$url" || { echo "    FAILED"; return 1; }
  mv "$local.part" "$local"
}

echo "=== STT: 流式 Zipformer 中英双语 int8 ==="
for f in tokens.txt bpe.model \
         encoder-epoch-99-avg-1.int8.onnx \
         decoder-epoch-99-avg-1.int8.onnx \
         joiner-epoch-99-avg-1.int8.onnx; do
  fetch "$STT_REPO" "$f" "$STT_DIR/$f"
done
for w in 0.wav 1.wav 2.wav 3.wav; do
  fetch "$STT_REPO" "test_wavs/$w" "$WAV_DIR/stt_$w"
done

echo "=== TTS: VITS zh-ll ==="
for f in model.onnx tokens.txt lexicon.txt date.fst number.fst phone.fst new_heteronym.fst \
         dict/jieba.dict.utf8 dict/hmm_model.utf8 dict/idf.utf8 \
         dict/pos_dict/char_state_tab.utf8 dict/pos_dict/prob_emit.utf8 dict/pos_dict/prob_trans.utf8; do
  fetch "$TTS_REPO" "$f" "$TTS_DIR/$f"
done

echo
echo "=== 结果 ==="
du -sh "$STT_DIR" "$TTS_DIR" "$WAV_DIR" 2>/dev/null
echo "下载完成。"