"""comfyui_run_foreign_workflow_oneshot — ejecuta un workflow ComfyUI ajeno end-to-end.

Promocion de la secuencia del roadmap (issue 0064/0087): traer un workflow
foraneo (URL de Drive/GitHub/Civitai/HuggingFace, PNG/JSON local, o dict en API
format) y ejecutarlo en NUESTRO servidor ComfyUI en una sola llamada, resolviendo
antes sus dependencias para no fallar a ciegas. Compone funciones del registry del
grupo `comfyui`:

    comfyui_download_workflow_py_ml          (cualquier fuente -> API format)
    comfyui_resolve_workflow_deps_py_ml      (detecta nodos/modelos faltantes)
    comfyui_install_custom_node_py_ml        (instala nodos custom, opt-in + confiable)
    comfyui_validate_workflow_py_ml          (revalida tras instalar)
    comfyui_submit_workflow_py_ml            (POST /prompt)
    comfyui_wait_result_py_ml                (poll /history)
    comfyui_fetch_output_image_py_ml         (GET /view -> disco, imagenes)
    comfyui_fetch_output_video_py_ml         (GET /view -> disco, video/animacion)
    comfyui_fetch_output_mesh_py_ml          (GET /view -> disco, mallas 3D)

SEGURIDAD: un workflow foraneo es codigo de terceros que se ejecuta como grafo en
nuestro servidor. Por eso el pipeline SIEMPRE resuelve dependencias antes del
submit y NUNCA instala nada a ciegas:
  - Modelos faltantes -> se REPORTAN en `missing`, jamas se descargan
    automaticamente (no hay forma de saber si la fuente es confiable).
  - Nodos custom faltantes -> solo se instalan si install_nodes=True Y el caller
    pasa su URL de repo en `node_repos` (fuente confiable explicita). Tras
    instalar, ComfyUI debe reiniciarse para cargarlos, asi que el pipeline NO
    completa el submit en esa misma llamada: devuelve ok=False con la instruccion
    de reiniciar y reintentar.
Si tras resolver quedan dependencias faltantes, el pipeline NO encola: devuelve
ok=False con `missing` poblado.

Pipeline impuro: red (HTTP) + escritura en disco (+ subprocess git/pip si se
instalan nodos custom).
"""

from __future__ import annotations

import json
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_download_workflow import comfyui_download_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_fetch_output_video import comfyui_fetch_output_video
from ml.comfyui_install_custom_node import comfyui_install_custom_node
from ml.comfyui_resolve_workflow_deps import comfyui_resolve_workflow_deps
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_validate_workflow import comfyui_validate_workflow
from ml.comfyui_wait_result import comfyui_wait_result

# Clasificacion de los outputs por extension del filename.
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tiff")
_VIDEO_EXTS = (".mp4", ".webm", ".webp", ".gif", ".mkv", ".mov", ".avif")
_MESH_EXTS = (".glb", ".gltf", ".obj", ".ply", ".fbx", ".stl", ".usdz", ".splat")


def _classify_outputs(outputs: dict) -> dict:
    """Reparte los items de los outputs de /history en image/video/mesh por extension.

    Devuelve {"image": [items], "video": [items], "mesh": [items]} donde cada
    item es {filename, subfolder, type}. Un .webp se considera video/animacion
    (los nodos de animacion de ComfyUI lo usan); las imagenes estaticas usan png/jpg.
    """
    buckets: dict[str, list] = {"image": [], "video": [], "mesh": []}
    for node_out in outputs.values():
        if not isinstance(node_out, dict):
            continue
        for items in node_out.values():
            if not isinstance(items, list):
                continue
            for item in items:
                if not isinstance(item, dict):
                    continue
                fn = (item.get("filename") or "").lower()
                rec = {
                    "filename": item.get("filename", ""),
                    "subfolder": item.get("subfolder", ""),
                    "type": item.get("type", "output"),
                }
                if fn.endswith(_MESH_EXTS):
                    buckets["mesh"].append(rec)
                elif fn.endswith(_VIDEO_EXTS):
                    buckets["video"].append(rec)
                elif fn.endswith(_IMAGE_EXTS):
                    buckets["image"].append(rec)
    return buckets


