#!/usr/bin/env python3
"""Convert a ComfyUI workflow from UI format (Save/Workflow menu) to API format.

Equivalent to the frontend's Workflow > Export (API), but headless. Needs a
running ComfyUI instance because widget-to-input mapping comes from its
GET /object_info (so the instance must have the same custom nodes installed
as the workflow uses).

Handles: subgraphs (flattened, ids become "<instance>:<inner>"), Reroute
splicing, bypassed (mode 4) nodes with type-matched pass-through, muted
(mode 2) nodes, widget-converted-to-input entries, and the phantom
control_after_generate value that follows every seed widget.

Example:
    python convert_workflow.py --url http://127.0.0.1:7865 \
        ../head-swap-flux-klein.json -o ../api/head-swap-flux-klein.json
"""
import argparse
import json
import pathlib
import sys

import requests

VIRTUAL_TYPES = {"Reroute", "Note", "MarkdownNote", "PrimitiveNode"}
WIDGET_SCALAR_TYPES = {"INT", "FLOAT", "STRING", "BOOLEAN", "COMBO"}
SUBGRAPH_INPUT_ID = -10
SUBGRAPH_OUTPUT_ID = -20


class WidgetDefault:
    """Literal carried across an unwired promoted subgraph input."""

    def __init__(self, value):
        self.value = value


class Scope:
    """One graph namespace: either the root workflow or a subgraph instance."""

    def __init__(self, graph, prefix="", parent=None, instance_id=None):
        self.nodes = {n["id"]: n for n in graph.get("nodes", [])}
        self.links = {}
        for l in graph.get("links", []):
            if isinstance(l, dict):  # subgraph defs store links as objects
                self.links[l["id"]] = (l["origin_id"], l["origin_slot"], l["target_id"], l["target_slot"])
            else:  # root stores [id, src, src_slot, dst, dst_slot, type]
                self.links[l[0]] = (l[1], l[2], l[3], l[4])
        self.prefix = prefix
        self.parent = parent            # (parent Scope, instance node id) or None
        self.instance_id = instance_id
        self.children = {}              # instance node id -> child Scope

    def fid(self, node_id):
        return f"{self.prefix}{node_id}"


