from egogym.wrappers.episode_monitor import EpisodeMonitor
from egogym.wrappers.unprivileged_molmo import UnprivilegedMolmo
from egogym.wrappers.unprivileged_gemini import UnprivilegedGemini
from egogym.wrappers.unprivileged_moondream import UnprivilegedMoondream
from egogym.wrappers.unprivileged_chatgpt import UnprivilegedChatGPT

VLM_WRAPPERS = {
    "molmo": UnprivilegedMolmo,
    "gemini": UnprivilegedGemini,
    "moondream": UnprivilegedMoondream,
    "chatgpt": UnprivilegedChatGPT,
}


def get_vlm_wrapper(vlm_name: str):
    if vlm_name not in VLM_WRAPPERS:
        available = ", ".join(f"'{k}'" for k in VLM_WRAPPERS.keys())
        raise ValueError(f"Unknown VLM option: '{vlm_name}'. Choose from {available}.")
    return VLM_WRAPPERS[vlm_name]

