"""Importa un workflow ComfyUI desde una URL (http/https) o un path local.

Detecta el formato:
- API format: dict {node_id: {class_type, inputs}} -> se devuelve tal cual.
- UI graph: dict {nodes, links, ...} (lo que exporta "Save" en la UI) -> se
  normaliza a API format. La normalizacion de los valores de widget necesita el
  catalogo /object_info del servidor; si el servidor responde, los widgets se
  mapean por nombre; si no, solo se conservan las conexiones entre nodos.

Compone comfyui_object_info para el mapeo de widgets del UI graph.

Impura: red (HTTP GET si source es URL) + lectura de disco. Solo stdlib.
"""
import json
import os
import sys
import urllib.error
import urllib.request

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

from comfyui_object_info import comfyui_object_info  # noqa: E402


def comfyui_import_workflow_json(
    source: str,
    *,
    server: str = "127.0.0.1:8188",
    timeout: float = 15.0,
) -> dict:
    """Lee un workflow JSON y lo normaliza a API format.

    Args:
        source: URL http(s) de un JSON de workflow, o ruta de un archivo local.
        server: host:port de ComfyUI usado solo para mapear los widgets cuando
            la fuente viene en formato UI graph. keyword-only.
        timeout: timeout HTTP en segundos. keyword-only.

    Returns:
        dict {ok, workflow, format_detected, error}. format_detected es "api",
        "ui_graph" o "". Si falla, ok=False y error explica el motivo.
    """
    try:
        if source.startswith(("http://", "https://")):
            with urllib.request.urlopen(source, timeout=timeout) as resp:
                raw = resp.read()
        else:
            with open(source, "rb") as f:
                raw = f.read()
        data = json.loads(raw)
    except (urllib.error.URLError, OSError) as exc:
        return {"ok": False, "workflow": {}, "format_detected": "",
                "error": f"no se pudo leer {source!r}: {exc}"}
    except json.JSONDecodeError as exc:
        return {"ok": False, "workflow": {}, "format_detected": "",
                "error": f"JSON invalido en {source!r}: {exc}"}

    if not isinstance(data, dict):
        return {"ok": False, "workflow": {}, "format_detected": "",
                "error": "el JSON no es un objeto de workflow"}

    # API format: todos los valores son dicts con class_type
    if data and all(isinstance(v, dict) and "class_type" in v for v in data.values()):
        return {"ok": True, "workflow": data, "format_detected": "api", "error": ""}

    # UI graph: tiene la clave "nodes"
    if "nodes" in data:
        obj_info = None
        try:
            obj_info = comfyui_object_info(server=server, timeout=min(timeout, 5.0))
        except Exception:
            obj_info = None
        api = _ui_graph_to_api(data, obj_info)
        return {"ok": True, "workflow": api, "format_detected": "ui_graph", "error": ""}

    return {"ok": False, "workflow": {}, "format_detected": "",
            "error": "formato de workflow no reconocido (ni API ni UI graph)"}


# Node types virtuales del editor de ComfyUI: solo existen en el UI graph y se
# descartan al pasar UI -> API (ComfyUI hace lo mismo). Note/MarkdownNote son
# anotaciones; PrimitiveNode inyecta un valor de widget; Reroute es un passthrough
# de una conexion (se resuelve reconectando origen real -> destino).
_NOTE_TYPES = {"Note", "MarkdownNote"}


def _is_reroute(ctype) -> bool:
    """True si el node type es un Reroute (nativo 'Reroute' o variantes custom)."""
    return isinstance(ctype, str) and ctype.startswith("Reroute")


def _is_virtual(ctype) -> bool:
    """True si el node type es virtual del editor (no va al API format)."""
    return ctype in _NOTE_TYPES or ctype == "PrimitiveNode" or _is_reroute(ctype)


