from typing import Callable
import torch.nn as nn

# Maps model name → (builder, kind); populated by each model module at import time
_REGISTRY: dict[str, tuple[Callable, str]] = {}


# Called by each model module to advertise its name and kind to get_model
def register(name: str, builder: Callable, *, kind: str) -> None:
    _REGISTRY[name] = (builder, kind)


# Returns (model_or_pair, kind) for the model named in cfg["model"]
# kind in {"dcgan", "wgan"} → (generator, critic)  |  kind in {"vae", "ddpm"} → model
def get_model(cfg: dict) -> tuple:
    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


# Importing each module triggers its register() calls
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

