Nexum-Expanded/runtime/src/nexum_runtime/bundle.py

"""Nexum release-folder validation."""

from __future__ import annotations

import hashlib
import json
import re
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any


REQUIRED_FILES: tuple[str, ...] = (
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "chat_template.jinja",
    "generation_config.json",
    "state_config.json",
    "tensor_map.json",
)
NUMBERED_TENSOR_RE = re.compile(r"^\d{6}\.safetensors$")
SHA256_RE = re.compile(r"^[0-9a-f]{64}$")


@dataclass(frozen=True)
class BundleFile:
    name: str
    exists: bool
    bytes: int = 0
    sha256: str = ""
    keys: int | None = None


@dataclass(frozen=True)
class BundleReport:
    ok: bool
    model_dir: str
    files: tuple[BundleFile, ...]
    missing: tuple[str, ...]
    tensor_map_ok: bool
    load_map_ok: bool
    tensor_count: int
    tensor_bytes: int
    tensor_hashes_verified: bool
    tokenizer_ok: bool
    errors: tuple[str, ...]

    def to_dict(self) -> dict[str, Any]:
        return {
            "ok": self.ok,
            "model_dir": Path(self.model_dir).name,
            "files": [asdict(item) for item in self.files],
            "missing": list(self.missing),
            "tensor_map_ok": self.tensor_map_ok,
            "load_map_ok": self.load_map_ok,
            "tensor_count": self.tensor_count,
            "tensor_bytes": self.tensor_bytes,
            "tensor_hashes_verified": self.tensor_hashes_verified,
            "tokenizer_ok": self.tokenizer_ok,
            "errors": list(self.errors),
        }


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


def _count_tensor_keys(path: Path) -> int | None:
    if path.suffix != ".safetensors":
        return None
    from safetensors import safe_open

    with safe_open(  # type: ignore[no-untyped-call]
        str(path), framework="pt", device="cpu"
    ) as handle:
        return len(handle.keys())


