"""Encola N variantes de un workflow ComfyUI, una por seed, y recoge los prompt_ids.

Funcion impura: hace red (POST /prompt por variante, via comfyui_submit_workflow).
Compone comfyui_submit_workflow.

Para cada seed de la lista, copia el workflow (deepcopy, no muta el original),
parchea el campo de semilla de los nodos sampler (KSampler.seed, KSamplerAdvanced.
noise_seed, SamplerCustom.noise_seed — en general cualquier input "seed"/"noise_seed")
y lo encola. Util para barridos de re-roll: misma escena, varias semillas, una sola
llamada. Devuelve los prompt_ids en el mismo orden que la lista de seeds; cada uno
se sigue con comfyui_wait_result.
"""
import copy
import os
import sys

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
if _THIS_DIR not in sys.path:
    sys.path.insert(0, _THIS_DIR)

from comfyui_submit_workflow import comfyui_submit_workflow  # noqa: E402

# Campos de semilla conocidos en los nodos sampler de ComfyUI.
_SEED_KEYS = ("seed", "noise_seed")


def _patch_seed(workflow: dict, seed: int) -> dict:
    """Copia el workflow y fija `seed` en todos los inputs de semilla (no muta el original)."""
    wf = copy.deepcopy(workflow)
    for node in wf.values():
        inputs = node.get("inputs")
        if not isinstance(inputs, dict):
            continue
        for key in _SEED_KEYS:
            if key in inputs:
                inputs[key] = seed
    return wf


def comfyui_batch_generate(
    workflow: dict,
    *,
    seeds: list | None = None,
    server: str = "127.0.0.1:8188",
) -> dict:
    """Encola una variante del workflow por cada seed y devuelve los prompt_ids.

    Args:
        workflow: dict en API format (resultado de un builder). No se muta: cada
            variante es una copia profunda con la semilla parcheada.
        seeds: lista de semillas (int). Cada una produce una variante encolada. Si
            es None o vacia, se encola el workflow tal cual una sola vez (sin
            parchear semilla). keyword-only.
        server: host:port del servidor ComfyUI sin esquema. keyword-only.

    Returns:
        dict con:
        - ok (bool): True si TODAS las variantes se encolaron sin error.
        - prompt_ids (list[str]): prompt_id de cada variante encolada, en orden.
        - count (int): numero de variantes encoladas con exito.
        - error (str): primer error encontrado; cadena vacia si todo OK. Si una
          variante falla, se detiene el barrido y se devuelven los prompt_ids ya
          encolados.
    """
    out = {"ok": False, "prompt_ids": [], "count": 0, "error": ""}
    variants = [(s, _patch_seed(workflow, s)) for s in seeds] if seeds else [(None, workflow)]

    for seed, wf in variants:
        try:
            resp = comfyui_submit_workflow(wf, server=server)
        except RuntimeError as exc:
            label = "tal cual" if seed is None else f"seed={seed}"
            out["error"] = f"variante {label} fallo al encolar: {exc}"
            return out
        out["prompt_ids"].append(resp["prompt_id"])

    out["count"] = len(out["prompt_ids"])
    out["ok"] = True
    return out


if __name__ == "__main__":
    from comfyui_build_txt2img_workflow import comfyui_build_txt2img_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",
    )
    res = comfyui_batch_generate(wf, seeds=[1, 2])
    print(f"ok={res['ok']} count={res['count']} ids={res['prompt_ids']} error={res['error']!r}")