from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig

from .dalle2_transformation import DallE2ImageEditConfig
from .transformation import OpenAIImageEditConfig

__all__ = ["OpenAIImageEditConfig", "DallE2ImageEditConfig", "get_openai_image_edit_config"]


def get_openai_image_edit_config(model: str) -> BaseImageEditConfig:
    """
    Get the appropriate OpenAI image edit config based on the model.
    
    Args:
        model: The model name (e.g., "dall-e-2", "gpt-image-1")
        
    Returns:
        The appropriate config instance for the model
    """
    model_normalized = model.lower().replace("-", "").replace("_", "")
    
    if model_normalized == "dalle2":
        return DallE2ImageEditConfig()
    else:
        # Default to standard OpenAI config for gpt-image-1 and other models
        return OpenAIImageEditConfig()