#!/usr/bin/env python3
"""学習の最後のチェックポイントを Hugging Face の非公開リポジトリへ上げる (Pod で動かす. 2026-10-31 依頼 49).

lab_train.sh が学習のあとに Pod で流す:
  python3 hf_upload_ckpt.py --dir /workspace/openpi/checkpoints/<設定>/<exp>/<ステップ> --repo PixelstoMotors/<名前>

上げるもの: train_state/ 以外の全部 (params/・assets/・_CHECKPOINT_METADATA). リポジトリの一番上に置く
(lab_start.sh の 3c が hf download で /workspace/ckpt/3999 に取る 3999 と同じ形).
鍵は hf が読む ~/.cache/huggingface/token (lab_start.sh / lab_train.sh の 3b が置く). 鍵は出さない.

1. リポジトリを非公開で作る (あればそのまま). 非公開でなければ上げずに止まる
2. リポジトリにもうファイルがあれば, 上げるものと名前・バイト数が全部同じなら上げない (same). 違えば止まる
3. upload_folder で 1 回のコミットで上げる
4. リポジトリの一覧 (名前・バイト数) が手元と全部合うかを確かめる

最後の行: HFUP <結果> <説明>. 終了コード 0 = ok / same, 3 = not_private, 4 = other_content, 5 = mismatch, 6 = error
"""
import argparse
import os
import sys

SKIP_TOP = ("train_state",)
SKIP_REMOTE = (".gitattributes",)


def local_files(root):
    """train_state/ を除くファイルの (パス, バイト数) を名前順に"""
    out = []

    def err(e):
        raise e             # 読めないフォルダを黙って飛ばさない

    for d, dirs, files in os.walk(root, onerror=err):
        rel = os.path.relpath(d, root)
        if rel == ".":
            dirs[:] = [x for x in dirs if x not in SKIP_TOP]
        for f in files:
            p = os.path.normpath(os.path.join(rel, f))
            out.append((p, os.path.getsize(os.path.join(d, f))))
    return sorted(out, key=lambda x: x[0].encode())


def remote_files(api, repo):
    """リポジトリのファイルの (パス, バイト数) を名前順に (.gitattributes は数えない)"""
    out = []
    for x in api.list_repo_tree(repo, recursive=True, repo_type="model"):
        size = getattr(x, "size", None)
        if size is None or x.path in SKIP_REMOTE:      # フォルダには size がない
            continue
        out.append((x.path, int(size)))
    return sorted(out, key=lambda x: x[0].encode())


def diff_lines(a, b, n=10):
    da, db = dict(a), dict(b)
    lines = []
    for p in sorted(set(da) | set(db), key=str.encode):
        if da.get(p) != db.get(p):
            lines.append(f"  {p}: 手元 {da.get(p, 'なし')} / リポジトリ {db.get(p, 'なし')}")
    return lines[:n] + ([f"  ほか {len(lines) - n} 件"] if len(lines) > n else [])


def done(kind, msg, code):
    print(f"HFUP {kind} {msg}", flush=True)
    return code


def main(argv=None, api=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", required=True, help="上げるチェックポイントのフォルダ (<ステップ>/)")
    ap.add_argument("--repo", required=True, help="Hugging Face のリポジトリ (例 PixelstoMotors/pi05-mecharm-trim-corr-lora-v1)")
    ap.add_argument("--message", default="", help="コミットの文")
    args = ap.parse_args(argv)
    if not os.path.isfile(os.path.join(args.dir, "_CHECKPOINT_METADATA")):
        return done("error", f"{args.dir} に _CHECKPOINT_METADATA がない (書き終えたチェックポイントではない)", 6)
    try:
        mine = local_files(args.dir)
    except OSError as e:
        return done("error", f"手元の一覧を作れない: {e}", 6)
    total = sum(s for _, s in mine)
    print(f"[HF    ] 上げるもの: {len(mine)} ファイル {total:,} バイト (train_state/ は上げない)", flush=True)
    try:
        if api is None:
            from huggingface_hub import HfApi
            api = HfApi()
        api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True)
        if not api.model_info(args.repo).private:
            return done("not_private", f"{args.repo} が非公開になっていない. 上げずに止めた", 3)
        have = remote_files(api, args.repo)
        if have:
            if have == mine:
                return done("same", f"{args.repo} にはもう同じ中身がある ({len(mine)} ファイル). 上げ直さない", 0)
            print("\n".join(diff_lines(mine, have)), flush=True)
            return done("other_content", f"{args.repo} に別の中身がある. 上書きしないで止めた", 4)
        api.upload_folder(folder_path=args.dir, repo_id=args.repo, repo_type="model",
                          ignore_patterns=[f"{t}/**" for t in SKIP_TOP] + [f"{t}/*" for t in SKIP_TOP],
                          commit_message=args.message or f"upload {os.path.basename(os.path.normpath(args.dir))}")
        have = remote_files(api, args.repo)
    except Exception as e:      # 鍵は例外の文に入らない (hf はヘッダーの中身を出さない)
        return done("error", f"{type(e).__name__}: {str(e).splitlines()[0] if str(e) else ''}", 6)
    if have != mine:
        print("\n".join(diff_lines(mine, have)), flush=True)
        return done("mismatch", f"上げたあとのリポジトリの一覧が手元と合わない", 5)
    return done("ok", f"{args.repo} に {len(mine)} ファイル {total:,} バイトを上げ, 名前・バイト数が全部合うのを確かめた", 0)


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