# Base client
from ....base import OpenAI_Image_Client

# Mixin
from ....mixins import Image_Context_Manager_Mixin

# Typing
from ....enums import IMAGE_RESPONSE_FORMAT, IMAGE_SIZE
from typing import Union, List
from io import BufferedReader


class URL_Image_Variation_Client(OpenAI_Image_Client):
    """ Create image variations from an image. Return a URL to the image

        Takes the path (or opened path) to an image 
    """

    def __init__(self, image:Union[str, BufferedReader], pre_prompt_context:List[str] = None, post_prompt_context:List[str] = None, api_key: str = None, max_retries=3, ms_between_retries=500, default_size:IMAGE_SIZE = None, 
                    default_number_of_images:int=1, default_response_format:str=None) -> None:
        
        super().__init__(api_key, max_retries, ms_between_retries, default_size, default_number_of_images, default_response_format)
        Image_Context_Manager_Mixin.__init__(self, pre_prompt_context=pre_prompt_context, post_prompt_context=post_prompt_context)
        
        # Open image from specified location or use already opened image
        self._image = self._open_image(image)


    def run_prompt(self, image:Union[str, BufferedReader]=None, mask:Union[str, BufferedReader]=None, number_of_images:int=None, size:IMAGE_SIZE=None):
        """ Sends an image to OpenAI to have variations generated by DALLE. You can specify the 
            number of images to generate and the dimensions of the image.

            If more than one image should be generated, this method will return a list of URLs 
            or b64 JSON for each image. 
        """

        # Allow the number_of_images to generate and image size to be specified per call, 
        # but fall back on defaults
        number_of_images = number_of_images or self._number_of_images
        size = self._check_size(size or self._size)
        
        image = self._open_image(image) if image else self._image

        result = self._api.create_variation(image, n=number_of_images, size=size, response_format=self._response_format)

        return self.process_response(number_of_images, result)


class B64_Image_Variation_Client(URL_Image_Variation_Client):
    """ Create image variations from an image. Return the images in b64 JSON """

    _response_format = IMAGE_RESPONSE_FORMAT.B64_JSON
