import logging
import asyncio
import aiohttp
from .ai_responder import exponential_backoff, AIResponderBase
from io import BytesIO


class LeonardoAIDrawMixIn(AIResponderBase):
    async def draw_leonardo(self, description: str) -> BytesIO:
        error_backoff = exponential_backoff(max_attempts=12)
        generation_id = None
        image_url = None
        image_bytes = None
        while True:
            error_sleep = next(error_backoff)
            try:
                async with aiohttp.ClientSession() as session:
                    if generation_id is None:
                        async with session.post("https://cloud.leonardo.ai/api/rest/v1/generations",
                                                json={"prompt": description,
                                                      "modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3",
                                                      "num_images": 1,
                                                      "sd_version": "v2",
                                                      "promptMagic": True,
                                                      "unzoomAmount": 1,
                                                      "width": 512,
                                                      "height": 512},
                                                headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
                                                         "Accept": "application/json",
                                                         "Content-Type": "application/json"},
                                                ) as response:
                            response = await response.json()
                            if "sdGenerationJob" not in response:
                                logging.warning(f"No 'sdGenerationJob' found in response, sleep for {error_sleep}s: {repr(response)}")
                                await asyncio.sleep(error_sleep)
                                continue
                            generation_id = response["sdGenerationJob"]["generationId"]
                    if image_url is None:
                        async with session.get(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
                                               headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
                                                        "Accept": "application/json"},
                                               ) as response:
                            response = await response.json()
                            if "generations_by_pk" not in response:
                                logging.warning(f"Unexpected response, sleep for {error_sleep}s: {repr(response)}")
                                await asyncio.sleep(error_sleep)
                                continue
                            if len(response["generations_by_pk"]["generated_images"]) == 0:
                                await asyncio.sleep(error_sleep)
                                continue
                            image_url = response["generations_by_pk"]["generated_images"][0]["url"]
                    if image_bytes is None:
                        async with session.get(image_url) as response:
                            image_bytes = BytesIO(await response.read())
                    async with session.delete(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
                                              headers={"Authorization": f"Bearer {self.config['leonardo-token']}"},
                                              ) as response:
                        await response.json()
                    logging.info(f'Drawed a picture with leonardo AI on this description: {repr(description)}')
                    return image_bytes
            except Exception as err:
                logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}\n{repr(err)}")
            else:
                logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}")
            await asyncio.sleep(error_sleep)
        raise RuntimeError(f"Failed to generate image {repr(description)}")