"""
down_embedding_model.py
下载 Qwen3-Embedding-4B 嵌入模型
"""

import os
import sys
from pathlib import Path

def download_model():
    """下载 Qwen3-Embedding-4B 模型"""
    project_root = Path(__file__).resolve().parent
    model_dir = project_root / "models" / "Qwen3-Embedding-4B"

    print("=" * 60)
    print("Qwen3-Embedding-4B 模型下载")
    print("=" * 60)
    print(f"模型目录: {model_dir}")
    print()

    # 使用国内镜像加速
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    print("使用 HuggingFace 镜像: https://hf-mirror.com")
    print()

    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("❌ 未安装 huggingface_hub，正在安装...")
        import subprocess
        subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub"])
        from huggingface_hub import snapshot_download

    # 创建模型目录
    os.makedirs(model_dir, exist_ok=True)

    print(f"开始下载 Qwen3-Embedding-4B 模型...")
    print(f"模型大小约 4GB，请耐心等待...")
    print()

    try:
        snapshot_download(
            repo_id="Qwen/Qwen3-Embedding-4B",
            local_dir=str(model_dir),
            local_dir_use_symlinks=False,
            resume_download=True
        )
        print()
        print("=" * 60)
        print("✅ 模型下载完成！")
        print("=" * 60)
        print(f"模型路径: {model_dir}")
        return True
    except Exception as e:
        print()
        print("=" * 60)
        print(f"❌ 下载失败: {e}")
        print("=" * 60)
        return False

if __name__ == "__main__":
    download_model()