import asyncio
import requests
import os, tempfile
from .runpod import RunpodWrapper

import io
import base64
from PIL import Image, PngImagePlugin

import logging

logger = logging.getLogger(__name__)


class RunpodImageAutomaticWrapper(RunpodWrapper):

    async def generate(self, input_prompt: str, negative_prompt: str, endpoint_name: str, typing_fn, timeout=180):

        # Define your inputs
        input_data = {
            "input": {
                "prompt": input_prompt,
                "nagative_prompt": negative_prompt,
                "steps": 25,
                "cfg_scale": 7,
                "seed": -1,
                "width": 512,
                "height": 768,
                "batch_size": 3,
        #        "sampler_index": "DPM++ 2M Karras",
        #        "enable_hr": True,
        #        "hr_scale": 2,
        #        "hr_upscaler": "ESRGAN_4x", # "Latent"
        #        "denoising_strength": 0.5,
        #        "hr_second_pass_steps": 15,
                "restore_faces": True,
        #        "gfpgan_visibility": 0.5,
        #        "codeformer_visibility": 0.5,
        #        "codeformer_weight": 0.5,
        ##        "override_settings": {
        ##            "filter_nsfw": False,
        ##        },
                "api_endpoint": "txt2img",
            },
            "cmd": "txt2img"
        }

        output = await super().generate(input_data, endpoint_name, typing_fn, timeout)

        upscale = False
        if upscale:
            count = 0
            for i in output['images']:
                payload = {
                    "init_images": [i],
                    "prompt": input_prompt,
                    "nagative_prompt": negative_prompt,
                    "steps": 20,
                    "seed": -1,
                    #"sampler_index": "Euler",
    #                                  tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding, upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur, seams_fix_type, target_size_type, custom_width, custom_height, custom_scale
    #                "script_args": ["",512,0,8,32,64,0.275,32,3,False,0,True,8,3,2,1080,1440,1.875],
    #                "script_name": "Ultimate SD upscale",
                }
                upscaled_output = await serverless_automatic_request(payload, "img2img", api_url, api_key, typing_fn)
                output['images'][count] = upscaled_output['images'][count]


        os.makedirs("./.data/images", exist_ok=True)
        files = []
        for i in output['images']:
            temp_name = next(tempfile._get_candidate_names())
            filename = "./.data/images/" + temp_name + ".png"
            image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
            info = output['info']
            parameters = output['parameters']
            pnginfo = PngImagePlugin.PngInfo()
            pnginfo.add_text("parameters", info)
            image.save(filename, pnginfo=pnginfo)
            files.append(filename)

        return files