import time

import asyncio
import aiohttp

from urllib.request import urlopen as urlopen
from urllib.parse import urlencode as urlencode

async def get_comfyui_generations(api_url, workflow):
    async with aiohttp.ClientSession() as session:
        async with session.post(f"{api_url}/prompt", json={"prompt": workflow}) as resp:
            resp_json = await resp.json()

        prompt_id = resp_json["prompt_id"]

        # Loop endlessly until ComfyUI confirms our prompt's completion:
        while True:
            async with session.get(f"{api_url}/history/{prompt_id}") as resp:
                resp_json = await resp.json()
            if not resp_json:
                time.sleep(1)
                continue
            break

        # Read the output history anmd fetch each image:
        history         = resp_json[prompt_id]
        output_images   = []

        for o in history["outputs"]:
            for node_id in history["outputs"]:
                node_output     = history["outputs"][node_id]
                files_output    = []

                if "images" in node_output:
                    for image in node_output["images"]:
                        url_params = urlencode({
                            "filename":     image["filename"],
                            "subfolder":    image["subfolder"],
                            "type":         image["type"]
                        })

                        async with session.get(f"{api_url}/view?{url_params}") as resp:
                            output_images.append(await resp.content.read())

        return output_images