#!/usr/bin/env python3
"""将 ComfyUI workflow.json (UI格式) 转换为 workflow_api.json (API格式)

用法:
    python convert_workflow.py refcontrol_pose.json
    python convert_workflow.py refcontrol_pose.json -o output_api.json
"""

import argparse
import json
from pathlib import Path


def convert(workflow: dict) -> dict:
    """
    workflow.json → workflow_api.json

    UI格式: nodes[] 每个节点有 id/type/inputs[]/outputs[]/widgets_values[]/links[]
    API格式: {node_id: {class_type, inputs: {param: value_or_link}}}

    连接关系: links[] = [link_id, src_node, src_slot, dst_node, dst_slot]
    widgets_values 按顺序对应节点没有连线的输入参数
    """
    nodes = workflow.get("nodes", [])
    links = workflow.get("links", [])

    # 建立 link_id → [src_node_id, src_slot] 的映射
    link_map = {}
    for link in links:
        link_id, src_node, src_slot, dst_node, dst_slot = link[:5]
        link_map[link_id] = [str(src_node), src_slot]

    api = {}

    for node in nodes:
        node_id = str(node["id"])
        class_type = node.get("type", "")

        # 跳过纯 UI 节点（Note、MarkdownNote 等无实际计算的节点）
        if class_type in ("Note", "MarkdownNote", "PrimitiveNode"):
            continue
        # 跳过 mode=4（disabled）的节点
        if node.get("mode") == 4:
            continue

        inputs_def = node.get("inputs", [])      # 有连线的输入槽
        widgets_values = node.get("widgets_values", [])  # 无连线的参数值

        inputs = {}

        # 先处理有连线的输入槽
        linked_slots = set()
        for inp in inputs_def:
            link_id = inp.get("link")
            if link_id is not None and link_id in link_map:
                inputs[inp["name"]] = link_map[link_id]
                linked_slots.add(inp["name"])

        # 再处理 widgets_values（按顺序填入没有连线的参数）
        # 需要知道节点有哪些 widget 参数，通过排除已连线的输入来推断
        # widgets 对应的参数名无法从 workflow.json 直接获取，只能按顺序赋值
        # 用 widget_idx_0, widget_idx_1 ... 作为 key，实际使用时按需重命名
        for i, val in enumerate(widgets_values):
            key = f"widget_{i}"
            inputs[key] = val

        api[node_id] = {
            "class_type": class_type,
            "inputs": inputs,
        }

    return api


def main():
    parser = argparse.ArgumentParser(description="Convert workflow.json to workflow_api.json")
    parser.add_argument("input", help="输入 workflow.json 路径")
    parser.add_argument("-o", "--output", help="输出路径，默认在原文件名加 _api 后缀")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output) if args.output else input_path.with_stem(input_path.stem + "_api")

    with open(input_path, "r", encoding="utf-8") as f:
        workflow = json.load(f)

    # 支持两种格式：直接是 nodes[] 的对象，或者包含 nodes[] 的对象
    if "nodes" not in workflow:
        print("ERROR: 不是有效的 workflow.json，缺少 nodes 字段")
        return

    api = convert(workflow)

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(api, f, indent=2, ensure_ascii=False)

    print(f"转换完成: {output_path}")
    print(f"节点数: {len(api)}")
    for node_id, node in api.items():
        print(f"  [{node_id}] {node['class_type']}")


if __name__ == "__main__":
    main()