def _resolve_workflow(source, server, civitai_token, hf_token):
    """Resuelve el `source` (dict | str) a un workflow en API format.

    Devuelve (workflow, source_type, error). Si source ya es un dict en API
    format se usa tal cual; si es str se delega a comfyui_download_workflow.
    """
    if isinstance(source, dict):
        if source and all(
            isinstance(v, dict) and "class_type" in v for v in source.values()
        ):
            return source, "dict", ""
        return {}, "dict", "el dict pasado no esta en API format ({node_id: {class_type, inputs}})"
    if not isinstance(source, str):
        return {}, "", f"source debe ser str (URL/path) o dict (API format), no {type(source).__name__}"
    dl = comfyui_download_workflow(
        source, server=server, civitai_token=civitai_token, hf_token=hf_token
    )
    if not dl.get("ok"):
        return {}, dl.get("source_type", ""), dl.get("error", "no se pudo importar el workflow")
    return dl.get("workflow", {}), dl.get("source_type", ""), ""


def comfyui_run_foreign_workflow_oneshot(
    source,
    *,
    server: str = "127.0.0.1:8188",
    dest: str | None = None,
    output_kind: str = "auto",
    install_nodes: bool = False,
    node_repos: dict | None = None,
    wait_timeout: float = 600.0,
    civitai_token: str | None = None,
    hf_token: str | None = None,
) -> dict:
    """Importa, valida, ejecuta y descarga los outputs de un workflow ComfyUI ajeno.

    Args:
        source: workflow de entrada. Puede ser una URL (Google Drive, GitHub,
            Civitai, HuggingFace o directa a .json/.png/.webp), una ruta local
            (.json/.png), o un dict ya en API format ({node_id: {class_type,
            inputs}}).
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        dest: directorio destino de los outputs descargados. None = cwd.
            keyword-only.
        output_kind: que outputs descargar: "auto" (todos: imagenes + el primer
            video + la primera malla), "image", "video" o "mesh". keyword-only.
        install_nodes: si True, instala los nodos custom faltantes cuyo class_type
            este mapeado a una URL de repo confiable en `node_repos`. Tras instalar
            hay que reiniciar ComfyUI para cargarlos, asi que el pipeline NO
            completa el submit en esa llamada. Por defecto False. keyword-only.
        node_repos: mapa {class_type: repo_url} con las URLs de repo confiables de
            los nodos custom que se permite instalar. Solo se usa si
            install_nodes=True. Los modelos faltantes NUNCA se instalan. keyword-only.
        wait_timeout: segundos maximos esperando a que el server termine.
            keyword-only.
        civitai_token / hf_token: tokens opcionales para fuentes gated. keyword-only.

    Returns:
        dict {ok, prompt_id, outputs, missing, source_type, error}:
          - outputs: lista de rutas locales de los archivos descargados.
          - missing: lista de dependencias faltantes (suggestions de
            comfyui_resolve_workflow_deps): cada una {kind: "node"|"model", name,
            action, hint, ...}. No vacia cuando el workflow pide algo que no
            tenemos; en ese caso ok=False y NO se encola.
          - source_type: de donde vino el workflow ("github", "drive", "local",
            "dict", ...).
        Si falla en cualquier paso, ok=False y error explica el motivo. Nunca
        instala modelos ni nodos de fuentes no confiables a ciegas.
    """
    if output_kind not in ("auto", "image", "video", "mesh"):
        return {"ok": False, "prompt_id": "", "outputs": [], "missing": [],
                "source_type": "", "error": f"output_kind {output_kind!r} no valido"}

    # 1. Importar el workflow a API format desde cualquier fuente.
    workflow, source_type, err = _resolve_workflow(source, server, civitai_token, hf_token)
    if err:
        return {"ok": False, "prompt_id": "", "outputs": [], "missing": [],
                "source_type": source_type, "error": f"import fallo: {err}"}
    if not workflow:
        return {"ok": False, "prompt_id": "", "outputs": [], "missing": [],
                "source_type": source_type, "error": "el workflow importado esta vacio"}

    # 2. Resolver dependencias (nodos/modelos faltantes) ANTES de encolar.
    deps = comfyui_resolve_workflow_deps(workflow, server=server)
    if not deps.get("ok"):
        return {"ok": False, "prompt_id": "", "outputs": [], "missing": [],
                "source_type": source_type,
                "error": f"no se pudieron resolver dependencias (¿server vivo?): {deps.get('error')}"}

    missing_nodes = list(deps.get("missing_nodes", []))
    missing_models = list(deps.get("missing_models", []))
    suggestions = list(deps.get("suggestions", []))

    # 3. Instalar SOLO nodos custom confiables (opt-in + URL provista por el caller).
    install_notes = []
    repos = node_repos or {}
    if install_nodes and missing_nodes and repos:
        for ctype in list(missing_nodes):
            repo = repos.get(ctype)
            if not repo:
                continue  # sin URL confiable -> no se instala, queda en missing.
            res = comfyui_install_custom_node(repo, restart=False)
            if res.get("ok"):
                install_notes.append(f"nodo '{ctype}' instalado desde {repo}")
            else:
                install_notes.append(f"fallo instalando '{ctype}': {res.get('error')}")

    # 4. Si quedan dependencias faltantes (modelos, o nodos sin cargar) -> NO encolar.
    #    Los nodos recien instalados requieren reiniciar ComfyUI; no se cargan ahora.
    if missing_nodes or missing_models:
        note = "dependencias faltantes; no se encola"
        if install_notes:
            note += ". " + "; ".join(install_notes)
            note += ". Reinicia ComfyUI (cuando no haya generaciones en curso) y reintenta"
        return {"ok": False, "prompt_id": "", "outputs": [], "missing": suggestions,
                "source_type": source_type, "error": note}

    # 5. Encolar el workflow.
    try:
        sub = comfyui_submit_workflow(workflow, server=server)
        prompt_id = sub["prompt_id"]
    except (RuntimeError, KeyError) as exc:
        return {"ok": False, "prompt_id": "", "outputs": [], "missing": [],
                "source_type": source_type, "error": f"submit fallo: {exc}"}

    # 6. Esperar a que termine.
    try:
        outputs = comfyui_wait_result(prompt_id, server=server, timeout=wait_timeout)
    except (TimeoutError, RuntimeError) as exc:
        return {"ok": False, "prompt_id": prompt_id, "outputs": [], "missing": [],
                "source_type": source_type, "error": f"wait fallo: {exc}"}

    # 7. Descargar los outputs segun output_kind. `dest` se trata SIEMPRE como
    #    directorio (el workflow puede producir varios archivos): se crea de
    #    antemano para que fetch_video/fetch_mesh escriban el basename dentro
    #    (si el dir no existe, esos helpers interpretarian la ruta como archivo).
    out_dir = dest or "."
    try:
        os.makedirs(out_dir, exist_ok=True)
    except OSError as exc:
        return {"ok": False, "prompt_id": prompt_id, "outputs": [], "missing": [],
                "source_type": source_type, "error": f"no se pudo crear dest {out_dir!r}: {exc}"}

    buckets = _classify_outputs(outputs)
    paths: list[str] = []
    want_image = output_kind in ("auto", "image")
    want_video = output_kind in ("auto", "video")
    want_mesh = output_kind in ("auto", "mesh")

    if want_image:
        for item in buckets["image"]:
            r = comfyui_fetch_output_image(
                item["filename"], subfolder=item["subfolder"],
                type_=item["type"], server=server, dest_dir=out_dir,
            )
            if r.get("ok"):
                paths.append(r["path"])
    if want_video and buckets["video"]:
        r = comfyui_fetch_output_video(prompt_id, server=server, dest=out_dir, outputs=outputs)
        if r.get("ok"):
            paths.append(r["path"])
    if want_mesh and buckets["mesh"]:
        r = comfyui_fetch_output_mesh(prompt_id, server=server, dest=out_dir)
        if r.get("ok"):
            paths.append(r["path"])

    if not paths:
        return {"ok": False, "prompt_id": prompt_id, "outputs": [], "missing": [],
                "source_type": source_type,
                "error": f"el workflow termino pero no se descargo ningun output de tipo {output_kind!r}"}

    return {"ok": True, "prompt_id": prompt_id, "outputs": paths, "missing": [],
            "source_type": source_type, "error": ""}


if __name__ == "__main__":
    # Smoke: ejecuta un workflow ya validado de la libreria local (deps presentes).
    src = sys.argv[1] if len(sys.argv) > 1 else os.path.expanduser(
        "~/ComfyUI/workflows_library/txt2img_flux_schnell.api.json")
    res = comfyui_run_foreign_workflow_oneshot(src, dest="/tmp/comfy_foreign")
    print(json.dumps(res, indent=2))
