"""Verify the source/artifacts and generate the public manifest and checksums.

License: CC-BY-NC-SA-4.0. Run after conversion/validation and documentation changes.
"""
import gc
import hashlib
import importlib.metadata
import json
import platform
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path

import onnx

ROOT = Path(__file__).resolve().parents[1]


def sha256(path):
    with path.open("rb") as f:
        return hashlib.file_digest(f, "sha256").hexdigest()


def main():
    source = json.loads((ROOT / "upstream/metadata.json").read_text())
    for info in source["files"]:
        name = info["name"]
        target = ROOT / ("upstream/README.md" if name == "README.md" else name)
        if name == ".gitattributes":
            continue  # Our generated model files need their own LFS rules.
        if sha256(target) != info["sha256"]:
            raise RuntimeError("Original source file changed: " + name)
    variants = []
    for key, filename in [("fp32", "model.onnx"), ("fp16", "model_fp16.onnx"), ("int8", "model_int8.onnx")]:
        path = ROOT / "onnx" / filename
        onnx.checker.check_model(str(path), full_check=True)
        model = onnx.load(path)
        report = json.loads((ROOT / "validation" / (key + ".json")).read_text())
        if report["model_sha256"] != sha256(path) or report["reference_npz_sha256"] != sha256(ROOT / "validation/reference.npz") or report["fixtures_sha256"] != sha256(ROOT / "validation/fixtures.json"):
            raise RuntimeError("Stale validation report: " + key)
        if key != "int8" and not report["passed"]:
            raise RuntimeError("Required numerical check failed: " + key)
        variants.append({
            "id": key, "path": "onnx/" + filename,
            "recommended_default": key == "fp32", "auto_select": key == "fp32",
            "status": "experimental-parity-failed" if not report["passed"] else "numerical-smoke-tests-passed",
            "size_bytes": path.stat().st_size, "sha256": sha256(path),
            "opsets": {v.domain or "ai.onnx": v.version for v in model.opset_import},
            "ir_version": model.ir_version,
            "inputs": [{"name": v.name, "element_type": onnx.TensorProto.DataType.Name(v.type.tensor_type.elem_type),
                        "shape": [d.dim_param or d.dim_value for d in v.type.tensor_type.shape.dim]} for v in model.graph.input],
            "outputs": [{"name": v.name, "element_type": onnx.TensorProto.DataType.Name(v.type.tensor_type.elem_type),
                         "shape": [d.dim_param or d.dim_value for d in v.type.tensor_type.shape.dim]} for v in model.graph.output],
            "operators": dict(sorted(Counter((n.domain + ":" if n.domain else "") + n.op_type for n in model.graph.node).items())),
            "external_tensor_files": [], "validated_provider": report["provider"],
            "numerical_check_passed": report["passed"], "validation_report": "validation/" + key + ".json",
            "accuracy_evaluated": False, "accelerator_execution_tested": False,
        })
        if any(v.data_location == onnx.TensorProto.EXTERNAL for v in model.graph.initializer):
            raise RuntimeError("Unexpected external tensor data")
        del model
        gc.collect()
    versions = {name: importlib.metadata.version(name) for name in [
        "torch", "transformers", "tokenizers", "huggingface-hub", "safetensors", "numpy",
        "onnx", "onnxruntime", "onnxconverter-common", "protobuf", "ml-dtypes"]}
    manifest = {
        "schema_version": 1, "repository": "CoderBak/editlens_roberta_modelkit",
        "license": "CC-BY-NC-SA-4.0", "public": True, "gated": False,
        "source_repository": source["repo_id"], "source_revision": source["revision"],
        "original_weights": {"path": "model.safetensors", "precision": "float32",
                             "size_bytes": (ROOT / "model.safetensors").stat().st_size,
                             "sha256": sha256(ROOT / "model.safetensors"), "unchanged_from_upstream": True},
        "generated_utc": datetime.now(timezone.utc).isoformat(),
        "build_environment": {"python": platform.python_version(), "os": platform.system(),
                              "os_version": platform.mac_ver()[0], "architecture": platform.machine(), "packages": versions},
        "maximum_sequence_tokens_including_special_tokens": 512, "num_labels": 4,
        "default_variant": "fp32", "variants": variants,
        "limitations": ["Numerical conversion checks only; no labeled accuracy benchmark.",
                        "INT8 is experimental and failed the documented numerical acceptance gate.",
                        "No cross-platform or accelerated-provider compatibility certification.",
                        "Checksums establish file integrity; they are not an independent publisher signature."],
    }
    (ROOT / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    files = sorted(p for p in ROOT.rglob("*") if p.is_file() and p.name != "SHA256SUMS"
                   and not any(part in {"__pycache__", ".git", ".cache", ".venv"} for part in p.relative_to(ROOT).parts)
                   and p.name != ".DS_Store" and p.suffix != ".pyc")
    (ROOT / "SHA256SUMS").write_text("".join(sha256(p) + "  " + p.relative_to(ROOT).as_posix() + "\n" for p in files))
    print(json.dumps({"verified_source": source["revision"], "published_files": len(files) + 1,
                      "variants": [{"id": v["id"], "bytes": v["size_bytes"], "status": v["status"]} for v in variants]}, indent=2))


if __name__ == "__main__":
    main()