fn_registry/python/functions/pipelines/comfyui_generate_until_quality.py

"""comfyui_generate_until_quality — loop evaluator-optimizer (GAN sin entrenar).

Genera una imagen con un builder del registry, la juzga con el panel multi-juez
(`comfyui_judge_image`), y si no alcanza la calidad pedida REFINA (nueva seed,
mas calidad, prompt corregido con el feedback del juez) y regenera, hasta que
pasa el umbral (`verdict == 'good'`) o se agotan los intentos. Siempre devuelve
la MEJOR candidata por score (best-of-N): nunca devuelve basura por agotar
iteraciones.

Es la doctrina del issue 0087 (promover una secuencia repetida a un pipeline
one-shot) aplicada al bucle de mejora del grupo `comfyui-skill`: build -> submit
-> wait -> fetch -> judge -> (refine) -> repeat. Compone funciones del registry:

    <builder>_py_ml                       (workflow de nodos en API format)
    comfyui_submit_workflow_py_ml         (POST /prompt)
    comfyui_wait_result_py_ml             (poll /history)
    comfyui_fetch_output_image_py_ml      (GET /view -> disco)
    comfyui_judge_image_py_ml             (panel estetico + CLIP + critica LLM)
    ask_llm_py_core                       (refine del prompt con el feedback)

Pipeline impuro: red (HTTP), GPU (generacion), disco (PNG), y API (juez critico
+ refine de prompt). Determinista en estructura: nunca lanza excepcion cruda,
siempre devuelve un dict de estado.
"""

from __future__ import annotations

import importlib
import inspect
import os
import sys

# Importa las funciones del registry (mismo arbol python/functions).
_FUNCTIONS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _FUNCTIONS_ROOT not in sys.path:
    sys.path.insert(0, _FUNCTIONS_ROOT)

from ml.comfyui_fetch_output_image import comfyui_fetch_output_image
from ml.comfyui_judge_image import comfyui_judge_image
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_wait_result import comfyui_wait_result


# Primo grande para derrochar el espacio de seeds entre rerolls de forma
# determinista (mismo subject + mismo base_seed -> misma traza de seeds).
_SEED_STRIDE = 101_117


def _resolve_builder(builder):
    """Devuelve el callable del builder.

    Acepta un callable directo o el nombre de la funcion (string), que se
    resuelve desde el paquete `ml` (convencion del registry: el modulo se llama
    igual que la funcion, p.ej. `comfyui_build_ui_hud_workflow`).
    """
    if callable(builder):
        return builder
    if isinstance(builder, str):
        mod = importlib.import_module(f"ml.{builder}")
        return getattr(mod, builder)
    raise TypeError(
        f"builder debe ser callable o str (nombre de funcion ml.*), no {type(builder)}"
    )


def _extract_positive_prompt(workflow: dict) -> str:
    """Extrae el prompt positivo textual del workflow para pasarselo al juez.

    Sigue el input `positive` del KSampler hasta su CLIPTextEncode. Fallback: el
    CLIPTextEncode con el texto mas largo (heuristica: el positive suele serlo).
    """
    if not isinstance(workflow, dict):
        return ""
    for node in workflow.values():
        if not isinstance(node, dict):
            continue
        if node.get("class_type") in ("KSampler", "KSamplerAdvanced"):
            pos = node.get("inputs", {}).get("positive")
            if isinstance(pos, list) and pos:
                tgt = workflow.get(str(pos[0]))
                if isinstance(tgt, dict) and tgt.get("class_type") == "CLIPTextEncode":
                    txt = tgt.get("inputs", {}).get("text")
                    if isinstance(txt, str) and txt.strip():
                        return txt
    texts = [
        n["inputs"]["text"]
        for n in workflow.values()
        if isinstance(n, dict)
        and n.get("class_type") == "CLIPTextEncode"
        and isinstance(n.get("inputs", {}).get("text"), str)
    ]
    return max(texts, key=len) if texts else ""


def _builder_default(sig: inspect.Signature, name: str, fallback):
    """Default declarado de un parametro del builder, o el fallback dado."""
    p = sig.parameters.get(name)
    if p is None or p.default is inspect.Parameter.empty:
        return fallback
    return p.default if isinstance(p.default, (int, float)) else fallback


