from pathlib import Path
from typing import Callable, Union

import torch
import torch.nn as nn

# Maps backbone name -> builder function; populated by each model module at import time
_REGISTRY: dict[str, Callable[[dict], nn.Module]] = {}


# Called by each model module to advertise its backbone(s) to get_model
def register(name: str, builder: Callable[[dict], nn.Module]) -> None:
    _REGISTRY[name] = builder


# Instantiates the backbone requested in cfg["backbone"]
def get_model(cfg: dict) -> nn.Module:
    backbone = cfg.get("backbone", "simple_cnn")
    builder = _REGISTRY.get(backbone)
    if builder is None:
        available = ", ".join(sorted(_REGISTRY))
        raise ValueError(f"Unknown backbone: {backbone!r}. Available: {available}")
    return builder(cfg)


# Loads a saved state-dict into model in-place and returns it
def load_checkpoint(model: nn.Module, path: Union[Path, str], device) -> nn.Module:
    model.load_state_dict(torch.load(path, map_location=device, weights_only=True))
    return model


# Importing the modules triggers their register() calls
from src.models import simple_cnn, resnet, efficientnet  # noqa: E402, F401

__all__ = ["get_model", "load_checkpoint", "register"]
