ComfyUI_Skills_OpenClaw/ui/workflow_format.py

from __future__ import annotations

import re
from typing import Any


def normalize_string(value: Any, default: str = "") -> str:
    if value is None:
        return default
    return str(value).strip()


def is_editor_workflow(workflow_data: Any) -> bool:
    return (
        isinstance(workflow_data, dict)
        and isinstance(workflow_data.get("nodes"), list)
        and isinstance(workflow_data.get("links"), list)
    )


def is_api_workflow(workflow_data: Any) -> bool:
    if not isinstance(workflow_data, dict) or not workflow_data:
        return False
    if is_editor_workflow(workflow_data):
        return False

    for key, value in workflow_data.items():
        if not isinstance(key, str) or not key.strip():
            return False
        if not isinstance(value, dict) or "class_type" not in value:
            return False
    return True


def _get_type_guess(value: Any) -> str:
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, int) and not isinstance(value, bool):
        return "int"
    if isinstance(value, float):
        return "float"
    return "string"


def _get_auto_mapping(node_class: str, field: str, node_id: str) -> dict[str, Any]:
    if "KSampler" in node_class:
        if field == "seed":
            return {"exposed": True, "required": False, "name": field, "description": "Random seed (for reproducibility)"}
        if field == "steps":
            return {"exposed": True, "required": False, "name": field, "description": "Generation steps"}

    if "CLIPTextEncode" in node_class or "Text" in node_class or "Prompt" in node_class:
        if field in {"text", "prompt"}:
            return {"exposed": True, "required": True, "name": f"prompt_{node_id}", "description": "Text prompt description"}

    if node_class == "EmptyLatentImage" and field in {"width", "height", "batch_size"}:
        return {"exposed": True, "required": False, "name": field, "description": f"Image {field}"}

    if node_class == "SaveImage" and field == "filename_prefix":
        return {"exposed": True, "required": False, "name": field, "description": "Output file prefix"}

    if node_class == "LightCCDoubaoImageNode":
        if field == "prompt":
            return {"exposed": True, "required": True, "name": field, "description": "Positive image prompt"}
        if field == "size":
            return {"exposed": True, "required": False, "name": field, "description": "e.g., 1:1,2048x2048"}
        if field == "seed":
            return {"exposed": True, "required": False, "name": field, "description": "Random seed"}
        if field == "num":
            return {"exposed": True, "required": False, "name": field, "description": "Number of images to generate"}

    if field in {"text", "prompt"}:
        return {"exposed": True, "required": True, "name": f"prompt_{node_id}", "description": "Text prompt"}
    if field == "seed":
        return {"exposed": True, "required": False, "name": f"seed_{node_id}", "description": "Random seed"}
    if field in {"width", "height", "batch_size", "size", "num", "num_images", "max_images"}:
        return {"exposed": True, "required": False, "name": f"{field}_{node_id}", "description": f"Workflow parameter: {field}"}
    if field == "filename_prefix":
        return {"exposed": True, "required": False, "name": f"filename_prefix_{node_id}", "description": "Output file prefix"}

    return {"exposed": False, "required": False, "name": field, "description": ""}


def extract_schema_params(workflow_data: dict[str, Any]) -> dict[str, dict[str, Any]]:
    schema_params: dict[str, dict[str, Any]] = {}

    for node_id, node_object in workflow_data.items():
        if not isinstance(node_object, dict):
            continue
        inputs = node_object.get("inputs")
        if not isinstance(inputs, dict):
            continue

        node_class = normalize_string(node_object.get("class_type"))
        for field, value in inputs.items():
            if isinstance(value, list):
                continue

            auto_mapping = _get_auto_mapping(node_class, field, node_id)
            schema_params[f"{node_id}_{field}"] = {
                "exposed": auto_mapping["exposed"],
                "node_id": int(node_id),
                "field": field,
                "name": auto_mapping["name"],
                "type": _get_type_guess(value),
                "required": auto_mapping["required"],
                "description": auto_mapping["description"],
                "default": value,
                "example": value,
                "choices": [],
                "currentVal": value,
                "nodeClass": node_class or "UnknownNode",
            }

    return schema_params


