from typing import Callable
import torch.nn as nn

_REGISTRY: dict[str, tuple[Callable, str]] = {}


def register(name: str, builder: Callable, *, kind: str) -> None:
    _REGISTRY[name] = (builder, kind)


def get_model(cfg: dict) -> tuple:
    """Return (model_or_pair, kind).

    kind="dcgan" -> (generator, discriminator)
    """
    name = cfg.get("model")
    entry = _REGISTRY.get(name)
    if entry is None:
        available = ", ".join(sorted(_REGISTRY))
        raise ValueError(f"Unknown model: {name!r}. Available: {available}")
    builder, kind = entry
    return builder(cfg), kind


from src.models import dcgan  # noqa: E402, F401
from src.models import wgan   # noqa: E402, F401
from src.models import vae    # noqa: E402, F401
from src.models import unet   # noqa: E402, F401

