OBLITERATUS/obliteratus/checkpoint_capabilities.py

"""Closed, inert checkpoint capability and dependency diagnostics.

This module resolves declarative capability metadata only.  It never imports an
optional producer framework, discovers plugins, reads checkpoint payloads, or
authorizes a trusted operation.
"""

from __future__ import annotations

import json
import re
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from hashlib import sha256
from importlib import metadata
from typing import Any

from obliteratus import __version__


_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
_VERSION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.+!-]*$")
_TRUST_REQUIRED_FORMATS = frozenset(
    {
        "hf_pytorch_pickle",
        "pytorch_dcp",
        "fsdp_state_dict",
        "megatron_torch_dist",
        "megatron_torch_dcp",
        "megatron_fsdp_dtensor",
        "deepspeed_zero",
        "deepspeed_universal",
    }
)


def _require_identifier(value: str, field: str) -> None:
    if (
        not isinstance(value, str)
        or len(value) > 64
        or _IDENTIFIER.fullmatch(value) is None
    ):
        raise ValueError(f"{field} must be a portable package identifier")


def _require_version(value: str, field: str) -> None:
    if (
        not isinstance(value, str)
        or len(value) > 64
        or not any(character.isdigit() for character in value)
        or _VERSION.fullmatch(value) is None
    ):
        raise ValueError(f"{field} must be an exact portable version")


@dataclass(frozen=True, order=True)
class ExactDependency:
    """One distribution and the only version accepted by a capability."""

    distribution: str
    version: str

    def __post_init__(self) -> None:
        _require_identifier(self.distribution, "distribution")
        _require_version(self.version, "version")

    @property
    def requirement(self) -> str:
        return f"{self.distribution}=={self.version}"


@dataclass(frozen=True)
class AdapterCapability:
    """Inert metadata for one future exact adapter contract.

    Possessing this record does not provide an adapter, reader, trust profile,
    or authorization.  Product registrations remain empty until a separately
    approved adapter supplies an exact record.
    """

    adapter_id: str
    adapter_version: str
    producer: str
    producer_version: str
    formats: tuple[str, ...]
    required_extras: tuple[str, ...]
    required_dependencies: tuple[ExactDependency, ...]

    def __post_init__(self) -> None:
        if any(
            type(value) is not tuple
            for value in (self.formats, self.required_extras, self.required_dependencies)
        ):
            raise TypeError("capability collections must be immutable tuples")
        _require_identifier(self.adapter_id, "adapter_id")
        _require_version(self.adapter_version, "adapter_version")
        _require_identifier(self.producer, "producer")
        _require_version(self.producer_version, "producer_version")
        if not self.formats:
            raise ValueError("formats must not be empty")
        if not self.required_extras:
            raise ValueError("required_extras must not be empty")
        if not self.required_dependencies:
            raise ValueError("required_dependencies must not be empty")
        if len(self.formats) > 8:
            raise ValueError("formats exceeds the bounded capability limit")
        if len(self.required_extras) > 4:
            raise ValueError("required_extras exceeds the bounded capability limit")
        if len(self.required_dependencies) > 4:
            raise ValueError("required_dependencies exceeds the bounded capability limit")
        for checkpoint_format in self.formats:
            _require_identifier(checkpoint_format, "format")
            if checkpoint_format not in _TRUST_REQUIRED_FORMATS:
                raise ValueError("capabilities may target only trust-required formats")
        for extra in self.required_extras:
            _require_identifier(extra, "required_extra")
        if len(set(self.formats)) != len(self.formats):
            raise ValueError("formats must be unique")
        if len(set(self.required_extras)) != len(self.required_extras):
            raise ValueError("required_extras must be unique")
        if not all(isinstance(item, ExactDependency) for item in self.required_dependencies):
            raise TypeError("required_dependencies must contain ExactDependency records")
        distributions = [item.distribution for item in self.required_dependencies]
        if len(set(distributions)) != len(distributions):
            raise ValueError("required dependency distributions must be unique")

    @property
    def capability_digest(self) -> str:
        record = {
            "adapter_id": self.adapter_id,
            "adapter_version": self.adapter_version,
            "producer": self.producer,
            "producer_version": self.producer_version,
            "formats": sorted(self.formats),
            "required_extras": sorted(self.required_extras),
            "required_dependencies": [
                {"distribution": item.distribution, "version": item.version}
                for item in sorted(self.required_dependencies)
            ],
        }
        payload = json.dumps(record, sort_keys=True, separators=(",", ":")).encode()
        return f"sha256:{sha256(payload).hexdigest()}"


@dataclass(frozen=True)
class AdapterResolution:
    """Descriptor-compatible result of an inert registry lookup."""

    status: str
    adapter_id: str | None
    adapter_version: str | None
    capability_digest: str | None
    reason: str
    dependency_unavailable: bool = False

    def to_dict(self) -> dict[str, Any]:
        return {
            "status": self.status,
            "adapter_id": self.adapter_id,
            "adapter_version": self.adapter_version,
            "capability_digest": self.capability_digest,
            "reason": self.reason,
        }


def _installed_version(distribution: str) -> str | None:
    try:
        return metadata.version(distribution)
    except metadata.PackageNotFoundError:
        return None


