from litellm._logging import verbose_logger
from litellm.llms.azure_ai.image_generation import AzureFoundryMAIImageGenerationConfig
from litellm.llms.base_llm.image_generation.transformation import (
    BaseImageGenerationConfig,
)

from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig
from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig
from .gpt_transformation import AzureGPTImageGenerationConfig
from .http_utils import azure_deployment_image_generation_json_body

__all__ = [
    "AzureDallE2ImageGenerationConfig",
    "AzureDallE3ImageGenerationConfig",
    "AzureGPTImageGenerationConfig",
    "azure_deployment_image_generation_json_body",
]


def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig:
    model = model.lower()
    model = model.replace("-", "")
    model = model.replace("_", "")
    if model == "" or "dalle2" in model:  # empty model is dall-e-2
        return AzureDallE2ImageGenerationConfig()
    elif "dalle3" in model:
        return AzureDallE3ImageGenerationConfig()
    elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
        return AzureFoundryMAIImageGenerationConfig()
    else:
        verbose_logger.debug(
            "Using AzureGPTImageGenerationConfig for model: %s. This follows the gpt-image model format.", model
        )
        return AzureGPTImageGenerationConfig()
