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

        print(resp_json)

        # 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())


        messages = resp_json[prompt_id]["status"]["messages"]
        timestamp_start = 0
        timestamp_finish = 0

        for m in messages:
            label   = m[0]
            obj     = m[1]
            if label == "execution_start":      timestamp_start = obj["timestamp"]
            if label == "execution_success":    timestamp_finish = obj["timestamp"]

        execution_time = (timestamp_finish - timestamp_start) / 1000.0

        return {
            "execution_time":   execution_time,
            "images":           output_images
        }