# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
"""Content-bound loading of the published mixed MXFP8 checkpoint."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

from safetensors import safe_open

from ax_doclayout_mlx._identity import EXPECTED_FILES, EXPECTED_TENSORS, MODEL_ID
from ax_doclayout_mlx.errors import ArtifactError


def validate_bundle(directory: Path) -> dict:
    """Verify every required file, then the exact safetensors inventory."""
    try:
        root = Path(directory).expanduser()
        if not root.is_dir() or root.is_symlink():
            raise ArtifactError("regular_bundle_directory_required")
        for name, expected in EXPECTED_FILES.items():
            path = root / name
            if path.is_symlink() or not path.is_file():
                raise ArtifactError("regular_bundle_files_required")
            with path.open("rb") as stream:
                actual = hashlib.file_digest(stream, "sha256").hexdigest()
            if actual != expected:
                raise ArtifactError("bundle_checksum_mismatch")
        manifest = json.loads((root / "manifest.json").read_text())
        if (manifest["model_id"] != MODEL_ID or manifest["mode"] != "mxfp8"
                or manifest["group_size"] != 32
                or manifest["convolution_precision"] != "float16"
                or manifest["convolutions_quantized"] is not False):
            raise ArtifactError("bundle_precision_contract_mismatch")
        with safe_open(root / "model.safetensors", framework="np") as weights:
            if set(weights.keys()) != set(EXPECTED_TENSORS):
                raise ArtifactError("checkpoint_inventory_mismatch")
            for name, expected in EXPECTED_TENSORS.items():
                item = weights.get_slice(name)
                if (list(item.get_shape()) != expected["shape"]
                        or item.get_dtype() != expected["dtype"]):
                    raise ArtifactError("checkpoint_tensor_contract_mismatch")
        return manifest
    except ArtifactError:
        raise
    except Exception:  # noqa: BLE001 - redact third-party paths and details
        raise ArtifactError("bundle_validation_failed") from None