def _refine_subject(subject: str, judge_prompt: str, reasons, model: str) -> str:
    """Reescribe el subject corrigiendo lo que el juez senalo, via ask_llm.

    Devuelve el subject mejorado (string corto) o el original si el LLM falla.
    """
    from core.ask_llm import ask_llm

    complaints = "; ".join(str(r) for r in (reasons or []) if r) or "(sin razones)"
    system = (
        "Eres un prompt-engineer de generacion de imagenes. Recibes el SUBJECT de "
        "una imagen rechazada por un juez de calidad y la lista de quejas del juez. "
        "Devuelve un SUBJECT mejorado y conciso (una frase, en ingles) que conserve la "
        "intencion original pero corrija las quejas anadiendo descriptores visuales "
        "concretos (p.ej. 'clean vector UI, sharp edges, high contrast, crisp lines' "
        "si era borroso). NO escribas explicaciones, NO uses comillas: responde SOLO "
        "con el subject mejorado."
    )
    user = (
        f"SUBJECT original: {subject}\n"
        f"Prompt completo generado: {judge_prompt}\n"
        f"Quejas del juez: {complaints}\n"
        "SUBJECT mejorado:"
    )
    try:
        out = ask_llm(user, model=model, system=system, echo=False)
        out = (out or "").strip().strip('"').strip()
        return out or subject
    except Exception:  # noqa: BLE001 — refine es best-effort; nunca rompe el loop.
        return subject