def build_final_schema(schema_params: dict[str, dict[str, Any]]) -> dict[str, dict[str, Any]]:
    final_schema: dict[str, dict[str, Any]] = {}
    for parameter in schema_params.values():
        if not bool(parameter.get("exposed")):
            continue

        alias = normalize_string(parameter.get("name"))
        if not alias:
            continue

        target: dict[str, Any] = {
            "node_id": int(parameter["node_id"]),
            "field": normalize_string(parameter.get("field")),
            "required": bool(parameter.get("required", False)),
            "type": normalize_string(parameter.get("type"), "string") or "string",
            "description": normalize_string(parameter.get("description")),
        }
        if "default" in parameter:
            target["default"] = parameter["default"]
        if "example" in parameter:
            target["example"] = parameter["example"]
        choices = parameter.get("choices")
        if isinstance(choices, list) and choices:
            target["choices"] = choices
        final_schema[alias] = target

    return final_schema


def suggest_workflow_id(workflow_data: dict[str, Any], file_name: str = "") -> str:
    def normalize_candidate(value: Any) -> str:
        if not isinstance(value, str):
            return ""
        normalized = re.sub(r"[./\\]+", "-", value.strip())
        normalized = re.sub(r"[^\w-]+", "-", normalized, flags=re.UNICODE)
        normalized = re.sub(r"-+", "-", normalized)
        return normalized.strip("-_")

    def first_node_title() -> str:
        for node_object in workflow_data.values():
            if not isinstance(node_object, dict):
                continue
            meta = node_object.get("_meta")
            if isinstance(meta, dict):
                title = normalize_string(meta.get("title"))
                if title:
                    return title
        return ""

    base_file_name = re.sub(r"\.[^.]+$", "", file_name).strip() if file_name else ""
    candidates = [
        workflow_data.get("workflow_name"),
        workflow_data.get("name"),
        workflow_data.get("title"),
        workflow_data.get("_meta", {}).get("title") if isinstance(workflow_data.get("_meta"), dict) else "",
        workflow_data.get("extra", {}).get("workflow_name") if isinstance(workflow_data.get("extra"), dict) else "",
        workflow_data.get("extra", {}).get("name") if isinstance(workflow_data.get("extra"), dict) else "",
        workflow_data.get("extra", {}).get("title") if isinstance(workflow_data.get("extra"), dict) else "",
        workflow_data.get("metadata", {}).get("workflow_name") if isinstance(workflow_data.get("metadata"), dict) else "",
        workflow_data.get("metadata", {}).get("name") if isinstance(workflow_data.get("metadata"), dict) else "",
        workflow_data.get("metadata", {}).get("title") if isinstance(workflow_data.get("metadata"), dict) else "",
        base_file_name,
        first_node_title(),
    ]

    for candidate in candidates:
        normalized = normalize_candidate(candidate)
        if normalized:
            return normalized
    return "workflow"


def _list_value(value: Any) -> list[Any]:
    if isinstance(value, list):
        return value
    return []


class WorkflowImportError(ValueError):
    pass