def _render_observed(value: str | None) -> str:
    if value is None:
        return "<missing>"
    if (
        not isinstance(value, str)
        or len(value) > 64
        or not any(character.isdigit() for character in value)
        or _VERSION.fullmatch(value) is None
    ):
        return "<invalid>"
    return value


@dataclass(frozen=True)
class AdapterRegistry:
    """A closed table; resolution performs no plugin or package imports."""

    capabilities: tuple[AdapterCapability, ...] = ()

    def __post_init__(self) -> None:
        if type(self.capabilities) is not tuple:
            raise TypeError("capabilities must be an immutable tuple")
        if not all(isinstance(item, AdapterCapability) for item in self.capabilities):
            raise TypeError("capabilities must contain AdapterCapability records")
        if len(self.capabilities) > 16:
            raise ValueError("adapter registry exceeds the bounded capability limit")
        identifiers = [item.adapter_id for item in self.capabilities]
        if len(set(identifiers)) != len(identifiers):
            raise ValueError("adapter capability identifiers must be unique")

    def resolve(
        self,
        checkpoint_format: str,
        *,
        version_provider: Callable[[str], str | None] = _installed_version,
        project_name: str = "obliteratus",
        project_version: str = __version__,
        observed_producer: str | None = None,
        observed_producer_version: str | None = None,
    ) -> AdapterResolution:
        """Resolve one format and report exact dependency state without importing it."""
        _require_identifier(checkpoint_format, "checkpoint_format")
        _require_identifier(project_name, "project_name")
        _require_version(project_version, "project_version")
        if (observed_producer is None) != (observed_producer_version is None):
            raise ValueError("observed producer and version must be provided together")
        if observed_producer is not None and observed_producer_version is not None:
            _require_identifier(observed_producer, "observed_producer")
            _require_version(observed_producer_version, "observed_producer_version")
        format_matches = tuple(
            sorted(
                (
                    capability
                    for capability in self.capabilities
                    if checkpoint_format in capability.formats
                ),
                key=lambda capability: capability.adapter_id,
            )
        )
        if not format_matches:
            return AdapterResolution(
                status="missing",
                adapter_id=None,
                adapter_version=None,
                capability_digest=None,
                reason=(
                    "No approved exact capability is registered; structural recognition "
                    "does not select an adapter or dependency set."
                ),
            )
        source_identity_verified = observed_producer is not None
        matches = (
            tuple(
                capability
                for capability in format_matches
                if capability.producer == observed_producer
                and capability.producer_version == observed_producer_version
            )
            if source_identity_verified
            else format_matches
        )
        if not matches:
            return AdapterResolution(
                status="missing",
                adapter_id=None,
                adapter_version=None,
                capability_digest=None,
                reason=(
                    "No exact capability matches the observed producer/version; "
                    "do not install or select a format-only candidate."
                ),
            )
        if len(matches) > 1:
            identifiers = ",".join(item.adapter_id for item in matches)
            return AdapterResolution(
                status="ambiguous",
                adapter_id=None,
                adapter_version=None,
                capability_digest=None,
                reason=f"Multiple exact capabilities match: {identifiers}.",
            )

        capability = matches[0]
        observed = {
            dependency.distribution: version_provider(dependency.distribution)
            for dependency in sorted(capability.required_dependencies)
        }
        unavailable = [
            dependency
            for dependency in sorted(capability.required_dependencies)
            if observed[dependency.distribution] != dependency.version
        ]
        if unavailable:
            extras = ",".join(sorted(capability.required_extras))
            install_extra = f"{project_name}[{extras}]=={project_version}"
            requirements = ",".join(
                item.requirement for item in sorted(capability.required_dependencies)
            )
            observed_versions = ",".join(
                f"{item.distribution}={_render_observed(observed[item.distribution])}"
                for item in sorted(capability.required_dependencies)
            )
            return AdapterResolution(
                status="missing",
                adapter_id=capability.adapter_id,
                adapter_version=capability.adapter_version,
                capability_digest=capability.capability_digest,
                reason=(
                    (
                        "source_identity=verified; "
                        if source_identity_verified
                        else "source_identity=unverified; "
                    )
                    + "dependency_status=missing_or_incompatible; "
                    f"install_extra={install_extra}; required_versions={requirements}; "
                    f"observed_versions={observed_versions}"
                ),
                dependency_unavailable=True,
            )
        if not source_identity_verified:
            return AdapterResolution(
                status="missing",
                adapter_id=capability.adapter_id,
                adapter_version=capability.adapter_version,
                capability_digest=capability.capability_digest,
                reason=(
                    "source_identity=unverified; exact producer/version evidence is "
                    "required before this format-only capability candidate can match."
                ),
            )
        return AdapterResolution(
            status="matched",
            adapter_id=capability.adapter_id,
            adapter_version=capability.adapter_version,
            capability_digest=capability.capability_digest,
            reason="Exact capability dependencies are present; trust authorization is still required.",
        )


def registry_from(capabilities: Iterable[AdapterCapability]) -> AdapterRegistry:
    """Construct a deterministic closed registry from explicit records."""
    return AdapterRegistry(tuple(sorted(capabilities, key=lambda item: item.adapter_id)))