# 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_Edit_Client(OpenAI_Image_Client, Image_Context_Manager_Mixin):
    """ Create image edits from an image, mask and prompt. Return a URL to the image

        Takes the path (or opened path) to an image and a mask to indicate what part of the image should be replaced 
    """

    def __init__(self, image:Union[str, BufferedReader], mask: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)
        # # Open image mask from specified location or use already opened image
        self._mask = self._open_image(mask)


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

            Takes an optional path (or opened path) to an image and a mask to indicate what part
            of the image should be replaced if it was not passed during class instantiation
             
            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
        mask = self._open_image(mask) if mask else self._mask

        result = self._api.create_edit(image, mask, prompt=self.get_prompt_with_context(prompt), n=number_of_images, size=size, response_format=self._response_format)

        return self.process_response(number_of_images, result)


class B64_JSON_Image_Edit_Client(URL_Image_Edit_Client):
    """ Create image edits from an image, mask and prompt. Return the image in b64 JSON """

    _response_format = IMAGE_RESPONSE_FORMAT.B64_JSON
