import aiohttp
import asyncio

class DiffuseAPI:
    def __init__(self, url, nsfw_enabled, styles, steps):
        self.url = url
        self.nsfw_enabled = nsfw_enabled
        self.styles = styles
        self.steps = steps
        self.seed = -1

    def _generate_payload(self):
        payload = {
            "prompt": "",
            "negative_prompt": "",
            "steps": self.steps,
            "seed": self.seed,
            "styles": self.styles,
            "height": 1024,
            "width": 512
        }
        settings = {
            "filter_nsfw": not self.nsfw_enabled,
            "samples_save": True,
        }

        override_payload = {
            "override_settings": settings
        }
        payload.update(override_payload)
        return payload

    async def generate_image(self, prompt, negative_prompt=""):
        async with aiohttp.ClientSession(self.url) as session:
            payload = self._generate_payload()
            payload.update({"prompt": prompt, "negative_prompt": negative_prompt})
            print(payload)
            async with session.post('/sdapi/v1/txt2img', json=payload) as image_handler:
                image_data = await image_handler.json()

        print(image_data)