"""Sondea GET /history/{prompt_id} hasta que un workflow ComfyUI termina.

Funcion impura: hace red (HTTP GET en bucle) y duerme entre sondeos. Solo
stdlib (urllib, json, time).

Usa polling de /history como mecanismo principal (no WebSocket): es mas robusto
porque no depende de websocket-client, que no esta garantizado en el venv. Para
saber si el resultado esta listo (no streaming de progreso paso a paso) el
polling de /history es suficiente y portable.
"""
import json
import time
import urllib.error
import urllib.request


def comfyui_wait_result(
    prompt_id: str,
    server: str = "127.0.0.1:8188",
    timeout: float = 180.0,
    poll_interval: float = 1.0,
) -> dict:
    """Espera a que ComfyUI termine de ejecutar un prompt y devuelve sus outputs.

    Sondea GET /history/{prompt_id} cada poll_interval segundos hasta que
    status.completed es True o status.status_str es "success"/"error", o hasta
    agotar el timeout.

    Args:
        prompt_id: id devuelto por comfyui_submit_workflow.
        server: host:port del servidor ComfyUI (sin esquema).
        timeout: maximo de segundos a esperar antes de fallar.
        poll_interval: segundos entre sondeos.

    Returns:
        dict de outputs {node_id: {"images": [{filename, subfolder, type}, ...]}}
        tal como ComfyUI los expone en history[prompt_id]["outputs"]. Puede
        contener otros tipos de output (gifs, texto) segun los nodos del
        workflow.

    Raises:
        TimeoutError: si se agota el timeout sin que el prompt complete.
        RuntimeError: si la ejecucion termina con status_str "error", si no se
            puede conectar, o si la respuesta no es JSON valido.
    """
    url = f"http://{server}/history/{prompt_id}"
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=timeout) as resp:
                hist = json.loads(resp.read())
        except urllib.error.URLError as exc:
            raise RuntimeError(
                f"comfyui_wait_result: no se pudo conectar a {url}: {exc.reason}"
            ) from exc
        except json.JSONDecodeError as exc:
            raise RuntimeError(
                f"comfyui_wait_result: respuesta no es JSON valido desde {url}: {exc}"
            ) from exc

        entry = hist.get(prompt_id)
        if entry:
            status = entry.get("status", {})
            status_str = status.get("status_str")
            if status_str == "error":
                raise RuntimeError(
                    f"comfyui_wait_result: ejecucion fallo para {prompt_id}: "
                    f"{json.dumps(status)[:500]}"
                )
            if status.get("completed") or status_str == "success":
                return entry.get("outputs", {})
        time.sleep(poll_interval)

    raise TimeoutError(
        f"comfyui_wait_result: timeout de {timeout}s esperando {prompt_id}"
    )


if __name__ == "__main__":
    import sys

    from comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
    from comfyui_submit_workflow import comfyui_submit_workflow

    wf = comfyui_build_txt2img_workflow(
        ckpt_name="v1-5-pruned-emaonly-fp16.safetensors",
        positive="a red apple on a wooden table, sharp focus",
        negative="blurry, low quality",
        steps=20,
        seed=42,
    )
    resp = comfyui_submit_workflow(wf)
    pid = resp["prompt_id"]
    print(f"esperando prompt_id={pid} ...", file=sys.stderr)
    outputs = comfyui_wait_result(pid)
    for node_id, out in outputs.items():
        for img in out.get("images", []):
            print(f"OUTPUT node={node_id} filename={img['filename']}")