class Converter:
    def __init__(self, workflow, object_info):
        self.oi = object_info
        self.subgraph_defs = {
            sg["id"]: sg for sg in workflow.get("definitions", {}).get("subgraphs", [])
        }
        self.root = Scope(workflow)
        self._build_children(self.root)

    def _build_children(self, scope):
        for nid, node in scope.nodes.items():
            sg = self.subgraph_defs.get(node["type"])
            if sg is not None:
                child = Scope(sg, prefix=f"{scope.fid(nid)}:", parent=(scope, nid), instance_id=nid)
                scope.children[nid] = child
                self._build_children(child)

    def all_scopes(self):
        stack = [self.root]
        while stack:
            s = stack.pop()
            yield s
            stack.extend(s.children.values())

    def resolve(self, scope, link_id):
        """Follow a link upstream to a real executable (node_fid, output_slot), or None."""
        if link_id is None or link_id not in scope.links:
            return None
        origin_id, origin_slot, _, _ = scope.links[link_id]

        if origin_id == SUBGRAPH_INPUT_ID:  # crossed into this subgraph from outside
            parent_scope, instance_id = scope.parent
            instance = parent_scope.nodes[instance_id]
            # instances serialize only a subset of the def's inputs (widget-promoted
            # ones live in proxyWidgets instead), so match by name, not slot index
            name = self.subgraph_defs[instance["type"]]["inputs"][origin_slot]["name"]
            entry = next((i for i in instance.get("inputs", []) if i["name"] == name), None)
            if entry is not None and entry.get("link") is not None:
                return self.resolve(parent_scope, entry["link"])
            # unwired promoted widget: the UI backs it with the widget of the
            # first interior node this boundary slot fans out to, so inputs
            # without a widget of their own (e.g. math variables) still get a
            # value -- mirror that lookup here
            return self._boundary_default(scope, origin_slot)

        node = scope.nodes[origin_id]
        ntype = node["type"]

        if ntype in self.subgraph_defs:  # source is a subgraph instance: descend
            child = scope.children[origin_id]
            for lid, (_, _, tid, tslot) in child.links.items():
                if tid == SUBGRAPH_OUTPUT_ID and tslot == origin_slot:
                    return self.resolve(child, lid)
            return None

        if ntype == "Reroute":
            return self.resolve(scope, node["inputs"][0].get("link"))

        mode = node.get("mode", 0)
        if mode == 2:  # muted: acts as removed
            return None
        if mode == 4:  # bypassed: pass through first linked input of matching type
            out_type = node["outputs"][origin_slot].get("type", "*")
            for inp in node.get("inputs", []):
                if inp.get("link") is None:
                    continue
                in_type = inp.get("type", "*")
                if in_type == out_type or "*" in (in_type, out_type):
                    return self.resolve(scope, inp["link"])
            return None

        return (scope.fid(origin_id), origin_slot)

    def _boundary_default(self, scope, origin_slot):
        """Widget value backing an unwired promoted subgraph input, taken from
        the first fan-out target of the slot that has that widget itself."""
        for _, (oid, oslot, tid, tslot) in sorted(scope.links.items()):
            if oid != SUBGRAPH_INPUT_ID or oslot != origin_slot:
                continue
            target = scope.nodes.get(tid)
            if (target is None or target.get("mode", 0) in (2, 4)
                    or target["type"] in self.subgraph_defs
                    or target["type"] not in self.oi):
                continue
            if tslot >= len(target.get("inputs", [])):
                continue
            widgets = self.widget_inputs(target)
            name = target["inputs"][tslot]["name"]
            if name in widgets:
                return WidgetDefault(widgets[name])
        return None

    def widget_inputs(self, node):
        """Map a node's widgets_values onto named inputs using /object_info order."""
        if "Power Lora Loader" in node["type"]:
            # rgthree serializes lora rows as dicts among placeholder widgets;
            # the API expects them as lora_1..lora_N dict inputs
            return {
                f"lora_{i}": row
                for i, row in enumerate(
                    (v for v in node.get("widgets_values") or []
                     if isinstance(v, dict) and "lora" in v),
                    start=1,
                )
            }
        spec = self.oi[node["type"]].get("input", {})
        ordered = list(spec.get("required", {}).items()) + list(spec.get("optional", {}).items())
        values = node.get("widgets_values") or []
        if isinstance(values, dict):  # some packs (rgthree) serialize by name
            widget_names = {name for name, s in ordered if self._is_widget(s)}
            return {k: v for k, v in values.items() if k in widget_names}
        result = {}
        self._consume_widgets(ordered, list(values), result)
        return result

    def _consume_widgets(self, ordered, queue, result, prefix=""):
        for name, s in ordered:
            if not queue:
                return
            opts = s[1] if len(s) > 1 and isinstance(s[1], dict) else {}
            if s[0] == "COMFY_DYNAMICCOMBO_V3":
                # the selection is one widget value; the chosen option pulls in
                # nested inputs addressed with dotted ids ("resize_type.megapixels")
                selected = queue.pop(0)
                result[f"{prefix}{name}"] = selected
                option = next((o for o in opts.get("options", []) if o["key"] == selected), None)
                if option:
                    inner = option["inputs"]
                    inner_ordered = (list(inner.get("required", {}).items())
                                     + list(inner.get("optional", {}).items()))
                    self._consume_widgets(inner_ordered, queue, result, prefix=f"{prefix}{name}.")
            elif self._is_widget(s):
                result[f"{prefix}{name}"] = queue.pop(0)
                # discard the UI-only seed-control widget value; some packs
                # (e.g. RES4LYF) add it in JS without declaring it in the spec
                if queue and (opts.get("control_after_generate")
                              or (name in ("seed", "noise_seed")
                                  and queue[0] in ("fixed", "increment", "decrement", "randomize"))):
                    queue.pop(0)

    @staticmethod
    def _is_widget(s):
        kind = s[0]
        opts = s[1] if len(s) > 1 and isinstance(s[1], dict) else {}
        if opts.get("forceInput"):
            return False
        return isinstance(kind, list) or kind in WIDGET_SCALAR_TYPES

    def convert(self):
        prompt = {}
        for scope in self.all_scopes():
            for nid, node in scope.nodes.items():
                ntype = node["type"]
                if (ntype in VIRTUAL_TYPES or ntype in self.subgraph_defs
                        or node.get("mode", 0) in (2, 4)):
                    continue
                if ntype not in self.oi:
                    raise SystemExit(
                        f"ERROR: node {scope.fid(nid)} has type {ntype!r} which the ComfyUI "
                        "instance does not know -- install its custom nodes first")
                inputs = self.widget_inputs(node)
                for inp in node.get("inputs", []):
                    source = self.resolve(scope, inp.get("link"))
                    if isinstance(source, WidgetDefault):
                        # the node's own widget value (if any) already won via
                        # widget_inputs; only fill inputs that have none
                        inputs.setdefault(inp["name"], source.value)
                    elif source is not None:
                        inputs[inp["name"]] = list(source)
                    elif inp.get("link") is not None and not inp.get("widget"):
                        print(f"warning: {scope.fid(nid)} ({ntype}) input {inp['name']!r} "
                              "leads to a muted/bypassed dead end; leaving unconnected",
                              file=sys.stderr)
                entry = {"class_type": ntype, "inputs": inputs}
                title = node.get("title")
                if title:
                    entry["_meta"] = {"title": title}
                prompt[scope.fid(nid)] = entry
        return prompt


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--url", required=True, help="ComfyUI base URL (source of /object_info)")
    ap.add_argument("workflow", help="UI-format workflow JSON")
    ap.add_argument("-o", "--output", help="where to write the API JSON (default: stdout)")
    args = ap.parse_args()

    wf = json.loads(pathlib.Path(args.workflow).read_text())
    if "nodes" not in wf:
        sys.exit("ERROR: input already looks like API format (no 'nodes' array)")

    object_info = requests.get(f"{args.url.rstrip('/')}/object_info", timeout=120).json()
    prompt = Converter(wf, object_info).convert()
    out = json.dumps(prompt, indent=2, sort_keys=True)
    if args.output:
        pathlib.Path(args.output).write_text(out + "\n")
        print(f"wrote {args.output} ({len(prompt)} nodes)", file=sys.stderr)
    else:
        print(out)


if __name__ == "__main__":
    main()