"""Plugin registry for any4d data libraries.

External packages register Dataset classes used by experiment configs without
editing any4d source:

    from custom.dataset.registry import register_data_library
    from my_pkg.dataset import MyDataset

    register_data_library('my_lib', MyDataset)

Then any4d experiments can set `any4d_config.data_library='my_lib'` to use it.

The registry is a global singleton keyed by string name. Registering twice
with the same name overwrites. Built-in defaults (`anydata`, `vidar`) are
loaded lazily by `custom/experiment/template.py` so this module stays
import-light.
"""

from typing import Type

_DATA_LIBRARIES: dict[str, Type] = {}


def register_data_library(name: str, dataset_cls: Type) -> None:
    """Register a Dataset class under `name`.

    The class must accept the AnyDataset-style constructor kwargs that
    `custom/experiment/template.py` passes (dataset_dir, phase, data_overrides,
    data_defaults, shuffle).
    """
    _DATA_LIBRARIES[name] = dataset_cls


def get_data_library(name: str) -> Type | None:
    """Look up a registered Dataset class. Returns None if not registered."""
    return _DATA_LIBRARIES.get(name)