"""
企业级 ComfyUI 异步 TaskWorker
Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119)
"""

import asyncio
import json
import logging
import uuid
import aiohttp
from typing import Dict, Any, List

logger = logging.getLogger("VerseFlow.ComfyUIWorker")

class ComfyUIExecutionError(Exception):
    pass

class ComfyUITaskWorker:
    def __init__(self, host: str = "127.0.0.1", port: int = 8188):
        self.base_url = f"http://{host}:{port}"
        self.ws_url = f"ws://{host}:{port}/ws"
        self.client_id = str(uuid.uuid4())

    async def execute_prompt(self, workflow_prompt: Dict[str, Any], task_id: str) -> List[Dict[str, Any]]:
        async with aiohttp.ClientSession() as session:
            async with session.ws_connect(f"{self.ws_url}?clientId={self.client_id}") as ws:
                payload = {"prompt": workflow_prompt, "client_id": self.client_id}
                async with session.post(f"{self.base_url}/prompt", json=payload) as resp:
                    if resp.status != 200:
                        raise ComfyUIExecutionError(f"提交失败: {await resp.text()}")
                    prompt_id = (await resp.json())["prompt_id"]

                while True:
                    msg = await ws.receive(timeout=120.0)
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        event = json.loads(msg.data)
                        if event.get("type") == "executing" and event.get("data", {}).get("prompt_id") == prompt_id:
                            if event.get("data", {}).get("node") is None:
                                break

                async with session.get(f"{self.base_url}/history/{prompt_id}") as hist_resp:
                    history = (await hist_resp.json()).get(prompt_id, {})
                    return self._extract_outputs(history)

    def _extract_outputs(self, history: Dict[str, Any]) -> List[Dict[str, Any]]:
        outputs = []
        for node_id, node_out in history.get("outputs", {}).items():
            if "images" in node_out:
                for img in node_out["images"]:
                    outputs.append({"node_id": node_id, "filename": img["filename"], "type": "image"})
        return outputs