"""comfyui_text_to_3d_oneshot — prompt de texto -> malla 3D GLB en una sola llamada.

Promocion de la secuencia repetida (issue 0087): generar una imagen desde un
prompt con Stable Diffusion -> reconstruir una malla 3D desde esa imagen. Antes
eran dos invocaciones encadenadas a mano (build_txt2img + submit + wait + fetch
imagen, luego build_image_to_3d + submit + wait + fetch malla). Aqui se compone
en una sola llamada las funciones del registry del grupo `comfyui`:

    comfyui_build_txt2img_workflow_py_ml                 (imagen desde texto)
    comfyui_submit_workflow_py_ml                        (POST /prompt)
    comfyui_wait_result_py_ml                            (poll /history)
    comfyui_fetch_output_image_py_ml                     (GET /view -> disco)
    comfyui_build_image_to_3d_workflow_py_ml             (malla nativa Hunyuan3D)
    comfyui_build_textured_3d_multiview_workflow_py_ml   (malla texturizada PBR)
    comfyui_fetch_output_mesh_py_ml                      (GET /view -> disco)

La imagen generada queda en el output/ del servidor; para alimentar el paso 3D
(cuyo nodo LoadImage lee del input/) se sube al input/ reutilizando el helper
`_upload_image` del pipeline hermano comfyui_image_to_3d_oneshot (POST
/upload/image, no se reescribe la logica multipart).

Pipeline impuro: red (HTTP) + escritura en disco.
"""

from __future__ import annotations

import json
import os
import sys
import tempfile

# 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_build_image_to_3d_workflow import comfyui_build_image_to_3d_workflow
from ml.comfyui_build_textured_3d_multiview_workflow import (
    comfyui_build_textured_3d_multiview_workflow,
)
from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
from ml.comfyui_fetch_output_image import comfyui_fetch_output_image
from ml.comfyui_fetch_output_mesh import comfyui_fetch_output_mesh
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_wait_result import comfyui_wait_result
from pipelines.comfyui_image_to_3d_oneshot import _upload_image

# variant -> nombre del checkpoint Hunyuan3D-2 para el paso 3D no texturizado.
_VARIANT_CKPT = {
    "mini": "3D_hunyuan3d-dit-v2-mini.safetensors",
    "standard": "3D_hunyuan3d-dit-v2-0.safetensors",
    "mv": "3D_hunyuan3d-dit-v2-mv.safetensors",
}

# Claves de **gen que reenviamos a comfyui_build_txt2img_workflow; cualquier otra
# clave en **gen se ignora silenciosamente (el paso 3D usa sus defaults).
_TXT2IMG_KEYS = {
    "steps", "cfg", "width", "height", "seed",
    "sampler_name", "scheduler", "filename_prefix",
}


def _first_image(outputs: dict) -> dict | None:
    """Devuelve la primera imagen de los outputs de un prompt (o None)."""
    for out in outputs.values():
        for img in out.get("images", []) or []:
            if img.get("filename"):
                return img
    return None


def _image_dir(dest: str | None) -> str:
    """Resuelve el directorio donde guardar la imagen intermedia.

    dest None -> directorio temporal del sistema; dest con extension de archivo
    -> su carpeta; dest sin extension -> se trata como carpeta.
    """
    if not dest:
        return os.path.join(tempfile.gettempdir(), "comfy_text_to_3d")
    dest = os.path.expanduser(dest)
    _, ext = os.path.splitext(dest)
    if ext:
        return os.path.dirname(dest) or "."
    return dest