def comfyui_generate_until_quality(
    builder,
    subject: str,
    *,
    threshold: float = 6.0,
    clip_threshold: float = 0.24,
    max_iters: int = 4,
    strategy: str = "reroll+escalate+refine_prompt",
    server: str = "127.0.0.1:8188",
    dest_dir: str = "~/ComfyUI/output",
    judge_prompt: str | None = None,
    seed: int = 0,
    refine_model: str = "claude-haiku-4-5-20251001",
    judge_model: str = "claude-opus-4-8",
    wait_timeout: float = 300.0,
    **builder_kwargs,
) -> dict:
    """Genera y refina hasta alcanzar la calidad pedida (o agotar intentos).

    Args:
        builder: callable o nombre (str) de un builder `comfyui_build_*_workflow`
            del registry. El `subject` se pasa como PRIMER positional del builder
            (caso de los builders de asset: ui_hud, item_icon, enemy_creature...,
            cuyo primer arg es el elemento/sujeto).
        subject: descripcion del elemento a generar (p.ej. "RPG health and mana
            bars" para `comfyui_build_ui_hud_workflow`). Se inyecta en el builder
            y, si se refina, se reescribe con el feedback del juez.
        threshold: umbral estetico (0-10) que el juez usa para votar good/bad.
            keyword-only.
        clip_threshold: umbral de fidelidad CLIP (0-1) del juez. keyword-only.
        max_iters: numero maximo de iteraciones de generacion. keyword-only.
        strategy: combinacion de tacticas de mejora separadas por '+':
            'reroll' (seed nueva cada iter), 'escalate' (mas steps/cfg en iters
            tardias) y 'refine_prompt' (reescribe el subject con ask_llm usando
            las razones del juez). keyword-only.
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        dest_dir: directorio local donde guardar los PNG. keyword-only.
        judge_prompt: texto que se pasa al juez para medir fidelidad. Si None,
            se extrae el prompt positivo del workflow construido. keyword-only.
        seed: semilla base; los rerolls derivan de ella de forma determinista.
            keyword-only.
        refine_model: modelo de ask_llm para el refine del prompt (barato).
        judge_model: modelo del juez critico LLM-vision. keyword-only.
        wait_timeout: segundos maximos esperando cada generacion. keyword-only.
        **builder_kwargs: parametros extra del builder (ui_style, checkpoint,
            size, transparent...). Solo se pasan los que el builder acepta.

    Returns:
        dict {ok, converged, best_image_path, best_score, best_verdict,
        iterations, error}. `iterations` es una lista de
        {iter, seed, params, score, verdict, reasons, image, error}. `converged`
        True si alguna iteracion logro verdict 'good'. `best_*` apuntan a la
        candidata de mayor score (aunque ninguna convergiera). Si nada se pudo
        generar, ok=False y error explica.
    """
    parts = {p.strip() for p in str(strategy).split("+") if p.strip()}
    do_reroll = "reroll" in parts
    do_escalate = "escalate" in parts
    do_refine = "refine_prompt" in parts

    try:
        builder_fn = _resolve_builder(builder)
    except (ImportError, AttributeError, TypeError) as exc:
        return {
            "ok": False, "converged": False, "best_image_path": "",
            "best_score": None, "best_verdict": "", "iterations": [],
            "error": f"no se pudo resolver el builder: {exc}",
        }

    sig = inspect.signature(builder_fn)
    accepts = set(sig.parameters)
    base_steps = builder_kwargs.get("steps", _builder_default(sig, "steps", 28))
    base_cfg = builder_kwargs.get("cfg", _builder_default(sig, "cfg", 7.0))
    prefix = builder_kwargs.get("filename_prefix", "until_quality")

    dest = os.path.expanduser(dest_dir)
    subject_cur = subject
    iterations: list[dict] = []
    best: dict | None = None
    converged = False

    for i in range(max(1, int(max_iters))):
        # --- parametros de esta iteracion segun la estrategia ---
        cur_seed = (seed + i * _SEED_STRIDE) if do_reroll else seed
        kw = dict(builder_kwargs)
        if "seed" in accepts:
            kw["seed"] = cur_seed
        if do_escalate and i > 0:
            if "steps" in accepts:
                kw["steps"] = int(base_steps) + i * 8  # mas pasos = mas nitidez
            if "cfg" in accepts:
                kw["cfg"] = round(min(float(base_cfg) + i * 0.5, 12.0), 2)
        if "filename_prefix" in accepts:
            kw["filename_prefix"] = f"{prefix}_i{i}"
        # Solo pasamos kwargs que el builder acepta (evita TypeError entre builders).
        kw = {k: v for k, v in kw.items() if k in accepts}
        params = {
            "seed": cur_seed,
            "steps": kw.get("steps", base_steps),
            "cfg": kw.get("cfg", base_cfg),
            "subject": subject_cur,
        }

        rec = {"iter": i, "seed": cur_seed, "params": params, "score": None,
               "verdict": "", "reasons": [], "image": "", "error": ""}

        # --- build ---
        try:
            workflow = builder_fn(subject_cur, **kw)
        except Exception as exc:  # noqa: BLE001 — registra y reintenta siguiente iter.
            rec["error"] = f"build fallo: {exc}"
            iterations.append(rec)
            continue

        jp = judge_prompt if judge_prompt else _extract_positive_prompt(workflow)

        # --- submit ---
        try:
            sub = comfyui_submit_workflow(workflow, server=server)
            prompt_id = sub["prompt_id"]
        except (RuntimeError, KeyError) as exc:
            rec["error"] = f"submit fallo: {exc}"
            iterations.append(rec)
            continue

        # --- wait ---
        try:
            outputs = comfyui_wait_result(prompt_id, server=server, timeout=wait_timeout)
        except (TimeoutError, RuntimeError) as exc:
            rec["error"] = f"wait fallo: {exc}"
            iterations.append(rec)
            continue

        # --- localizar el PNG ---
        img = None
        for node_out in outputs.values():
            images = node_out.get("images") if isinstance(node_out, dict) else None
            if images:
                img = images[0]
                break
        if img is None:
            rec["error"] = f"el workflow no produjo imagenes (outputs={list(outputs)})"
            iterations.append(rec)
            continue

        # --- fetch ---
        fetched = comfyui_fetch_output_image(
            img["filename"], subfolder=img.get("subfolder", ""),
            type_=img.get("type", "output"), server=server, dest_dir=dest,
        )
        if not fetched.get("ok"):
            rec["error"] = f"fetch fallo: {fetched.get('error')}"
            iterations.append(rec)
            continue
        rec["image"] = fetched["path"]

        # --- judge (degrada con gracia si un juez cae) ---
        try:
            verdict = comfyui_judge_image(
                fetched["path"], jp, threshold=threshold,
                clip_threshold=clip_threshold, server=server, model=judge_model,
            )
        except Exception as exc:  # noqa: BLE001 — un juez caido no debe tumbar el loop.
            verdict = {"ok": False, "verdict": "unknown", "score": 0.0,
                       "reasons": [f"juez no disponible: {exc}"]}

        rec["score"] = float(verdict.get("score") or 0.0)
        rec["verdict"] = verdict.get("verdict", "unknown")
        rec["reasons"] = list(verdict.get("reasons") or [])
        iterations.append(rec)

        # --- best-of-N: guarda siempre la mejor por score ---
        if best is None or rec["score"] > best["score"]:
            best = rec

        # --- convergencia ---
        if rec["verdict"] == "good":
            converged = True
            break

        # --- refine para la siguiente iteracion ---
        if do_refine and i < max_iters - 1:
            subject_cur = _refine_subject(subject_cur, jp, rec["reasons"], refine_model)

    if best is None:
        last_err = iterations[-1]["error"] if iterations else "sin iteraciones"
        return {
            "ok": False, "converged": False, "best_image_path": "",
            "best_score": None, "best_verdict": "", "iterations": iterations,
            "error": f"ninguna iteracion produjo imagen ({last_err})",
        }

    return {
        "ok": True,
        "converged": converged,
        "best_image_path": best["image"],
        "best_score": best["score"],
        "best_verdict": best["verdict"],
        "iterations": iterations,
        "error": "",
    }


if __name__ == "__main__":
    import json

    res = comfyui_generate_until_quality(
        "comfyui_build_ui_hud_workflow",
        "RPG health and mana bars, clean game UI",
        ui_style="fantasy game UI, clean vector, high contrast",
        threshold=6.5,
        max_iters=3,
        dest_dir="/tmp/comfy_until_quality",
        transparent=False,
    )
    print(json.dumps(res, indent=2))