import os
import time
from io import BytesIO

import requests
from openai import OpenAI
from openai.types.images_response import ImagesResponse
from PIL import Image

from .base_model import BaseModel

from .config import DEFAULT_ITI_MODELS
from .config import MAX_SEED
from .config import OPENAI_KEY_NAME
from .config import OPENAI_ITI_ACCEPTABLE_COLOR_MODE
from .config import OPENAI_ITI_DEFAULT_N
from .config import OPENAI_ITI_MAX_N
from .config import OPENAI_ITI_MODE_LIST
from .config import OPENAI_ITI_SIZE_LIST
from .config import REQUEST_ACCEPTED
from .config import STABLE_DIFFUSION_BASE_EP
from .config import STABLE_DIFFUSION_KEY_NAME
from .config import STABLE_DIFFUSION_ITI_EDIT_ERASE_EP
from .config import STABLE_DIFFUSION_ITI_EDIT_INPAINT_EP
from .config import STABLE_DIFFUSION_ITI_EDIT_OUTPAINT_EP
from .config import STABLE_DIFFUSION_ITI_EDIT_REMOVE_BACKGROUND_EP
from .config import STABLE_DIFFUSION_ITI_EDIT_SEARCH_AND_REPLACE_EP
from .config import STABLE_DIFFUSION_ITI_OUTPUT_FORMATS
from .config import STABLE_DIFFUSION_ITI_UPSCALE_CONSERVATIVE_EP
from .config import STABLE_DIFFUSION_ITI_UPSCALE_CREATIVE_START_EP
from .config import STABLE_DIFFUSION_ITI_UPSCALE_CREATIVE_RESULT_EP
from .config import STABLE_DIFFUSION_CONSERVATIVE_CREATIVITY_MIN
from .config import STABLE_DIFFUSION_CONSERVATIVE_CREATIVITY_MAX
from .config import STABLE_DIFFUSION_CREATIVE_CREATIVITY_MIN
from .config import STABLE_DIFFUSION_CREATIVE_CREATIVITY_MAX
from .config import STABLE_DIFFUSION_GROW_MASK_MIN
from .config import STABLE_DIFFUSION_GROW_MASK_MAX
from .config import STABLE_DIFFUSION_OUTPAINT_LEFT_MAX
from .config import STABLE_DIFFUSION_OUTPAINT_UP_MAX
from .config import STABLE_DIFFUSION_OUTPAINT_RIGHT_MAX
from .config import STABLE_DIFFUSION_OUTPAINT_DOWN_MAX
from .config import STABLE_DIFFUSION_OUTPAINT_CREATIVITY_MIN
from .config import STABLE_DIFFUSION_OUTPAINT_CREATIVITY_MAX
from .config import WAIT_FOR_UPSCALE_CREATIVE_RESULT


class OpenAIImageToImage(BaseModel):
    '''
    List of available models as of 2024-07-12:
      - dall-e-2
    Note: OpenAI image edit and variation accept only PNG format.
    '''
    def __init__(self, **kwargs):

        try:
            super().__init__(**kwargs)

        except Exception as e:
            msg = 'Error while verifying specific parameters for '
            msg += 'OpenAIImageToImage'
            raise Exception(msg) from e

    def run(self):
        '''
        Note:
        OpenAI Image-To-Image returns class ImagesResponse.
        Save the generated image by accesing data.url for both single and
        multiple outputs.
        When failed to generate image file, return value is given in str.
        Handle return value `answer` with care for different type in case
        of success and failure.
        '''
        answer = 'Editted image not found.'

        try:
            client = OpenAI(api_key=os.getenv(OPENAI_KEY_NAME))

            if self.parameters['mode'] == 'variations':
                response = client.images.create_variation(
                    model=self.parameters['model'],
                    image=self.parameters['image'],
                    n=self.parameters['n'],
                    size=self.parameters['size'])
            elif (self.parameters['mode'] == 'edits' and
                  'mask' in self.parameters):
                response = client.images.edit(
                    model=self.parameters['model'],
                    image=self.parameters['image'],
                    mask=self.parameters['mask'],
                    prompt=self.parameters['prompt'],
                    n=self.parameters['n'],
                    size=self.parameters['size'])
            elif (self.parameters['mode'] == 'edits' and
                  'mask' not in self.parameters):
                response = client.images.edit(
                    model=self.parameters['model'],
                    image=self.parameters['image'],
                    prompt=self.parameters['prompt'],
                    n=self.parameters['n'],
                    size=self.parameters['size'])
            else:
                raise Exception("Unexpected type of process selected.")

            if isinstance(response, ImagesResponse):
                # ImagesResponse class including generated image url
                answer = response

        except Exception as e:
            answer = str(e)

        self.response = answer

    def _verify_arguments(self, **kwargs):
        '''
        Expected inputs:
          - mode (required): either 'edits' or 'variations'
          - image (required): str, path to image file (PNG only)
          - mask: str, path to mask file (PNG only)
          - n: int, number of images to generate
          - size: str, size of image to generate
        '''
        parameters = kwargs

        if 'mode' not in kwargs:
            raise ValueError("'mode' parameter is required")
        elif kwargs['mode'] not in OPENAI_ITI_MODE_LIST:
            raise ValueError(f"'mode' {kwargs['mode']} is not valid")

        if 'image' not in kwargs:
            raise ValueError("'image' parameter is required")
        elif not os.path.isfile(kwargs['image']):
            raise ValueError(f"'image' {kwargs['image']} is not a valid file")
        else:
            parameters['image'] = self._get_byte_image(kwargs['image'])

        if 'mask' in kwargs and os.path.isfile(kwargs['mask']):
            parameters['mask'] = self._get_byte_image(kwargs['mask'])

        if 'size' in kwargs and kwargs['size'] in OPENAI_ITI_SIZE_LIST:
            pass
        else:
            parameters.update(size=OPENAI_ITI_SIZE_LIST[0])

        if 'n' not in kwargs:
            parameters.update(n=OPENAI_ITI_DEFAULT_N)
        else:
            buff = int(kwargs['n'])
            if isinstance(buff, int) and 0 < buff and buff <= OPENAI_ITI_MAX_N:
                pass
            else:
                parameters['n'] = OPENAI_ITI_DEFAULT_N

        return parameters

    def _get_byte_image(self, to_convert):
        '''
        Convert image to bytes in PNG format.
        OpenAI accepts RBGA, LA and L of color formats.
        '''
        byte_arr = BytesIO()
        with Image.open(to_convert) as image:
            if image.mode not in OPENAI_ITI_ACCEPTABLE_COLOR_MODE:
                image = image.convert('RGBA')
            image.save(byte_arr, format='PNG')
        return byte_arr.getvalue()


