#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path

from huggingface_hub import HfApi


SOURCE_ID = "MiniMaxAI/MiniMax-H3"
SOURCE_REVISION = "73372e6cf53e414edd3ab03e357717fb0602e758"
SOURCE_COMPONENTS = ("vae", "audio_vae")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--release", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    release = args.release.resolve()
    info = HfApi().model_info(SOURCE_ID, revision=SOURCE_REVISION, files_metadata=True)
    remote_lfs = {
        sibling.rfilename: getattr(sibling.lfs, "sha256", None)
        for sibling in info.siblings or []
        if getattr(sibling, "lfs", None) is not None
    }

    files = []
    for component in SOURCE_COMPONENTS:
        weights = sorted(
            path
            for path in (release / component).iterdir()
            if path.is_file() and path.suffix in {".safetensors", ".bin"}
        )
        if not weights:
            raise RuntimeError(f"source component has no weight files: {component}")
        for path in weights:
            relative = f"{component}/{path.name}"
            local_digest = sha256(path)
            remote_digest = remote_lfs.get(relative)
            if remote_digest != local_digest:
                raise RuntimeError(f"source-copy hash mismatch: {relative}")
            files.append(
                {
                    "file": relative,
                    "bytes": path.stat().st_size,
                    "sha256": local_digest,
                    "source_lfs_sha256": remote_digest,
                }
            )

    report = {
        "status": "pass",
        "source_model_id": SOURCE_ID,
        "source_revision": SOURCE_REVISION,
        "components": list(SOURCE_COMPONENTS),
        "verified_weight_files": files,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(report))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())