"""Verify the three-part Cyrillic FP8 release against Hub metadata."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from huggingface_hub import HfApi
from huggingface_hub.hf_api import RepoFile


def lfs_sha256(file: RepoFile) -> str | None:
    if file.lfs is None:
        return None
    if isinstance(file.lfs, dict):
        return file.lfs.get("sha256")
    return getattr(file.lfs, "sha256", None)


def remote_files(api: HfApi, repo_id: str) -> dict[str, RepoFile]:
    return {
        item.path: item
        for item in api.list_repo_tree(repo_id, repo_type="model", recursive=True, expand=True)
        if isinstance(item, RepoFile)
    }


def run(args: argparse.Namespace) -> dict[str, object]:
    api = HfApi()
    transformer = remote_files(api, args.transformer_repo)
    controlnet = remote_files(api, args.controlnet_repo)
    config = remote_files(api, args.config_repo)

    quant = json.loads(args.quantization_manifest.read_text(encoding="utf-8"))
    expected_transformer = {
        row["file"]: (int(row["output_bytes"]), row["output_sha256"])
        for row in quant["shards"]
    }
    transformer_mismatches = []
    for path, (size, sha256) in expected_transformer.items():
        remote = transformer.get(path)
        if remote is None or remote.size != size or lfs_sha256(remote) != sha256:
            transformer_mismatches.append(path)

    control = json.loads(args.controlnet_manifest.read_text(encoding="utf-8"))
    remote_control = controlnet.get(control["file"])
    controlnet_match = (
        remote_control is not None
        and remote_control.size == control["bytes"]
        and lfs_sha256(remote_control) == control["sha256"]
    )
    required_config = {
        "README.md",
        "training-config.json",
        "runtime-config.json",
        "benchmark-config.json",
        "benchmark-results.json",
        "quantization-manifest.json",
        "controlnet_benchmark.py",
        "controlnet_accuracy.py",
        "verify_hf_release.py",
    }
    missing_config = sorted(required_config - config.keys())
    report = {
        "transformer": {
            "repo": args.transformer_repo,
            "remote_files": len(transformer),
            "verified_shards": len(expected_transformer) - len(transformer_mismatches),
            "mismatches": transformer_mismatches,
        },
        "controlnet": {
            "repo": args.controlnet_repo,
            "remote_files": len(controlnet),
            "model_match": controlnet_match,
        },
        "config": {
            "repo": args.config_repo,
            "remote_files": len(config),
            "missing_required": missing_config,
        },
    }
    if transformer_mismatches or not controlnet_match or missing_config:
        raise RuntimeError(json.dumps(report, ensure_ascii=False))
    return report


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--transformer-repo", required=True)
    parser.add_argument("--controlnet-repo", required=True)
    parser.add_argument("--config-repo", required=True)
    parser.add_argument("--quantization-manifest", type=Path, required=True)
    parser.add_argument("--controlnet-manifest", type=Path, required=True)
    return parser.parse_args()


if __name__ == "__main__":
    print(json.dumps(run(parse_args()), ensure_ascii=False))