"""Helpers for importing original Ctrl-World modules without modifying that repo."""

from __future__ import annotations

import sys
from pathlib import Path


def ensure_ctrl_world_on_path(ctrl_world_root: str) -> Path:
    root = Path(ctrl_world_root).resolve()
    if not root.exists():
        raise FileNotFoundError(f"Ctrl-World root not found: {root}")
    root_str = str(root)
    if root_str not in sys.path:
        sys.path.insert(0, root_str)
    return root


def import_original_modules(ctrl_world_root: str):
    ensure_ctrl_world_on_path(ctrl_world_root)

    from models.pipeline_stable_video_diffusion import StableVideoDiffusionPipeline
    from models.pipeline_ctrl_world import CtrlWorldDiffusionPipeline
    from models.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel

    return {
        "StableVideoDiffusionPipeline": StableVideoDiffusionPipeline,
        "CtrlWorldDiffusionPipeline": CtrlWorldDiffusionPipeline,
        "UNetSpatioTemporalConditionModel": UNetSpatioTemporalConditionModel,
    }