from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.vertex_ai.common_utils import (
    VertexAIModelRoute,
    get_vertex_ai_model_route,
)

from .cost_calculator import cost_calculator
from .vertex_gemini_transformation import VertexAIGeminiImageEditConfig
from .vertex_imagen_transformation import VertexAIImagenImageEditConfig

__all__ = [
    "VertexAIGeminiImageEditConfig",
    "VertexAIImagenImageEditConfig",
    "get_vertex_ai_image_edit_config",
    "cost_calculator",
]


def get_vertex_ai_image_edit_config(model: str) -> BaseImageEditConfig:
    """
    Get the appropriate image edit config for a Vertex AI model.

    Routes to the correct transformation class based on the model type:
    - Gemini models use generateContent API (VertexAIGeminiImageEditConfig)
    - Imagen models use predict API (VertexAIImagenImageEditConfig)

    Args:
        model: The model name (e.g., "gemini-2.5-flash", "imagegeneration@006")

    Returns:
        BaseImageEditConfig: The appropriate configuration class
    """
    # Determine the model route
    model_route = get_vertex_ai_model_route(model)

    if model_route == VertexAIModelRoute.GEMINI:
        # Gemini models use generateContent API
        return VertexAIGeminiImageEditConfig()
    else:
        # Default to Imagen for other models (imagegeneration, etc.)
        # This includes NON_GEMINI models like imagegeneration@006
        return VertexAIImagenImageEditConfig()