class EditorWorkflowConverter:
    def __init__(self, object_info: dict[str, Any]):
        self.object_info = object_info

    def convert(self, workflow_data: dict[str, Any]) -> dict[str, Any]:
        nodes = workflow_data.get("nodes")
        links = workflow_data.get("links")
        if not isinstance(nodes, list) or not isinstance(links, list):
            raise WorkflowImportError("Unsupported ComfyUI editor workflow format.")

        node_by_id: dict[int, dict[str, Any]] = {}
        for node in nodes:
            if isinstance(node, dict) and node.get("id") is not None:
                node_by_id[int(node["id"])] = node

        link_map: dict[tuple[int, int], tuple[str, int]] = {}
        for link in links:
            if not isinstance(link, list) or len(link) < 5:
                continue
            source_node_id = self._resolve_link_source(int(link[1]), int(link[2]), node_by_id, links)
            source_slot = int(link[2])
            target_node_id = int(link[3])
            target_slot = int(link[4])
            if source_node_id is None:
                continue
            link_map[(target_node_id, target_slot)] = (source_node_id, source_slot)

        api_workflow: dict[str, Any] = {}
        for node in nodes:
            if not isinstance(node, dict):
                continue

            node_id = node.get("id")
            class_type = normalize_string(node.get("type"))
            if node_id is None or not class_type:
                continue
            if class_type in {"Reroute", "Note"}:
                continue

            node_object_info = self.object_info.get(class_type)
            if not isinstance(node_object_info, dict):
                raise WorkflowImportError(f"ComfyUI object_info is missing node type '{class_type}'.")

            title = normalize_string(node.get("title"))
            api_workflow[str(node_id)] = {
                "inputs": self._convert_node_inputs(node, class_type, node_object_info, link_map),
                "class_type": class_type,
                "_meta": {"title": title or class_type},
            }

        if not api_workflow:
            raise WorkflowImportError("No supported nodes were found in the editor workflow.")

        return api_workflow

    def _resolve_link_source(
        self,
        source_node_id: int,
        source_slot: int,
        node_by_id: dict[int, dict[str, Any]],
        links: list[Any],
    ) -> str | None:
        node = node_by_id.get(source_node_id)
        if not isinstance(node, dict):
            return str(source_node_id)

        class_type = normalize_string(node.get("type"))
        if class_type != "Reroute":
            return str(source_node_id)

        input_slots = _list_value(node.get("inputs"))
        if not input_slots:
            return None

        incoming_link_id = input_slots[0].get("link") if isinstance(input_slots[0], dict) else None
        if incoming_link_id is None:
            return None

        for link in links:
            if not isinstance(link, list) or len(link) < 5:
                continue
            if int(link[0]) != int(incoming_link_id):
                continue
            return self._resolve_link_source(int(link[1]), int(link[2]), node_by_id, links)

        return None

    def _convert_node_inputs(
        self,
        node: dict[str, Any],
        class_type: str,
        node_object_info: dict[str, Any],
        link_map: dict[tuple[int, int], tuple[str, int]],
    ) -> dict[str, Any]:
        node_id = int(node["id"])
        input_defs = self._ordered_input_defs(node_object_info)
        input_slots = _list_value(node.get("inputs"))
        widget_values = _list_value(node.get("widgets_values"))

        converted: dict[str, Any] = {}
        connected_names: set[str] = set()

        for slot_index, slot in enumerate(input_slots):
            if not isinstance(slot, dict):
                continue
            slot_name = normalize_string(slot.get("name"))
            if not slot_name:
                continue
            link_tuple = link_map.get((node_id, slot_index))
            if link_tuple is None:
                continue
            converted[slot_name] = [link_tuple[0], link_tuple[1]]
            connected_names.add(slot_name)

        widget_field_names = [
            normalize_string(slot.get("name"))
            for slot in input_slots
            if isinstance(slot, dict)
            and normalize_string(slot.get("name"))
            and normalize_string(slot.get("name")) not in connected_names
            and isinstance(slot.get("widget"), dict)
        ]

        if not widget_field_names:
            widget_field_names = [name for name in input_defs if name not in connected_names]

        control_after_generate_fields = self._control_after_generate_fields(node_object_info)
        widget_value_index = 0
        for field_name in widget_field_names:
            if widget_value_index >= len(widget_values):
                break

            converted[field_name] = widget_values[widget_value_index]
            widget_value_index += 1

            if field_name not in control_after_generate_fields:
                continue
            if widget_value_index >= len(widget_values):
                continue

            control_value = widget_values[widget_value_index]
            if isinstance(control_value, str) and control_value.lower() in {
                "fixed",
                "increment",
                "decrement",
                "randomize",
            }:
                widget_value_index += 1

        if not converted and widget_values and not widget_field_names:
            raise WorkflowImportError(f"Unable to map widget values for node type '{class_type}'.")

        return converted

    @staticmethod
    def _control_after_generate_fields(node_object_info: dict[str, Any]) -> set[str]:
        fields: set[str] = set()

        def visit_section(section: Any) -> None:
            if not isinstance(section, dict):
                return
            for field_name, definition in section.items():
                if (
                    isinstance(definition, list)
                    and len(definition) >= 2
                    and isinstance(definition[1], dict)
                    and definition[1].get("control_after_generate")
                ):
                    fields.add(field_name)

        input_section = node_object_info.get("input")
        if isinstance(input_section, dict):
            visit_section(input_section.get("required"))
            visit_section(input_section.get("optional"))

        visit_section(node_object_info.get("required"))
        visit_section(node_object_info.get("optional"))
        return fields

    @staticmethod
    def _ordered_input_defs(node_object_info: dict[str, Any]) -> list[str]:
        ordered: list[str] = []
        input_order = node_object_info.get("input_order")
        if isinstance(input_order, dict):
            for section_name in ("required", "optional"):
                section = input_order.get(section_name)
                if isinstance(section, list):
                    ordered.extend(normalize_string(item) for item in section if normalize_string(item))
        if ordered:
            return ordered

        input_section = node_object_info.get("input")
        if isinstance(input_section, dict):
            for section_name in ("required", "optional"):
                section = input_section.get(section_name)
                if isinstance(section, dict):
                    ordered.extend(section.keys())
        if ordered:
            return ordered

        for section_name in ("required", "optional"):
            section = node_object_info.get(section_name)
            if not isinstance(section, dict):
                continue
            ordered.extend(section.keys())
        return ordered