#!/usr/bin/env python3
"""test_hf_upload_ckpt.py — Pod で最後のチェックポイントを Hugging Face へ上げる hf_upload_ckpt.py (依頼 49) の確かめ.

  python3 pi05/test_hf_upload_ckpt.py -v

Hugging Face には触らない (HfApi を偽物にする). 一時フォルダにチェックポイントの形を作って流す.
"""
import io
import sys
import tempfile
import unittest
from contextlib import redirect_stdout
from pathlib import Path
from types import SimpleNamespace

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE / "openpi_addon"))

import hf_upload_ckpt as hu  # noqa: E402

REPO = "PixelstoMotors/pi05-mecharm-trim-corr-lora-v1"


class FakeApi:
    """リポジトリを {パス: バイト数} で持つ偽物. upload_folder は ignore_patterns の train_state を除いて写す"""

    def __init__(self, files=None, private=True, upload_error=None, drop=None):
        self.files = dict(files or {})
        self.private = private
        self.upload_error = upload_error
        self.drop = drop            # 上げたあと, これを落とす (一覧が合わない場合)
        self.calls = []

    def create_repo(self, repo, repo_type, private, exist_ok):
        self.calls.append(("create_repo", repo, private, exist_ok))

    def model_info(self, repo):
        return SimpleNamespace(private=self.private)

    def list_repo_tree(self, repo, recursive, repo_type):
        out = [SimpleNamespace(path=".gitattributes", size=1519)]
        dirs = {p.rsplit("/", 1)[0] for p in self.files if "/" in p}
        out += [SimpleNamespace(path=d) for d in sorted(dirs)]          # フォルダには size がない
        out += [SimpleNamespace(path=p, size=s) for p, s in self.files.items()]
        return out

    def upload_folder(self, folder_path, repo_id, repo_type, ignore_patterns, commit_message):
        self.calls.append(("upload_folder", repo_id, tuple(ignore_patterns), commit_message))
        if self.upload_error:
            raise self.upload_error
        for p, s in hu.local_files(folder_path):
            self.files[p] = s
        if self.drop:
            self.files.pop(self.drop, None)


def make_ckpt(root):
    d = Path(root) / "1499"
    for p, n in [("_CHECKPOINT_METADATA", 426), ("params/_METADATA", 300), ("params/d/abc", 1000),
                 ("assets/local/x/norm_stats.json", 50), ("train_state/opt", 9999)]:
        (d / p).parent.mkdir(parents=True, exist_ok=True)
        (d / p).write_bytes(b"\0" * n)
    return d


def run(api, d):
    buf = io.StringIO()
    with redirect_stdout(buf):
        rc = hu.main(["--dir", str(d), "--repo", REPO, "--message", "m"], api=api)
    return rc, buf.getvalue().strip().splitlines()[-1]


class T(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.d = make_ckpt(self.tmp.name)

    def tearDown(self):
        self.tmp.cleanup()

    def test_local_files_skip_train_state(self):
        names = [p for p, _ in hu.local_files(self.d)]
        self.assertEqual(names, ["_CHECKPOINT_METADATA", "assets/local/x/norm_stats.json", "params/_METADATA", "params/d/abc"])

    def test_ok(self):
        api = FakeApi()
        rc, last = run(api, self.d)
        self.assertEqual(rc, 0)
        self.assertTrue(last.startswith("HFUP ok "), last)
        self.assertEqual(api.calls[0], ("create_repo", REPO, True, True))     # 非公開で作る
        self.assertIn("train_state/**", api.calls[1][2])
        self.assertNotIn("train_state/opt", api.files)

    def test_not_private_stops_before_upload(self):
        api = FakeApi(private=False)
        rc, last = run(api, self.d)
        self.assertEqual((rc, last.split()[1]), (3, "not_private"))
        self.assertFalse([c for c in api.calls if c[0] == "upload_folder"])

    def test_same_content_skips(self):
        api = FakeApi(files=dict(hu.local_files(self.d)))
        rc, last = run(api, self.d)
        self.assertEqual((rc, last.split()[1]), (0, "same"))
        self.assertFalse([c for c in api.calls if c[0] == "upload_folder"])

    def test_other_content_stops(self):
        api = FakeApi(files={"params/old": 5})
        rc, last = run(api, self.d)
        self.assertEqual((rc, last.split()[1]), (4, "other_content"))
        self.assertEqual(api.files, {"params/old": 5})                       # 上書きしない

    def test_upload_error(self):
        rc, last = run(FakeApi(upload_error=RuntimeError("503 Server Error\nmore")), self.d)
        self.assertEqual((rc, last.split()[1]), (6, "error"))
        self.assertIn("503 Server Error", last)

    def test_mismatch_after_upload(self):
        rc, last = run(FakeApi(drop="params/d/abc"), self.d)
        self.assertEqual((rc, last.split()[1]), (5, "mismatch"))

    def test_unreadable_dir(self):
        locked = self.d / "params" / "locked"
        locked.mkdir()
        locked.chmod(0)
        try:
            rc, last = run(FakeApi(), self.d)
        finally:
            locked.chmod(0o755)
        self.assertEqual((rc, last.split()[1]), (6, "error"))

    def test_no_metadata(self):
        (self.d / "_CHECKPOINT_METADATA").unlink()
        rc, last = run(FakeApi(), self.d)
        self.assertEqual((rc, last.split()[1]), (6, "error"))


if __name__ == "__main__":
    unittest.main()