#!/usr/bin/env python3
"""
从 Hugging Face 下载模型到本地 model 文件夹
支持镜像源、自动按模型名创建子文件夹
"""

import argparse
import os
from pathlib import Path

# 项目根目录
PROJECT_ROOT = Path(__file__).parent.parent.resolve()
# 默认模型保存目录：项目根目录/model/weights
DEFAULT_MODEL_DIR = PROJECT_ROOT / "model" / "weights"


def download_model(model_id: str, output_dir: str = None, use_mirror: bool = True):
    """
    从 Hugging Face 下载模型

    Args:
        model_id: 模型 ID，格式为 "username/model_name"
        output_dir: 输出根目录，默认为项目根目录/model/weights
        use_mirror: 是否使用镜像源，默认 True
    """
    if output_dir is None:
        output_dir = DEFAULT_MODEL_DIR

    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("错误：未安装 huggingface_hub 库")
        print("请运行：pip install huggingface_hub")
        return

    # 提取模型名作为子文件夹名
    model_name = model_id.split("/")[-1]
    output_path = Path(output_dir).resolve() / model_name
    output_path.mkdir(parents=True, exist_ok=True)

    print(f"开始下载模型：{model_id}")
    print(f"保存路径：{output_path}")

    if use_mirror:
        print("使用镜像源：hf-mirror.com")

    print("-" * 50)

    try:
        # 设置镜像源环境变量
        if use_mirror:
            os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

        # 下载整个模型仓库
        downloaded_path = snapshot_download(
            repo_id=model_id,
            local_dir=str(output_path),
            local_dir_use_symlinks=False,  # Windows 兼容性
        )

        print("-" * 50)
        print(f"✓ 模型下载完成！")
        print(f"模型位置：{downloaded_path}")

        # 列出下载的文件
        files = list(output_path.iterdir())
        print(f"\n下载的文件 ({len(files)} 个):")
        for f in files:
            if f.is_file():
                size = f.stat().st_size
                if size > 1024 * 1024 * 1024:
                    size_str = f"{size / (1024*1024*1024):.2f} GB"
                elif size > 1024 * 1024:
                    size_str = f"{size / (1024*1024):.2f} MB"
                elif size > 1024:
                    size_str = f"{size / 1024:.2f} KB"
                else:
                    size_str = f"{size} B"
                print(f"  - {f.name} ({size_str})")

    except Exception as e:
        print(f"下载失败：{e}")
        raise


def main():
    parser = argparse.ArgumentParser(
        description="从 Hugging Face 下载模型到本地",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=f"""
示例:
  python downloadModel.py hf-internal-testing/tiny-random-BertModel
  python downloadModel.py bert-base-chinese --no-mirror
  python downloadModel.py Qwen/Qwen2.5-7B-Instruct -o ./my_models

默认保存目录：{DEFAULT_MODEL_DIR}
        """
    )

    parser.add_argument(
        "model_id",
        type=str,
        help="Hugging Face 模型 ID (格式：username/model_name)"
    )

    parser.add_argument(
        "-o", "--output",
        type=str,
        default=None,
        help=f"模型保存根目录 (默认：{DEFAULT_MODEL_DIR})"
    )

    parser.add_argument(
        "--no-mirror",
        action="store_true",
        help="不使用镜像源，直接使用 Hugging Face 官方源"
    )

    args = parser.parse_args()

    download_model(args.model_id, args.output, use_mirror=not args.no_mirror)


if __name__ == "__main__":
    main()