def _resolve_source(src_node, src_slot, node_by_id, link_src, _depth=0):
    """Resuelve el origen real de una conexion saltando los nodos Reroute.

    Un Reroute en el UI graph es un passthrough: su salida solo reenvia lo que
    llega a su unica entrada. Para producir API format hay que reconectar el
    consumidor directamente al origen real (origen -> destino, sin el Reroute).
    Devuelve (node_id, slot) del nodo no-Reroute al que se conecta, o None si la
    cadena de Reroutes esta rota (entrada sin link) o forma un ciclo.
    """
    if _depth > 64:
        return None  # ciclo de Reroutes: aborta la resolucion.
    node = node_by_id.get(src_node)
    if node is None or not _is_reroute(node.get("type")):
        return (src_node, src_slot)
    link = None
    for inp in node.get("inputs", []) or []:
        if inp.get("link") is not None:
            link = inp["link"]
            break
    if link is None or link not in link_src:
        return None  # Reroute sin entrada conectada: link muerto.
    nxt_node, nxt_slot = link_src[link]
    return _resolve_source(nxt_node, nxt_slot, node_by_id, link_src, _depth + 1)


def _ui_graph_to_api(graph: dict, obj_info) -> dict:
    """Convierte un UI graph de ComfyUI a API format (best-effort).

    Omite los nodos virtuales del editor (Note, MarkdownNote, PrimitiveNode,
    Reroute) tal como hace ComfyUI al pasar de UI a API: las anotaciones se
    descartan, los Reroute se resuelven reconectando la conexion directa
    origen->destino, y los PrimitiveNode se inyectan como valor de widget en el
    consumidor.
    """
    nodes = graph.get("nodes", []) or []
    links = graph.get("links", []) or []
    # link_id -> (src_node_id, src_slot)
    link_src = {}
    for lk in links:
        if isinstance(lk, list) and len(lk) >= 5:
            link_src[lk[0]] = (str(lk[1]), lk[2])
    # node_id (str) -> node dict, para TODOS los nodos (incluidos los virtuales),
    # necesario para resolver Reroutes e inyectar valores de PrimitiveNode.
    node_by_id = {str(n.get("id")): n for n in nodes if n.get("id") is not None}

    api = {}
    for node in nodes:
        ctype = node.get("type")
        if ctype is None or _is_virtual(ctype):
            continue  # los virtuales no existen en API format.
        nid = str(node.get("id"))
        inputs = {}
        connected = set()
        for inp in node.get("inputs", []) or []:
            name = inp.get("name")
            link = inp.get("link")
            if name is None or link is None or link not in link_src:
                continue
            src_node, src_slot = link_src[link]
            resolved = _resolve_source(src_node, src_slot, node_by_id, link_src)
            if resolved is None:
                continue  # cadena de Reroutes rota: el input queda sin conexion.
            rnode, rslot = resolved
            src = node_by_id.get(rnode)
            if src is not None and src.get("type") == "PrimitiveNode":
                # PrimitiveNode: inyecta su valor constante como widget, no como link.
                wv = src.get("widgets_values")
                if isinstance(wv, list) and wv:
                    inputs[name] = wv[0]
                    connected.add(name)
                continue
            inputs[name] = [rnode, rslot]
            connected.add(name)
        widgets = node.get("widgets_values")
        if isinstance(widgets, dict):
            inputs.update(widgets)
        elif isinstance(widgets, list) and widgets:
            for name, val in zip(_widget_input_names(ctype, obj_info, connected), widgets):
                inputs[name] = val
        api[nid] = {"class_type": ctype, "inputs": inputs}
    return api


def _widget_input_names(ctype, obj_info, connected) -> list:
    """Nombres de inputs que son widgets (no conexiones), en orden, via object_info."""
    if not obj_info or ctype not in obj_info:
        return []
    spec = obj_info[ctype].get("input", {})
    names = []
    for section in ("required", "optional"):
        for name, decl in (spec.get(section) or {}).items():
            if name in connected:
                continue
            t = decl[0] if isinstance(decl, list) and decl else decl
            if isinstance(t, list):
                names.append(name)  # combo/enum => widget
            elif t in ("INT", "FLOAT", "STRING", "BOOLEAN"):
                names.append(name)
    return names


if __name__ == "__main__":
    res = comfyui_import_workflow_json("/tmp/does_not_exist.json")
    print(res)