class StableDiffusionImageToImage(BaseModel):

    def __init__(self, **kwargs):

        try:
            super().__init__(**kwargs)

        except Exception as e:
            msg = 'Error while verifying specific parameters for '
            msg += 'StableDiffusionImageToImage'
            raise Exception(msg) from e

    def run(self):
        '''
        Note:
        Stable Diffusion Image-To-Image returns class.
        Save the generated image by accesing output[0].url.
        When failed to generate image file, return value is given in str.
        Handle return value `answer` with care for different type in case of
        success and failure.
        '''
        answer = 'Editted image not found.'

        try:
            response = self._execute_api_call(self.parameters['url'],
                                              self.parameters['headers'],
                                              self.parameters['files'],
                                              self.parameters['data'])

            if self.parameters['mode'] == 'upscale_creative':
                response = self._fetch_upscale_creative_result(
                    response.json().get('id'))

            if hasattr(response, 'content'):
                # binary data of generated image
                answer = response.content

            else:
                answer += str(response.json())

        except Exception as e:
            answer = str(e)

        self.response = answer

    def _verify_arguments(self, **kwargs):
        """
        Expected inputs:
          - model (required): str, model name, fixed v2beta
          - image (required): str, path to image file (PNG only)
          - for rest of arguments, see each function by mode
        """
        parameters = {}

        # header for REST API
        headers = {'authorization':
                   f'Bearer {os.getenv(STABLE_DIFFUSION_KEY_NAME)}'}
        headers.update(accept='image/*')

        parameters.update(headers=headers)

        # correct default model to fix v2beta
        # this process is critical to make endpoint for each mode
        default_model = DEFAULT_ITI_MODELS['stable_diffusion_iti']
        if kwargs['model'] not in default_model:
            kwargs['model'] = default_model
            parameters.update(model=default_model)

        # Check if the image file is valid
        if not os.path.isfile(kwargs['image']):
            msg = f"'image' {kwargs['image']} is not a valid file"
            raise ValueError(msg)
        else:
            with open(kwargs['image'], "rb") as f:
                parameters.update(files={'image': f.read()})

        # rest of parameters are verified in each mode
        if kwargs['mode'] == 'upscale_conservative':
            parameters.update(self._verify_upscale_conservative_args(**kwargs))
        elif kwargs['mode'] == 'upscale_creative':
            parameters.update(self._verify_upscale_creative_args(**kwargs))
        elif kwargs['mode'] == 'erase':
            parameters.update(self._verify_erase_args(**kwargs))
        elif kwargs['mode'] == 'inpaint':
            parameters.update(self._verify_inpaint_args(**kwargs))
        elif kwargs['mode'] == 'outpaint':
            parameters.update(self._verify_outpaint_args(**kwargs))
        elif kwargs['mode'] == 'search_and_replace':
            parameters.update(self._verify_search_and_replace_args(**kwargs))
        elif kwargs['mode'] == 'remove_background':
            parameters.update(self._verify_remove_background_args(**kwargs))
        else:
            raise ValueError(f"Invalid mode: {kwargs['mode']}")

        # add mask if exist
        if 'mask' in parameters:
            parameters['files'].update(mask=parameters['mask'])
            parameters.pop('mask')

        return parameters

    # common execute function for REST API call
    def _execute_api_call(self, endpoint='', headers={}, files={}, data={}):

        return requests.post(endpoint,
                             headers=headers,
                             files=files,
                             data=data)

    # fetch result function for upscale creative
    def _fetch_upscale_creative_result(self, id):

        answer = 'Editted file not found.'
        ep = self.parameters['url_result'].replace('{id}', id)

        flg = True

        while flg:
            response = requests.request(method='GET',
                                        url=ep,
                                        headers=self.parameters['headers'])

            if response.status_code == REQUEST_ACCEPTED:
                time.sleep(WAIT_FOR_UPSCALE_CREATIVE_RESULT)

            else:
                answer = response
                flg = False

        return answer

    # parameter check function for individual case
    def _verify_upscale_conservative_args(self, **kwargs):
        '''
        Expected options:
          - creativity: float, creativity value
          - negative_prompt: str, negative prompt
          - seed: int, seed value
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_UPSCALE_CONSERVATIVE_EP)

        parameters.update(url=endpoint)

        # data body
        data = {'prompt': kwargs['prompt']}

        if 'negative_prompt' in kwargs:
            data.update(negative_prompt=kwargs['negative_prompt'])

        if 'seed' in kwargs:
            seed = kwargs['seed']
            if isinstance(seed, int) and 0 <= seed and seed <= MAX_SEED:
                data.update(seed=seed)

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        if 'creativity' in kwargs:
            creativity = kwargs['creativity']
            if (isinstance(creativity, float) and
               STABLE_DIFFUSION_CONSERVATIVE_CREATIVITY_MIN <= creativity and
               creativity <= STABLE_DIFFUSION_CONSERVATIVE_CREATIVITY_MAX):
                data.update(creativity=creativity)

        parameters.update(data=data)

        return parameters

    def _verify_upscale_creative_args(self, **kwargs):
        '''
        Expected options:
          - creativity: float, creativity value
          - negative_prompt: str, negative prompt
          - seed: int, seed value
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_UPSCALE_CREATIVE_START_EP)

        parameters.update(url=endpoint)

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_UPSCALE_CREATIVE_RESULT_EP)

        parameters.update(url_result=endpoint)

        # data body
        data = {'prompt': kwargs['prompt']}

        if 'negative_prompt' in kwargs:
            data.update(negative_prompt=kwargs['negative_prompt'])

        if 'seed' in kwargs:
            seed = kwargs['seed']
            if isinstance(seed, int) and 0 <= seed and seed <= MAX_SEED:
                data.update(seed=seed)

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        if 'creativity' in kwargs:
            creativity = kwargs['creativity']
            if (isinstance(creativity, float) and
               STABLE_DIFFUSION_CREATIVE_CREATIVITY_MIN <= creativity and
               creativity <= STABLE_DIFFUSION_CREATIVE_CREATIVITY_MAX):
                data.update(creativity=creativity)

        parameters.update(data=data)

        return parameters

    def _verify_erase_args(self, **kwargs):
        '''
        Expected options:
          - mask: str, path to mask file
          - grow_mask: int, grow mask value
          - seed: int, seed value
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_EDIT_ERASE_EP)

        parameters.update(url=endpoint)

        # additional files of mask (optional)
        if 'mask' in kwargs and os.path.isfile(kwargs['mask']):
            with open(kwargs['mask'], "rb") as f:
                parameters.update(mask=f.read())

        # data body
        data = {}

        if 'grow_mask' in kwargs:
            grow_mask = kwargs['grow_mask']
            if (isinstance(grow_mask, int) and
               STABLE_DIFFUSION_GROW_MASK_MIN <= grow_mask and
               grow_mask <= STABLE_DIFFUSION_GROW_MASK_MAX):
                data.update(grow_mask=grow_mask)

        if 'seed' in kwargs:
            seed = kwargs['seed']
            if isinstance(seed, int) and 0 <= seed and seed <= MAX_SEED:
                data.update(seed=seed)

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        parameters.update(data=data)

        return parameters

    def _verify_inpaint_args(self, **kwargs):
        '''
        Expected options:
          - mask: str, path to mask file
          - grow_mask: int, grow mask value
          - negative_prompt: str, negative prompt
          - seed: int, seed value
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_EDIT_INPAINT_EP)

        parameters.update(url=endpoint)

        # additional files of mask (optional)
        if 'mask' in kwargs and os.path.isfile(kwargs['mask']):
            with open(kwargs['mask'], "rb") as f:
                parameters.update(mask=f.read())

        # data body
        data = {'prompt': kwargs['prompt']}

        if 'grow_mask' in kwargs:
            grow_mask = kwargs['grow_mask']
            if (isinstance(grow_mask, int) and
               STABLE_DIFFUSION_GROW_MASK_MIN <= grow_mask and
               grow_mask <= STABLE_DIFFUSION_GROW_MASK_MAX):
                data.update(grow_mask=grow_mask)

        if 'negative_prompt' in kwargs:
            data.update(negative_prompt=kwargs['negative_prompt'])

        if 'seed' in kwargs:
            seed = kwargs['seed']
            if isinstance(seed, int) and 0 <= seed and seed <= MAX_SEED:
                data.update(seed=seed)

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        parameters.update(data=data)

        return parameters

    def _verify_outpaint_args(self, **kwargs):
        '''
        Expected options:
          - left: int, left value
          - right: int, right value
          - up: int, up value
          - down: int, down value
          - creativity: float, creativity value
          - prompt: str, prompt
          - seed: int, seed value
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_EDIT_OUTPAINT_EP)

        parameters.update(url=endpoint)

        # data body
        data = {}

        if 'left' in kwargs:
            left = kwargs['left']
            if (isinstance(left, int) and
               0 <= left and left <= STABLE_DIFFUSION_OUTPAINT_LEFT_MAX):
                data.update(left=left)

        if 'right' in kwargs:
            right = kwargs['right']
            if (isinstance(right, int) and
               0 <= right and right <= STABLE_DIFFUSION_OUTPAINT_RIGHT_MAX):
                data.update(right=right)

        if 'up' in kwargs:
            up = kwargs['up']
            if (isinstance(up, int) and
               0 <= up and up <= STABLE_DIFFUSION_OUTPAINT_UP_MAX):
                data.update(up=up)

        if 'down' in kwargs:
            down = kwargs['down']
            if (isinstance(down, int) and
               0 <= down and down <= STABLE_DIFFUSION_OUTPAINT_DOWN_MAX):
                data.update(down=down)

        if 'creativity' in kwargs:
            creativity = kwargs['creativity']
            if (isinstance(creativity, float) and
               STABLE_DIFFUSION_OUTPAINT_CREATIVITY_MIN <= creativity and
               creativity <= STABLE_DIFFUSION_OUTPAINT_CREATIVITY_MAX):
                data.update(creativity=creativity)

        if 'prompt' in kwargs:
            data.update(prompt=kwargs['prompt'])

        if 'seed' in kwargs:
            seed = kwargs['seed']
            if isinstance(seed, int) and 0 <= seed and seed <= MAX_SEED:
                data.update(seed=seed)

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        parameters.update(data=data)

        return parameters

    def _verify_search_and_replace_args(self, **kwargs):
        '''
        Expected options:
          - search_prompt (required): str, search prompt
          - negative_prompt: str, negative prompt
          - grow_mask: int, grow mask value
          - seed: int, seed value
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_EDIT_SEARCH_AND_REPLACE_EP)

        parameters.update(url=endpoint)

        # data body
        data = {'prompt': kwargs['prompt']}

        if 'search_prompt' not in kwargs:
            raise ValueError("'search_prompt' is required")
        else:
            data.update(search_prompt=kwargs['search_prompt'])

        if 'negative_prompt' in kwargs:
            data.update(negative_prompt=kwargs['negative_prompt'])

        if 'grow_mask' in kwargs:
            grow_mask = kwargs['grow_mask']
            if (isinstance(grow_mask, int) and
               STABLE_DIFFUSION_GROW_MASK_MIN <= grow_mask and
               grow_mask <= STABLE_DIFFUSION_GROW_MASK_MAX):
                data.update(grow_mask=grow_mask)

        if 'seed' in kwargs:
            seed = kwargs['seed']
            if isinstance(seed, int) and 0 <= seed and seed <= MAX_SEED:
                data.update(seed=seed)

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        parameters.update(data=data)

        return parameters

    def _verify_remove_background_args(self, **kwargs):
        '''
        Expected options:
          - output_format: str, output format
        '''
        parameters = kwargs

        endpoint = (STABLE_DIFFUSION_BASE_EP + '/' + kwargs['model'] +
                    STABLE_DIFFUSION_ITI_EDIT_REMOVE_BACKGROUND_EP)

        parameters.update(url=endpoint)

        # data body
        data = {}

        if ('output_format' in kwargs and
           kwargs['output_format'] in STABLE_DIFFUSION_ITI_OUTPUT_FORMATS):
            data.update(output_format=kwargs['output_format'])

        parameters.update(data=data)

        return parameters
