# -*- coding: utf-8 -*-
"""
把 hub_package/ 上传到 Hugging Face Hub（对应 PPT 第 28 页第 4 条学习目标）。
================================================================
等价于 PPT 里的 package_to_hub(...)，只是拆成"先打包、再上传"两步，更透明。

准备：
  1) 注册 https://huggingface.co ，在 Settings -> Access Tokens 建一个 write token
  2) 登录（二选一）：
       hf auth login              # 交互式粘贴 token
       或者直接设环境变量：$env:HF_TOKEN = "hf_xxx"
  3) 打包：python build_hub_package.py --repo-id 你的用户名/value-based-qlearning

上传：
  python upload_to_hub.py --repo-id 你的用户名/value-based-qlearning
  python upload_to_hub.py --repo-id 你的用户名/value-based-qlearning --private

如果本机连不上 huggingface.co（本机实测会超时），有三条路：
  A) 开代理后重试本脚本；
  B) 换镜像（只能下载不能上传，不适用于上传）；
  C) 浏览器手动上传：https://huggingface.co/new 建仓库，把 hub_package/ 里的文件拖进去。
"""

import argparse
import os
import sys

BASE_DIR = os.path.dirname(os.path.abspath(__file__))


def list_files(folder):
    rows = []
    for root, _, files in os.walk(folder):
        for name in sorted(files):
            path = os.path.join(root, name)
            rows.append((os.path.relpath(path, folder), os.path.getsize(path)))
    return rows


def main():
    parser = argparse.ArgumentParser(description="上传 hub_package 到 Hugging Face Hub")
    parser.add_argument("--repo-id", required=True, help="用户名/仓库名，例如 zhangsan/value-based-qlearning")
    parser.add_argument("--folder", default=os.path.join(BASE_DIR, "hub_package"))
    parser.add_argument("--token", default=os.environ.get("HF_TOKEN"), help="默认读环境变量 HF_TOKEN")
    parser.add_argument("--private", action="store_true", help="创建私有仓库")
    parser.add_argument("--endpoint", default=os.environ.get("HF_ENDPOINT"),
                        help="默认 huggingface.co，可用环境变量 HF_ENDPOINT 覆盖")
    parser.add_argument("--dry-run", action="store_true", help="只列出要上传的文件，不真的上传")
    args = parser.parse_args()

    if not os.path.isdir(args.folder):
        raise SystemExit(f"找不到 {args.folder}，请先运行 python build_hub_package.py")

    files = list_files(args.folder)
    total_mb = sum(size for _, size in files) / (1024 * 1024)
    print(f"准备上传 {len(files)} 个文件（共 {total_mb:.1f} MB）到仓库 {args.repo_id}")
    for rel, size in files:
        print(f"  {rel}  ({size / 1024:.0f} KB)")

    if args.dry_run:
        print("\n[--dry-run] 不上传。")
        return

    if not args.token:
        print("\n[错误] 没有 token。请先执行 hf auth login，或设置 $env:HF_TOKEN。")
        print("       也可以去掉 --token 参数让 huggingface_hub 读本机缓存的登录信息。")
        return

    try:
        from huggingface_hub import HfApi
    except ImportError:
        print("\n[错误] 缺 huggingface_hub，请先：pip install huggingface_hub")
        return

    api = HfApi(endpoint=args.endpoint, token=args.token) if args.endpoint else HfApi(token=args.token)

    try:
        print(f"\n[1/2] 创建/确认仓库 {args.repo_id} ...")
        api.create_repo(repo_id=args.repo_id, repo_type="model",
                        private=args.private, exist_ok=True)
        print("[2/2] 上传文件 ...")
        api.upload_folder(
            repo_id=args.repo_id,
            repo_type="model",
            folder_path=args.folder,
            commit_message="value-based 实践：Q-Learning 智能体（FrozenLake + Taxi）"
                           " + 回放视频 + 评估分数",
        )
    except Exception as exc:                       # noqa: BLE001
        print(f"\n[上传失败] {type(exc).__name__}: {exc}")
        print("\n常见原因与对策：")
        print("  1) 网络连不上 huggingface.co（本机常见）-> 开代理后重试，或浏览器手动上传；")
        print("  2) token 没有 write 权限 -> 重新建一个 write token；")
        print("  3) 仓库名格式不对 -> 必须是 用户名/仓库名。")
        print(f"\n手动上传入口：https://huggingface.co/new  （仓库名建议 value-based-qlearning）")
        return

    url = f"https://huggingface.co/{args.repo_id}"
    print(f"\n[完成] 仓库地址：{url}")
    print(f"       视频和评估分数已经写进 README.md，打开仓库就能看到。")


if __name__ == "__main__":
    sys.exit(main())