from io import BufferedReader
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast

from httpx._types import RequestFiles

import litellm
from litellm.images.utils import ImageEditRequestUtils
from litellm.types.images.main import ImageEditRequestParams
from litellm.types.llms.openai import FileTypes
from litellm.types.router import GenericLiteLLMParams

from .transformation import OpenAIImageEditConfig

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj

    LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
    LiteLLMLoggingObj = Any


class DallE2ImageEditConfig(OpenAIImageEditConfig):
    """
    DALL-E-2 specific configuration for image edit API.
    
    DALL-E-2 only supports editing a single image (not an array).
    Uses "image" field name instead of "image[]".
    """

    def transform_image_edit_request(
        self,
        model: str,
        prompt: Optional[str],
        image: Optional[FileTypes],
        image_edit_optional_request_params: Dict,
        litellm_params: GenericLiteLLMParams,
        headers: dict,
    ) -> Tuple[Dict, RequestFiles]:
        """
        Transform image edit request for DALL-E-2.

        DALL-E-2 only accepts a single image with field name "image" (not "image[]").
        """        
        request_params = {
            "model": model,
            **image_edit_optional_request_params,
        }
        if image is not None:
            request_params["image"] = image
        if prompt is not None:
            request_params["prompt"] = prompt
            
        request = ImageEditRequestParams(**request_params)
        request_dict = cast(Dict, request)


        #########################################################
        # Separate images and masks as `files` and send other parameters as `data`
        #########################################################
        _image_list = request_dict.get("image")
        _mask = request_dict.get("mask")
        data_without_files = {
            k: v for k, v in request_dict.items() if k not in ["image", "mask"]
        }
        files_list: List[Tuple[str, Any]] = []

        # Handle image parameter - DALL-E-2 only supports single image
        if _image_list is not None:
            image_list = (
                [_image_list] if not isinstance(_image_list, list) else _image_list
            )

            # Validate only one image is provided
            if len(image_list) > 1:
                raise litellm.BadRequestError(
                    message="DALL-E-2 only supports editing a single image. Please provide one image.",
                    model=model,
                    llm_provider="openai",
                )

            # Use "image" field name (singular) for DALL-E-2
            for _image in image_list:
                if _image is not None:
                    self._add_image_to_files(
                        files_list=files_list,
                        image=_image,
                        field_name="image",
                    )

        # Handle mask parameter if provided
        if _mask is not None:
            # Handle case where mask can be a list (extract first mask)
            if isinstance(_mask, list):
                _mask = _mask[0] if _mask else None

            if _mask is not None:
                mask_content_type: str = ImageEditRequestUtils.get_image_content_type(
                    _mask
                )
                if isinstance(_mask, BufferedReader):
                    files_list.append(("mask", (_mask.name, _mask, mask_content_type)))
                else:
                    files_list.append(("mask", ("mask.png", _mask, mask_content_type)))

        return data_without_files, files_list