# bclm/registry.py
"""
Architecture registry.

Maps the ``"architecture"`` string found in a model's ``config.json`` to
the concrete classes needed to instantiate and run inference.  New model
families only need to call ``register_architecture`` — every downstream 
API picks them up automatically.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Optional, Type


# =============================================================================
# Registry entry
# =============================================================================

@dataclass(frozen=True)
class ArchitectureEntry:
    """Bundle of classes that together define a loadable architecture."""

    model_class: Type          # nn.Module subclass (e.g. BCLM1Model)
    config_class: Type         # Dataclass with a ``from_dict`` classmethod
    inference_class: Type      # Generation wrapper (e.g. BCLM1ForGeneration)


# =============================================================================
# Internal state
# =============================================================================

_REGISTRY: Dict[str, ArchitectureEntry] = {}
_BUILTINS_LOADED: bool = False


# =============================================================================
# Public API
# =============================================================================

def register_architecture(
    name: str,
    *,
    model_class: Type,
    config_class: Type,
    inference_class: Type,
) -> None:
    """
    Register an architecture so it can be loaded via ``bclm.load()``.

    Parameters
    ----------
    name : str
        Value that appears in ``config.json`` under ``"architecture"``.
    model_class : type
        The ``nn.Module`` subclass that implements the model.
    config_class : type
        A dataclass (or similar) with a ``from_dict(dict) -> config``
        classmethod and a ``to_dict() -> dict`` method.
    inference_class : type
        A wrapper class whose constructor accepts the raw model and
        exposes ``generate``, ``generate_stream``, ``prefill``, and
        ``decode_one_token``.
    """
    _REGISTRY[name] = ArchitectureEntry(
        model_class=model_class,
        config_class=config_class,
        inference_class=inference_class,
    )


def get_architecture(name: str) -> ArchitectureEntry:
    """
    Look up a registered architecture by name.

    Built-in architectures are registered lazily on first access.

    Raises
    ------
    ValueError
        If *name* is not found in the registry.
    """
    _ensure_builtins()

    entry = _REGISTRY.get(name)
    if entry is None:
        available = sorted(_REGISTRY.keys()) or ["(none)"]
        raise ValueError(
            f"Unknown architecture {name!r}. "
            f"Registered architectures: {', '.join(available)}"
        )
    return entry


def list_architectures() -> list[str]:
    """Return the names of all registered architectures."""
    _ensure_builtins()
    return sorted(_REGISTRY.keys())


# =============================================================================
# Built-in registration (lazy to avoid circular imports)
# =============================================================================

def _ensure_builtins() -> None:
    """Import and register all built-in architectures once."""
    global _BUILTINS_LOADED
    if _BUILTINS_LOADED:
        return
    _BUILTINS_LOADED = True

    _register_bclm1()


def _register_bclm1() -> None:
    """Register the BCLM-1 architecture."""
    from .config import BCLM1ModelConfig
    from .model import BCLM1Model
    from .inference import BCLM1ForGeneration

    register_architecture(
        "BCLM1Model",
        model_class=BCLM1Model,
        config_class=BCLM1ModelConfig,
        inference_class=BCLM1ForGeneration,
    )