#!/usr/bin/env bash

# ==========================================
# Whisper ggml 模型下载脚本
# 下载的模型将存放在 data/model 目录下
# ==========================================

# 确保在项目根目录运行
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TARGET_DIR="${SCRIPT_DIR}/data/model"

# 创建目标目录
mkdir -p "$TARGET_DIR"

# Hugging Face 的 whisper.cpp 模型仓库地址
# 使用 hf-mirror 镜像加速国内下载
BASE_URL="https://hf-mirror.com/ggerganov/whisper.cpp/resolve/main"

echo "=========================================="
echo "请选择要下载的 Whisper 语音识别模型："
echo "1) tiny (最小，速度最快，约 40MB)"
echo "2) base (基础，速度快，约 74MB)"
echo "3) small (当前默认，平衡，约 240MB)"
echo "4) medium (推荐，效果更好，约 770MB)"
echo "5) large-v3 (最大，效果最好，约 1.5GB)"
echo "=========================================="

read -p "请输入对应的数字 [1-5，默认 4]: " choice

MODEL_NAME=""
case $choice in
    1) MODEL_NAME="ggml-tiny-q8_0.bin" ;;
    2) MODEL_NAME="ggml-base-q8_0.bin" ;;
    3) MODEL_NAME="ggml-small-q8_0.bin" ;;
    4|"") MODEL_NAME="ggml-medium-q8_0.bin" ;;
    5) MODEL_NAME="ggml-large-v3-q8_0.bin" ;;
    *) 
        echo "无效的选择，使用默认选项 4 (medium)"
        MODEL_NAME="ggml-medium-q8_0.bin"
        ;;
esac

echo ""
echo "准备下载模型: $MODEL_NAME"
echo "保存路径: $TARGET_DIR/$MODEL_NAME"
echo ""

# 下载文件
DOWNLOAD_URL="${BASE_URL}/${MODEL_NAME}"

if command -v wget >/dev/null 2>&1; then
    echo "使用 wget 下载中..."
    wget -c --show-progress "$DOWNLOAD_URL" -O "$TARGET_DIR/$MODEL_NAME"
elif command -v curl >/dev/null 2>&1; then
    echo "使用 curl 下载中..."
    curl -L -C - -o "$TARGET_DIR/$MODEL_NAME" "$DOWNLOAD_URL"
else
    echo "错误：找不到 wget 或 curl 命令，无法下载文件。"
    exit 1
fi

if [ $? -eq 0 ]; then
    echo ""
    echo "✅ 下载成功！"
    echo "模型文件已保存到: $TARGET_DIR/$MODEL_NAME"
    echo ""
    echo "请记得更新 .env.prod 或 backend/.env 文件中的配置："
    echo "STT_MODEL=$MODEL_NAME"
    echo "修改后重启后端容器即可生效。"
else
    echo ""
    echo "❌ 下载失败，请检查网络连接或稍后重试。"
    exit 1
fi