"""WebSocket 통신 및 워크플로우 완료 대기."""

import json
import logging
import uuid
from collections.abc import Callable
from typing import Any, Optional

logger = logging.getLogger(__name__)


async def wait_for_completion(
    ws,
    target_prompt_id: str,
    progress_callback: Optional[Callable[[float], None]] = None,
    preview_callback: Optional[Callable[[bytes, int, int], None]] = None,
) -> None:
    """
    WebSocket 메시지를 수신하며 워크플로우 완료를 기다린다.

    완료 신호: { "type": "executing", "data": { "node": null } }
    에러 신호: { "type": "execution_error" }
    프리뷰 신호: 바이너리 메시지 (이미지)
    """
    async for raw in ws:
        try:
            if isinstance(raw, bytes):
                # 바이너리 프리뷰 메시지 (KSampler 중간 이미지)
                if preview_callback and raw:
                    # 프리뷰 형식: [단계, 총단계] (4바이트) + PNG 바이너리
                    if len(raw) >= 8:
                        import struct
                        step, total = struct.unpack(">II", raw[:8])
                        preview_callback(raw[8:], step, total)
                continue
            data = json.loads(raw)
        except (json.JSONDecodeError, Exception):
            continue

        msg_type = data.get("type")
        msg_data = data.get("data", {})

        # 다른 prompt_id 메시지 무시
        pid = msg_data.get("prompt_id")
        if pid and pid != target_prompt_id:
            continue

        if msg_type == "executing":
            node = msg_data.get("node")
            if node is None:
                logger.debug(f"[ComfyBridgeIO] Execution complete: {target_prompt_id}")
                return  # 완료

        elif msg_type == "progress" and progress_callback:
            value = msg_data.get("value", 0)
            maximum = msg_data.get("max", 1)
            if maximum > 0:
                progress_callback(value / maximum)

        elif msg_type == "execution_error":
            err = msg_data.get("exception_message", "Unknown error")
            raise RuntimeError(f"[ComfyBridgeIO] Execution error: {err}")