import sys

from .detr import DETRModelWrapper
from .fcos import FCOSModelWrapper
from .faster_rcnn import FasterRCNNModelWrapper
from .dino import DINOModelWrapper
from .yolo import YOLOModelWrapper

def build_model(model_name, dataset, config_path, device, save_path, distributed=False, local_rank=0):
    """
    Build the model wrapper based on the model name.

    Args:
        model_name (str): Name of the model to build.
        dataset (str): Name of the dataset.
        config_path (str): Path to the configuration file.

    Returns:
        BaseModelWrapper: An instance of the appropriate model wrapper.
    """

    if model_name == "detr":
        return DETRModelWrapper(dataset, config_path, device, save_path, distributed, local_rank)
    elif model_name == "fcos":
        return FCOSModelWrapper(dataset, config_path, device, save_path, distributed, local_rank)
    elif model_name == "faster_rcnn":
        return FasterRCNNModelWrapper(dataset, config_path, device, save_path, distributed, local_rank)
    elif model_name == "dino":
        return DINOModelWrapper(dataset, config_path, device, save_path, distributed, local_rank)
    elif model_name == "yolo":
        return YOLOModelWrapper(dataset, config_path, device, save_path, distributed, local_rank)
    else:
        raise ValueError(f"Unknown model name: {model_name}")