def _read_json(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError(f"{path.name} must contain a JSON object")
    return value


def _tensor_entries(map_path: Path) -> list[dict[str, Any]]:
    data = _read_json(map_path)
    tensors = data.get("tensors")
    if not isinstance(tensors, list):
        raise ValueError("tensor_map.json must contain a tensors list")
    return [item for item in tensors if isinstance(item, dict)]


def _product_contract(release_cfg: dict[str, Any]) -> tuple[str, int, bool]:
    product = release_cfg.get("product")
    if not isinstance(product, dict):
        return "Nexum-Universal", 113, True
    variant = product.get("variant")
    expected = {
        "Nexum-Lite": (113, True),
        "Nexum-Universal": (113, True),
        "Nexum-Expanded": (113, True),
    }
    if variant not in expected:
        raise ValueError("config.json product variant is unsupported")
    package_count, complete_authority = expected[str(variant)]
    if (
        product.get("tensor_package_count") != package_count
        or product.get("complete_none_authority") is not complete_authority
    ):
        raise ValueError("config.json product profile does not match its variant")
    return str(variant), package_count, complete_authority


def release_artifact_sha256(
    model_dir: str | Path, *, validated_report: BundleReport
) -> str:
    """Return one identity for the exact validated numbered model bundle."""

    root = Path(model_dir).expanduser().resolve()
    report_root = Path(validated_report.model_dir).expanduser().resolve()
    if report_root != root:
        raise ValueError("validated bundle report belongs to another model folder")
    if not validated_report.ok or not validated_report.tensor_hashes_verified:
        raise ValueError("release identity requires a deep validated model bundle")

    files = [
        {
            "path": name,
            "bytes": (root / name).stat().st_size,
            "sha256": _sha256(root / name),
        }
        for name in REQUIRED_FILES
    ]
    tensors = []
    for entry in _tensor_entries(root / "tensor_map.json"):
        tensors.append(
            {
                "id": str(entry.get("id") or ""),
                "file": str(entry.get("file") or ""),
                "bytes": int(entry.get("bytes") or 0),
                "sha256": str(entry.get("sha256") or ""),
                "binding_sha256": str(entry.get("binding_sha256") or ""),
                "descriptor_sha256": str(entry.get("descriptor_sha256") or ""),
                "key_count": int(entry.get("key_count") or 0),
            }
        )
    canonical = json.dumps(
        {
            "schema": "nexum.release-artifact.v1",
            "files": files,
            "tensors": tensors,
        },
        ensure_ascii=True,
        separators=(",", ":"),
        sort_keys=True,
    ).encode("utf-8")
    return hashlib.sha256(canonical).hexdigest()


def validate_bundle(model_dir: str | Path, *, deep: bool = False) -> BundleReport:
    root = Path(model_dir).expanduser().resolve()
    errors: list[str] = []
    missing: list[str] = []
    files: list[BundleFile] = []
    for name in REQUIRED_FILES:
        path = root / name
        if not path.exists():
            missing.append(name)
            files.append(BundleFile(name=name, exists=False))
            continue
        try:
            files.append(
                BundleFile(
                    name=name,
                    exists=True,
                    bytes=path.stat().st_size,
                    sha256=_sha256(path),
                    keys=_count_tensor_keys(path) if deep else None,
                )
            )
        except Exception as exc:
            errors.append(f"{name}: {type(exc).__name__}: {exc}")
            files.append(BundleFile(name=name, exists=True, bytes=path.stat().st_size))

    tensor_map_ok = False
    load_map_ok = False
    tensor_count = 0
    tensor_bytes = 0
    tensor_hashes_verified = False
    tensor_dir = root / "safetensors"
    if not tensor_dir.is_dir():
        errors.append("missing safetensors directory")
    elif (root / "tensor_map.json").is_file():
        try:
            entries = _tensor_entries(root / "tensor_map.json")
            tensor_count = len(entries)
            tensor_bytes = sum(int(item.get("bytes") or 0) for item in entries)
            indexed = {
                str(item.get("file"))
                for item in entries
                if isinstance(item.get("file"), str)
            }
            actual = {
                f"safetensors/{path.name}" for path in tensor_dir.glob("*.safetensors")
            }
            bad_names = sorted(
                path.name
                for path in tensor_dir.glob("*.safetensors")
                if not NUMBERED_TENSOR_RE.fullmatch(path.name)
            )
            unknown = sorted(actual - indexed)
            absent = sorted(indexed - actual)
            if bad_names:
                errors.append(f"non-numbered safetensors files: {bad_names[:8]}")
            if unknown:
                errors.append(f"unmapped safetensors files: {unknown[:8]}")
            if absent:
                errors.append(f"missing mapped safetensors files: {absent[:8]}")
            ids = [str(item.get("id") or "") for item in entries]
            expected_ids = [f"{index:06d}" for index in range(1, len(entries) + 1)]
            if ids != expected_ids:
                errors.append(
                    "tensor ids must be contiguous, ordered, and start at 000001"
                )
            for item in entries:
                tensor_id = str(item.get("id") or "")
                expected_file = f"safetensors/{tensor_id}.safetensors"
                if item.get("file") != expected_file:
                    errors.append(f"tensor {tensor_id} must map to {expected_file}")
                    break
                path = root / expected_file
                if (
                    path.is_file()
                    and int(item.get("bytes") or -1) != path.stat().st_size
                ):
                    errors.append(
                        f"tensor {tensor_id} byte size does not match tensor_map.json"
                    )
                    break
                if not SHA256_RE.fullmatch(str(item.get("sha256") or "")):
                    errors.append(f"tensor {tensor_id} has no valid sha256")
                    break
                if not SHA256_RE.fullmatch(
                    str(item.get("binding_sha256") or "")
                ) or not SHA256_RE.fullmatch(str(item.get("descriptor_sha256") or "")):
                    errors.append(f"tensor {tensor_id} has no binding fingerprints")
                    break
                if path.is_file():
                    from safetensors import safe_open

                    with safe_open(  # type: ignore[no-untyped-call]
                        str(path), framework="pt", device="cpu"
                    ) as handle:
                        keys = list(handle.keys())
                        metadata = handle.metadata()
                    expected_keys = [
                        f"{index:06d}" for index in range(1, len(keys) + 1)
                    ]
                    if keys != expected_keys:
                        errors.append(
                            f"tensor {tensor_id} keys must be contiguous and numeric"
                        )
                        break
                    if metadata:
                        errors.append(f"tensor {tensor_id} metadata must be empty")
                        break
                    if int(item.get("key_count") or -1) != len(keys):
                        errors.append(
                            f"tensor {tensor_id} key count does not match tensor_map.json"
                        )
                        break

            map_data = _read_json(root / "tensor_map.json")
            if map_data.get("schema") != "nexum.tensor_map.v4":
                errors.append("tensor_map.json schema is incompatible")
            release_cfg = _read_json(root / "config.json")
            _variant, expected_package_count, _complete_authority = (
                _product_contract(release_cfg)
            )
            if len(entries) != expected_package_count:
                errors.append(
                    "tensor map package count does not match the product profile"
                )
            load_map_ok = (
                bool(entries)
                and ids == expected_ids
                and len(indexed) == len(entries)
                and indexed == actual
                and not bad_names
                and all(
                    item.get("file")
                    == f"safetensors/{str(item.get('id') or '')}.safetensors"
                    for item in entries
                )
            )
            if not load_map_ok:
                errors.append(
                    "ordered tensor entries must map every numbered file exactly once"
                )

            if deep and not errors:
                for item in entries:
                    tensor_id = str(item["id"])
                    path = root / str(item["file"])
                    expected_hash = str(item.get("sha256") or "")
                    actual_hash = _sha256(path)
                    if actual_hash != expected_hash:
                        errors.append(
                            f"tensor {tensor_id} sha256 does not match tensor_map.json"
                        )
                        break
                    try:
                        key_count = _count_tensor_keys(path)
                    except Exception as exc:
                        errors.append(
                            f"tensor {tensor_id}: {type(exc).__name__}: {exc}"
                        )
                        break
                    if not key_count:
                        errors.append(f"tensor {tensor_id} contains no tensors")
                        break
                tensor_hashes_verified = not errors

            if (
                not bad_names
                and not unknown
                and not absent
                and entries
                and ids == expected_ids
                and load_map_ok
            ):
                tensor_map_ok = True
        except Exception as exc:
            errors.append(f"tensor_map: {type(exc).__name__}: {exc}")

    tokenizer_ok = False
    try:
        tokenizer = _read_json(root / "tokenizer.json")
        tokenizer_cfg = _read_json(root / "tokenizer_config.json")
        generation_cfg = _read_json(root / "generation_config.json")
        release_cfg = _read_json(root / "config.json")
        state_cfg = _read_json(root / "state_config.json")
        variant, expected_package_count, complete_authority = _product_contract(
            release_cfg
        )
        architecture = release_cfg.get("runtime_config")
        if not isinstance(architecture, dict):
            raise ValueError("config.json has no runtime configuration object")
        if release_cfg.get("schema") != "nexum.release.v4":
            raise ValueError("config.json schema is incompatible")
        if release_cfg.get("model_type") != "nexum_native":
            raise ValueError("config.json model type is incompatible")
        if release_cfg.get("architectures") != ["NexumForCausalLM"]:
            raise ValueError("config.json architecture entry point is incompatible")
        required_architecture = (
            "hidden_size",
            "num_hidden_layers",
            "vocab_size",
            "max_position_embeddings",
            "bos_token_id",
            "eos_token_id",
            "pad_token_id",
        )
        if not all(
            isinstance(architecture.get(name), int) and int(architecture[name]) > 0
            for name in required_architecture
        ):
            raise ValueError("config.json architecture values are incomplete")
        token_ids = {
            "bos_token_id": int(architecture["bos_token_id"]),
            "eos_token_id": int(architecture["eos_token_id"]),
            "pad_token_id": int(architecture["pad_token_id"]),
        }
        if any(generation_cfg.get(name) != value for name, value in token_ids.items()):
            raise ValueError("generation token ids do not match config.json")
        if (
            tokenizer_cfg.get("model_max_length")
            != architecture["max_position_embeddings"]
        ):
            raise ValueError("tokenizer context length does not match config.json")
        state_bindings = state_cfg.get("bindings")
        if state_cfg.get("schema") != "nexum.state.v2" or not isinstance(
            state_bindings, list
        ):
            raise ValueError("state_config.json is incompatible")
        if complete_authority and not state_bindings:
            raise ValueError("complete product profile has no numbered state bindings")
        if not complete_authority and state_bindings:
            raise ValueError("compact product profile has unexpected state bindings")
        if int(expected_package_count) != tensor_count:
            raise ValueError("product package count does not match tensor map")
        identity = release_cfg.get("identity")
        if (
            not isinstance(identity, dict)
            or identity.get("variant") not in (None, variant)
        ):
            raise ValueError("config.json identity does not match product variant")
        model_payload = tokenizer.get("model")
        base_vocab = (
            model_payload.get("vocab") if isinstance(model_payload, dict) else None
        )
        added_tokens = tokenizer.get("added_tokens")
        if not isinstance(base_vocab, dict) or not isinstance(added_tokens, list):
            raise ValueError("tokenizer vocabulary is incomplete")
        vocabulary_ids = {
            int(value) for value in base_vocab.values() if isinstance(value, int)
        }
        vocabulary_ids.update(
            int(item["id"])
            for item in added_tokens
            if isinstance(item, dict) and isinstance(item.get("id"), int)
        )
        expected_vocab_size = int(architecture["vocab_size"])
        if vocabulary_ids != set(range(expected_vocab_size)):
            raise ValueError("tokenizer does not cover the complete model vocabulary")
        if deep:
            from tokenizers import Tokenizer  # type: ignore[import-untyped]

            native_tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
            for token_id in range(expected_vocab_size):
                if not native_tokenizer.id_to_token(
                    token_id
                ) or not native_tokenizer.decode([token_id], skip_special_tokens=False):
                    raise ValueError(f"tokenizer id {token_id} is not reversible")
        tokenizer_ok = "eos_token" in tokenizer_cfg
    except Exception as exc:
        errors.append(f"tokenizer: {type(exc).__name__}: {exc}")

    ok = not missing and not errors and tensor_map_ok and load_map_ok and tokenizer_ok
    return BundleReport(
        ok=ok,
        model_dir=str(root),
        files=tuple(files),
        missing=tuple(missing),
        tensor_map_ok=tensor_map_ok,
        load_map_ok=load_map_ok,
        tensor_count=tensor_count,
        tensor_bytes=tensor_bytes,
        tensor_hashes_verified=tensor_hashes_verified,
        tokenizer_ok=tokenizer_ok,
        errors=tuple(errors),
    )