"""Per-backbone architecture registry for ``PytorchFp16Runtime`` (AZ-300).

The PyTorch baseline loads checkpoints saved as ``state_dict`` (the
``weights_only=True`` security path); ``model.load_state_dict`` requires
a pre-instantiated ``nn.Module``. This module owns the small
``model_name -> factory()`` indirection that lets the composition root
register each backbone's architecture class at startup without C7
importing component code directly (Risk-1 in the AZ-300 spec — would
violate the L2→L3 module-layout layering otherwise).

The registry is a plain ``dict``; ``register_architecture`` is
idempotent for the same factory and rejects re-registration with a
different one. The module does NOT import :mod:`torch` — keeping it
torch-free lets the composition root populate the registry even on
non-PyTorch tiers.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from torch import nn

__all__ = [
    "ArchitectureFactory",
    "default_registry",
    "register_architecture",
]

ArchitectureFactory = Callable[[], "nn.Module"]

_DEFAULT_REGISTRY: dict[str, ArchitectureFactory] = {}


def register_architecture(
    model_name: str, factory: ArchitectureFactory
) -> None:
    """Register a model-architecture factory under ``model_name``.

    Idempotent for the same ``factory`` object; raises ``ValueError`` if
    a *different* factory is registered under the same name (catches the
    common "two backbones colliding on the same stem" bug at composition
    time rather than at first ``deserialize_engine`` call).
    """
    if not model_name:
        raise ValueError("model_name must be non-empty")
    existing = _DEFAULT_REGISTRY.get(model_name)
    if existing is not None and existing is not factory:
        raise ValueError(
            f"Architecture {model_name!r} already registered with a "
            f"different factory; refusing to override"
        )
    _DEFAULT_REGISTRY[model_name] = factory


def default_registry() -> dict[str, ArchitectureFactory]:
    """Return the module-level singleton registry.

    Returned by reference — callers see live registrations. Tests that
    need an isolated registry SHOULD construct their own ``dict`` and
    pass it explicitly to :class:`PytorchFp16Runtime`.
    """
    return _DEFAULT_REGISTRY