def comfyui_text_to_3d_oneshot(
    prompt: str,
    *,
    server: str = "127.0.0.1:8188",
    ckpt_name: str = "IMG_v1-5-pruned-emaonly-fp16.safetensors",
    negative: str = "",
    textured: bool = False,
    variant: str = "mini",
    dest: str | None = None,
    image_wait_timeout: float = 300.0,
    wait_timeout: float = 600.0,
    **gen,
) -> dict:
    """Genera una imagen desde texto y reconstruye su malla 3D, end-to-end.

    Args:
        prompt: prompt positivo para la generacion de la imagen (txt2img).
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        ckpt_name: checkpoint Stable Diffusion para el txt2img (default SD1.5).
            keyword-only.
        negative: prompt negativo del txt2img. keyword-only.
        textured: si True usa el pipeline 3D texturizado PBR multi-vista
            (comfyui_build_textured_3d_multiview_workflow, requiere el wrapper
            Hunyuan3DWrapper); si False usa la malla nativa Hunyuan3D-2 segun
            `variant`. keyword-only.
        variant: "mini" (default), "standard" o "mv"; checkpoint Hunyuan3D-2 del
            paso 3D no texturizado (ignorado si textured=True). keyword-only.
        dest: ruta destino de la malla (None = cwd; dir = dentro; archivo = ahi).
            La imagen intermedia se guarda en una carpeta derivada de dest.
            keyword-only.
        image_wait_timeout: segundos maximos esperando la generacion de la
            imagen. keyword-only.
        wait_timeout: segundos maximos esperando la reconstruccion 3D.
            keyword-only.
        **gen: parametros del txt2img (steps, cfg, width, height, seed,
            sampler_name, scheduler, filename_prefix); el resto se ignora.

    Returns:
        dict {ok, image_path, mesh_path, prompt_ids, error}. image_path = ruta
        local del PNG generado; mesh_path = ruta local de la malla; prompt_ids =
        [id_txt2img, id_3d] de los trabajos encolados. Si falla, ok=False y error
        explica en que paso.
    """
    out = {"ok": False, "image_path": "", "mesh_path": "",
           "prompt_ids": [], "error": ""}

    if not textured and variant not in _VARIANT_CKPT:
        out["error"] = f"variant {variant!r} no valida; usa {sorted(_VARIANT_CKPT)}"
        return out

    # 1. Construir el workflow txt2img (funcion pura del registry).
    t2i_gen = {k: v for k, v in gen.items() if k in _TXT2IMG_KEYS}
    wf_img = comfyui_build_txt2img_workflow(ckpt_name, prompt, negative, **t2i_gen)

    # 2. Encolar el txt2img.
    try:
        sub = comfyui_submit_workflow(wf_img, server=server)
        pid_img = sub["prompt_id"]
    except (RuntimeError, KeyError) as exc:
        out["error"] = f"submit txt2img fallo: {exc}"
        return out
    out["prompt_ids"].append(pid_img)

    # 3. Esperar la imagen y localizarla en los outputs.
    try:
        outputs = comfyui_wait_result(pid_img, server=server, timeout=image_wait_timeout)
    except (TimeoutError, RuntimeError) as exc:
        out["error"] = f"wait txt2img fallo: {exc}"
        return out
    img = _first_image(outputs)
    if img is None:
        out["error"] = "txt2img no produjo ninguna imagen en los outputs"
        return out

    # 4. Descargar la imagen generada a disco local.
    fetched_img = comfyui_fetch_output_image(
        img["filename"],
        subfolder=img.get("subfolder", ""),
        type_=img.get("type", "output"),
        server=server,
        dest_dir=_image_dir(dest),
    )
    if not fetched_img.get("ok"):
        out["error"] = f"fetch de imagen fallo: {fetched_img.get('error')}"
        return out
    out["image_path"] = fetched_img["path"]

    # 5. Subir la imagen al input/ del servidor para el paso 3D (LoadImage).
    try:
        image_name = _upload_image(out["image_path"], server)
    except RuntimeError as exc:
        out["error"] = f"upload de imagen al input/ fallo: {exc}"
        return out

    # 6. Construir el workflow imagen->3D (texturizado multi-vista o malla nativa).
    if textured:
        wf_3d = comfyui_build_textured_3d_multiview_workflow(image_name)
    else:
        wf_3d = comfyui_build_image_to_3d_workflow(image_name, _VARIANT_CKPT[variant])

    # 7. Encolar el 3D.
    try:
        sub3 = comfyui_submit_workflow(wf_3d, server=server)
        pid_3d = sub3["prompt_id"]
    except (RuntimeError, KeyError) as exc:
        out["error"] = f"submit 3D fallo: {exc}"
        return out
    out["prompt_ids"].append(pid_3d)

    # 8. Esperar la reconstruccion 3D.
    try:
        comfyui_wait_result(pid_3d, server=server, timeout=wait_timeout)
    except (TimeoutError, RuntimeError) as exc:
        out["error"] = f"wait 3D fallo: {exc}"
        return out

    # 9. Descargar la malla.
    fetched_mesh = comfyui_fetch_output_mesh(pid_3d, server=server, dest=dest)
    if not fetched_mesh.get("ok"):
        out["error"] = f"fetch de malla fallo: {fetched_mesh.get('error')}"
        return out
    out["mesh_path"] = fetched_mesh["path"]

    out["ok"] = True
    return out


if __name__ == "__main__":
    res = comfyui_text_to_3d_oneshot(
        "a cute robot toy, full body, centered, plain background, sharp focus",
        dest="/tmp/comfy_text_to_3d",
        steps=20,
        seed=42,
    )
    print(json.dumps(res, indent=2))
