#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
将本地 DiffAM 文件夹中的所有内容上传到 Hugging Face Model Repository：

本地目录：
    /Users/gaobj/Downloads/DiffAM

目标仓库：
    https://huggingface.co/bingjie/DiffAM

安装依赖：
    pip install -U huggingface_hub

运行前设置 Token：
    export HF_TOKEN="你的新HuggingFaceToken"

运行：
    python upload_diffam_to_huggingface.py
"""

import os
import sys
from pathlib import Path

from huggingface_hub import HfApi
from huggingface_hub.errors import HfHubHTTPError


# ============================================================
# 配置
# ============================================================

LOCAL_FOLDER = Path("/Users/gaobj/Downloads/DiffAM")

# Hugging Face 模型仓库格式为：用户名/仓库名
REPO_ID = "bingjie/DiffAM"

# 可选：
# False：公开仓库
# True：私有仓库
PRIVATE_REPO = False

COMMIT_MESSAGE = "Upload DiffAM model files"


def get_hf_token() -> str:
    """
    从环境变量读取 Hugging Face Access Token。
    """
    token = os.environ.get("HF_TOKEN")

    if not token:
        raise RuntimeError(
            "没有找到环境变量 HF_TOKEN。\n\n"
            "请先在终端执行：\n"
            'export HF_TOKEN="你的HuggingFaceAccessToken"\n\n'
            "然后重新运行此脚本。"
        )

    return token.strip()


def check_local_folder(folder: Path) -> None:
    """
    检查本地目录是否存在且不为空。
    """
    if not folder.exists():
        raise FileNotFoundError(f"本地目录不存在：{folder}")

    if not folder.is_dir():
        raise NotADirectoryError(f"该路径不是文件夹：{folder}")

    files = [path for path in folder.rglob("*") if path.is_file()]

    if not files:
        raise RuntimeError(f"本地目录为空，没有可上传的文件：{folder}")

    total_size = sum(path.stat().st_size for path in files)

    print(f"本地目录：{folder}")
    print(f"文件数量：{len(files)}")
    print(f"文件总大小：{format_size(total_size)}")


def format_size(size_bytes: int) -> str:
    """
    将字节数转换为便于阅读的格式。
    """
    size = float(size_bytes)

    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if size < 1024 or unit == "TB":
            return f"{size:.2f} {unit}"
        size /= 1024

    return f"{size_bytes} B"


def upload_folder() -> None:
    """
    创建 Hugging Face 模型仓库，并上传整个本地文件夹。
    """
    check_local_folder(LOCAL_FOLDER)

    token = get_hf_token()
    api = HfApi(token=token)

    try:
        # 验证 Token，并获取当前账户信息
        user_info = api.whoami()
        username = user_info.get("name", "unknown")

        print(f"当前 Hugging Face 账户：{username}")

        if username != "bingjie":
            print(
                f"警告：当前 Token 所属账户是 {username}，"
                f"但目标仓库位于 bingjie 名下。"
            )

        # 创建仓库；exist_ok=True 表示仓库已存在时不会报错
        repo_url = api.create_repo(
            repo_id=REPO_ID,
            repo_type="model",
            private=PRIVATE_REPO,
            exist_ok=True,
        )

        print(f"目标仓库：{repo_url}")
        print("开始上传文件……")

        # 上传文件夹内的全部内容
        # path_in_repo="" 表示上传到仓库根目录
        commit_info = api.upload_folder(
            folder_path=str(LOCAL_FOLDER),
            repo_id=REPO_ID,
            repo_type="model",
            path_in_repo="",
            commit_message=COMMIT_MESSAGE,

            # 忽略常见的本地缓存和系统文件
            ignore_patterns=[
                ".git/**",
                ".DS_Store",
                "**/.DS_Store",
                "__pycache__/**",
                "**/__pycache__/**",
                "*.pyc",
                "**/*.pyc",
            ],
        )

        print("\n上传完成。")
        print(f"模型仓库：https://huggingface.co/{REPO_ID}")

        commit_url = getattr(commit_info, "commit_url", None)
        if commit_url:
            print(f"本次提交：{commit_url}")

    except HfHubHTTPError as exc:
        print("\nHugging Face 请求失败：", file=sys.stderr)
        print(str(exc), file=sys.stderr)

        if exc.response is not None:
            if exc.response.status_code == 401:
                print(
                    "\n可能原因：Token 无效、已过期或没有登录权限。",
                    file=sys.stderr,
                )
            elif exc.response.status_code == 403:
                print(
                    "\n可能原因：Token 没有写入权限，"
                    "或者当前账户无权写入 bingjie/DiffAM。",
                    file=sys.stderr,
                )

        raise

    except KeyboardInterrupt:
        print("\n上传已由用户中断。", file=sys.stderr)
        sys.exit(130)


if __name__ == "__main__":
    upload_folder()