#!/usr/bin/env bash
# 预下载 ChatTTS 模型到本地（走 HuggingFace 镜像，不依赖 GitHub）
# 用法: bash scripts/download_chattts_models.sh
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
MODEL_DIR="${CHATTTS_MODEL_DIR:-${ROOT}/models/ChatTTS}"
export MODEL_DIR
VENV_PY="${ROOT}/venv/bin/python"

export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
export HF_HUB_DOWNLOAD_TIMEOUT="${HF_HUB_DOWNLOAD_TIMEOUT:-600}"
HF_HOME="${HF_HOME:-${ROOT}/models/hf_cache}"

echo "[INFO] ChatTTS 模型目录: ${MODEL_DIR}"
echo "[INFO] HF 镜像: ${HF_ENDPOINT}"
mkdir -p "${MODEL_DIR}" "${HF_HOME}"

if [[ ! -x "${VENV_PY}" ]]; then
  echo "[ERROR] 未找到 venv，请先 bash deploy.sh deps"
  exit 1
fi

"${VENV_PY}" -m pip install -q huggingface_hub

"${VENV_PY}" << PY
import os
from pathlib import Path
from huggingface_hub import snapshot_download

target = Path(os.environ["MODEL_DIR"])
target.mkdir(parents=True, exist_ok=True)

print("[INFO] 正在从 HuggingFace 镜像下载 2Noise/ChatTTS（约 1-2GB）...")
snapshot_download(
    repo_id="2Noise/ChatTTS",
    local_dir=str(target),
    local_dir_use_symlinks=False,
)
print("[OK] 下载完成:", target)

# 简单校验
checks = [
    target / "asset",
    target / "config" / "path.yaml",
]
ok = any(p.exists() for p in checks)
if not ok:
    print("[WARN] 目录结构异常，请检查下载是否完整")
else:
    print("[OK] 模型目录结构校验通过")
PY

echo ""
echo "[OK] 请执行: pm2 restart trading_studio"
echo "     然后重新点击「锁定音色」"