"""Shared env registry and builder used by train.py and viz.py."""

from omegaconf import DictConfig, OmegaConf

from src.core.env import BaseEnv, BaseEnvConfig
from src.envs.rotary_cartpole import RotaryCartPoleEnv, RotaryCartPoleConfig

# Maps Hydra config-group name → (EnvClass, ConfigClass)
ENV_REGISTRY: dict[str, tuple[type[BaseEnv], type[BaseEnvConfig]]] = {
    "rotary_cartpole": (RotaryCartPoleEnv, RotaryCartPoleConfig),
}


def build_env(env_name: str, cfg: DictConfig) -> BaseEnv:
    """Instantiate the right env + config from the Hydra config-group name."""
    if env_name not in ENV_REGISTRY:
        raise ValueError(f"Unknown env '{env_name}'. Registered: {list(ENV_REGISTRY)}")

    env_cls, config_cls = ENV_REGISTRY[env_name]
    env_dict = OmegaConf.to_container(cfg.env, resolve=True)
    env_dict.pop("hpo", None)  # HPO range metadata — not an env config field
    return env_cls(config_cls(**env_dict))
