#!/usr/bin/env bash
# 预拉模型权重到 ./models volume，避免容器启动时才下载。
#
# 用法:
#   ./scripts/prefetch_models.sh                  # 读 config.yaml（当前激活配置）
#   ./scripts/prefetch_models.sh config.gpu.yaml  # 读指定配置文件
#
# 原理：用已构建的 audio2text 镜像跑一次性容器，挂载 ./models volume，
#       执行 scripts/prefetch_models.py 把模型下到 HF cache。
#       容器运行时 HF_HOME=/models/huggingface 命中缓存，秒级加载。
#
# 幂等：已下过的模型跳过。换 config 的 model 后重跑即可补下新模型。
set -euo pipefail

cd "$(dirname "$0")/.."
ROOT="$(pwd)"

CONFIG="${1:-config.yaml}"
VARIANT="${AUDIO2TEXT_VARIANT:-cpu}"

# 选镜像：优先用 gpu 镜像（GPU 模型大，可能要 cuda 才能 correctly 下载部分文件），
# 否则 cpu。两者都能下 HF 模型（下载是纯网络 IO，不依赖 GPU）。
IMAGE="audio2text:cpu"
if docker image inspect audio2text:gpu >/dev/null 2>&1 && [ "$VARIANT" = "gpu" ]; then
    IMAGE="audio2text:gpu"
fi

if ! docker image inspect "$IMAGE" >/dev/null 2>&1; then
    echo "✗ 镜像 $IMAGE 不存在，请先 ./setup.sh 构建。" >&2
    exit 1
fi

if [ ! -f "$ROOT/$CONFIG" ]; then
    echo "✗ 配置文件不存在: $ROOT/$CONFIG" >&2
    exit 1
fi

mkdir -p "$ROOT/models"

echo "==> 预拉模型"
echo "    镜像: $IMAGE"
echo "    配置: $CONFIG"
echo "    缓存: $ROOT/models (volume)"
echo

# 注意：CONFIG_PATH 指向容器内路径，挂载配置文件为只读
MSYS_NO_PATHCONV=1 docker run --rm \
    -v "$ROOT/models:/models" \
    -v "$ROOT/$CONFIG:/app/config.yaml:ro" \
    -e HF_HOME=/models/huggingface \
    -e CT2_CACHE=/models/ctranslate2 \
    -e CONFIG_PATH=/app/config.yaml \
    "$IMAGE" \
    python /app/scripts/prefetch_models.py /app/config.yaml

echo
echo "==> 完成。模型已缓存到 $ROOT/models/huggingface"
echo "    容器启动时会命中缓存，无需联网下载。"