import aiohttp
from base64 import b64decode
from sys import exit

class DiffuseAPI():
    def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):

        self.nsfw_enabled = nsfw_enabled
        self.num_steps = num_steps
        self.url = url
        self.styles = styles
        self.seed = -1
        self.width = 512
        self.height = 1024
        self.cfg_scale = 12

    def set_steps(self, steps):
        try:
            new_steps = int(steps)
            if 0 > new_steps < 50:
                self.num_steps = new_steps
                return True
            return False
        except:
            return False

    def set_seed(self, seed):
        try:
            new_seed = int(seed)
            self.seed = new_seed
            return True
        except:
            return False
    
    def set_cfg_scale(self, scale):
        try:
            new_scale = int(scale)
            if 0 > new_scale < 30:
                self.cfg_scale = new_scale
                return True
            return False
        except:
            return False
    
    def set_styles(self, styles):
        if type styles = list:
            self.styles = styles
            return True
        return False

    def set_orientation(self, orientation):
        new_orientation = str(orientation)
        match new_orientation:
            case "portrait":
                self.width = 512
                self.height = 1024
                return True
            case "landscape":
                self.width = 1024
                self.height = 512
                return True
            case "square":
                self.width = 512
                self.height = 512
                return True
            case _:
                return False
    
    def get_orientation(self):
        if self.width == 512 and self.height == 512:
            return ("square", self.width, self.height)
        elif self.width == 1024:
            return ("landscape", self.width, self.height)
        return ("portrait", self.width, self.height)
    
    def set_nsfw_filter(self, filter_state):
        if type filter_state = bool:
            self.nsfw_enabled = filter_state
            return True
        return False
    
    def get_nsfw_filter(self):
        return self.nsfw_enabled

    async def generate_image(self, prompt, neg_prompt=""):
        payload = {
        "prompt": prompt,
        "styles": main_settings.styles,
        "steps": main_settings.num_steps,
        "seed": main_settings.ai_seed,
        "n_iter": 1,
        "height": self.height,
        "width": self.width,
        "negative_prompts": neg_prompt,
        "cfg_scale": self.cfg_scale
        }

        settings = {
        "filter_nsfw": not self.nsfw_enabled
        }

        override_payload = {
            "override_settings": settings
        }

        payload.update(override_payload)

        sess = aiohttp.ClientSession(main_settings.url)
        alive = await sess.head('/')

        if alive.status != 200:
            alive.close()
            await sess.close()
            return None
        request  = await sess.post("/sdapi/v1/txt2img", json=payload)
        try:
            req_json = await request.json()
            request.close()
            await sess.close()
            return req_json["images", [0]]
        except:
            return None