#!/usr/bin/env python3
"""UI→API 转换：H3无损融合单采【官方满血API画布级本地流】.json
规则源: comfyui-ui-to-api-conversion skill + h3-workflow-case-bench skill
- SetNode/GetNode 解引用（本容器 :8190 无此类，图级行展开）
- DynamicCombo V3: selection="sol-attn" + selection.tau=1.3
- 槽位真值一律 UI links[]
"""
import json, sys, urllib.request

OI_URL = "http://127.0.0.1:8190/object_info"
UI_JSON = "/home/kai/下载/H3无损融合单采【官方满血API画布级本地流】.json"

oi = json.load(urllib.request.urlopen(OI_URL, timeout=60))
d = json.load(open(UI_JSON))
nodes = {n['id']: n for n in d['nodes']}
links = {l[0]: l for l in d['links']}  # id -> [id, src_node, src_slot, dst_node, dst_slot, type]

# ---- Set/Get 解析: content -> (owner_node_id) ----
set_bindings = {}   # content -> node_id (SetNode 定义)
get_bindings = {}   # node_id -> content (GetNode 引用)
for nid, n in nodes.items():
    wv = n.get('widgets_values') or []
    if isinstance(wv, dict):
        wv = list(wv.values())
    content = wv[0] if wv else None
    if n['type'] == 'SetNode':
        set_bindings[content] = nid
    elif n['type'] == 'GetNode':
        get_bindings[nid] = content

def resolve_source(nid, slot):
    """返回 (final_src_node_str, slot) 沿 link 回溯，穿透 GetNode."""
    n = nodes[nid]
    # 找到该节点输出 slot 对应的 link
    for lid, l in links.items():
        if l[1] == nid and l[2] == slot:
            dst_node, dst_slot = l[3], l[4]
            dn = nodes[dst_node]
            if dn['type'] == 'GetNode':
                content = get_bindings[dst_node]
                owner = set_bindings[content]
                # SetNode 的输出即其输入线
                return resolve_set_owner(owner)
            return (str(dst_node), dst_slot)
    raise SystemExit(f"no outgoing link from {nid} slot {slot}")

def resolve_set_owner(set_nid):
    """SetNode 的源 = 其输入 link 的源节点(可能是 GetNode，继续穿透)."""
    n = nodes[set_nid]
    for inp in n.get('inputs', []):
        lid = inp.get('link')
        if lid and lid in links:
            l = links[lid]
            src = nodes[l[1]]
            if src['type'] == 'GetNode':
                return resolve_set_owner(set_bindings[get_bindings[l[1]]])
            return (str(l[1]), l[2])
    raise SystemExit(f"SetNode {set_nid} no input")

# ---- 构建 API JSON ----
api = {}
for nid, n in nodes.items():
    ct = n['type']
    if ct in ('MarkdownNote', 'Note', 'SetNode', 'GetNode', '忽略多组孤海', 'Reroute', 'PrimitiveNode'):
        continue
    if ct not in oi:
        raise SystemExit(f"容器缺类: {ct}")
    schema = oi[ct]['input']
    inputs = {}
    # 1) 输入 link
    for inp in n.get('inputs', []) or []:
        lid = inp.get('link')
        if lid is None:
            continue
        if lid not in links:
            continue
        l = links[lid]
        src = nodes[l[1]]
        try:
            if src['type'] == 'GetNode':
                content = get_bindings[l[1]]
                owner = set_bindings[content]
                final = resolve_set_owner(owner)
                inputs[inp['name']] = list(final)
            else:
                inputs[inp['name']] = [str(l[1]), l[2]]
        except SystemExit:
            # 断线保存的 SetNode（如 231）→ 跳过该输入；若消费者仍需要则提交时会报缺键
            continue
    # 2) widgets_values 平铺
    wv = n.get('widgets_values') or []
    if isinstance(wv, dict):
        for k, v in wv.items():
            if k in schema.get('required', {}) or k in (schema.get('optional') or {}):
                inputs[k] = v
    else:
        order = [k for sec in ('required', 'optional') for k in (schema.get(sec) or {})]
        flat = [w for w in wv if not isinstance(w, dict)]
        wi = 0
        for k in order:
            if k in inputs or wi >= len(flat):
                continue
            # widget 挂线后其值不算（值在 link 源），但平铺顺序仍占位
            inputs[k] = flat[wi]; wi += 1
    api[str(nid)] = {'class_type': ct, 'inputs': inputs}

# ---- 手工修正段（schema 定谳后的特判）----
# DynamicCombo V3: BlockSparseAttention#165
bs = api['165']['inputs']
wv_bs = nodes[165]['widgets_values']  # ["sol-attn", 1.3, 0.2, 1, "", 12288, 256, "exact_kv_and_rows", true]
bs.clear()
bs['model'] = ["156", 0]
bs['selection'] = wv_bs[0]
bs['selection.tau'] = wv_bs[1]
bs['start_percent'] = wv_bs[2]
bs['end_percent'] = wv_bs[3]
bs['dense_blocks'] = wv_bs[4]
bs['min_tokens'] = wv_bs[5]
bs['extra_tokens'] = wv_bs[6]
bs['sink_conditioning'] = wv_bs[7]
bs['verbose'] = wv_bs[8] if len(wv_bs) > 8 else False

# ModelAttentionBackend#156
mv = nodes[156]['widgets_values']
api['156']['inputs'] = {'model': api['156']['inputs'].get('model', ["166",0]), 'attention': mv[0]}

# ManualSigmas#147
api['147']['inputs'] = {'sigmas': nodes[147]['widgets_values'][0]}

# ComfyMathExpression#224
api['224']['inputs'] = {'values.a': api['224']['inputs'].get('values.a', ["227",0]),
                        'expression': nodes[224]['widgets_values'][0]}

# ResolutionSelector#228: [aspect, megapixels, multiple] -> width/height 输出
# 实际 schema: aspect_ratio, megapixels, multiple(advanced)
api['228']['inputs'] = {'aspect_ratio': nodes[228]['widgets_values'][0],
                        'megapixels': nodes[228]['widgets_values'][1],
                        'multiple': nodes[228]['widgets_values'][2]}

# MiniMaxH3ReferenceToVideo#136 widgets ["",1344,1792,1008,"match"]:
# schema required: clip,prompt,width,height,length,ref_image_size + optional vae/audio_vae/ref_images/ref_videos/ref_audios
# widget 顺序实际是 [??, width, height, length, ref_image_size]，第一个空串疑似旧 prompt widget
rw = nodes[136]['widgets_values']
h3 = api['136']['inputs']
h3.pop('prompt', None)  # prompt 走 GetNode line —— 但转换后 219->136 的边已在（GetNode 解引用），此处恢复
# prompt 输入来自 GetNode#219(提示词<-Set225<-238)。转换脚本已把 136.prompt 解析为 ["238",0]
# 上面的 pop 会误删 —— 检查是否存在，若无则恢复
if 'prompt' not in h3:
    h3['prompt'] = ["238", 0]
h3['width'] = rw[1]; h3['height'] = rw[2]; h3['length'] = rw[3]; h3['ref_image_size'] = rw[4]

# autogrow flat keys: ref_images.ref_image_0..5, ref_videos.ref_video_0, ref_audios.ref_audio_0..2
# 已由 GetNode 边解析为 inputs['ref_images.ref_image_0'] = [src, slot]
# 修掉平铺错位产生的假键: ref_images/ref_videos/ref_video_audios 是 widgets_values 尾部数字错位产物
for k in ('ref_images', 'ref_videos', 'ref_video_audios'):
    h3.pop(k, None)
# 空 audio refs 不传（V3 optional 整键省略规则）
for k in list(h3):
    if k.startswith('ref_audios.'):
        v = h3[k]
        src_node = nodes[int(v[0])]
        wv_src = src_node.get('widgets_values')
        if wv_src and isinstance(wv_src, list) and not wv_src[0]:
            del h3[k]
h3.pop('ref_videos.ref_video_0', None)  # 冒烟臂不用参考视频  # 空 LoadAudio/LoadVideo → 整键省略

# PrimitiveFloat#227 (时长 7s)
api['227']['inputs'] = {'value': nodes[227]['widgets_values'][0]}

# PrimitiveStringMultiline#238 prompt
api['238']['inputs'] = {'value': nodes[238]['widgets_values'][0]}

# RandomNoise#129 seed
nv = nodes[129]['widgets_values']
api['129']['inputs'] = {'noise_seed': nv[0], 'control': nv[1]}

# KSamplerSelect#123
api['123']['inputs'] = {'sampler_name': nodes[123]['widgets_values'][0]}

# BasicScheduler#160 widgets ["beta",33,1]
sv = nodes[160]['widgets_values']
api['160']['inputs'].update({'scheduler': sv[0], 'steps': sv[1], 'denoise': sv[2]})

# ModelPreviewOverrideKJ#150 widgets [1024,80,true,48,8,"taeh3.safetensors",""]
# 本机 schema 需要 taeh3 tiny_vae（未安装），且该节点仅影响 UI 预览不影响采样结果 → 剪链:
# BasicScheduler/BasicGuider 直连 BlockSparseAttention 输出
del api['150']
api['160']['inputs']['model'] = ["165", 0]
api['126']['inputs']['model'] = ["165", 0]

# VHS_VideoCombine#148 — dict widgets: 提取标量
vhs = api['148']['inputs']
wv148 = nodes[148]['widgets_values']
if isinstance(wv148, dict):
    for k in ('save_output','filename_prefix','loop_count','pix_fmt','save_metadata','crf','trim_to_audio','format'):
        if k in wv148: vhs[k] = wv148[k]

# VHS_LoadVideo#185
wv185 = nodes[185]['widgets_values']
if isinstance(wv185, dict):
    lv = api['185']['inputs']
    for k in ('video','force_rate','custom_width','custom_height','frame_load_cap','skip_first_frames','select_every_nth','format'):
        if k in wv185: lv[k] = wv185[k]

# LoadImage/LoadAudio: widgets_values[0] = 文件名; 空(None/'')的 ref 槽 → 整节点删除+消费端断链
drop_nodes = set()
for nid in ('188','189','190','191','192','199','243','246'):
    if nid in api:
        fname = nodes[int(nid)]['widgets_values'][0]
        if not fname:
            drop_nodes.add(nid)
        else:
            api[nid]['inputs'] = {'image': fname.split('/')[-1]}
for nid in ('182','183','184'):
    if nid in api:
        w = nodes[int(nid)]['widgets_values']
        if not w or not w[0] or w[0] == 'None':
            drop_nodes.add(nid)
        else:
            api[nid]['inputs'] = {'audio': w[0]}
if '185' in api and not api['185']['inputs'].get('video'):
    drop_nodes.add('185')
# 清理 drop 节点及其消费引用
for nid in list(api):
    if nid in drop_nodes:
        del api[nid]
for nid, node in api.items():
    for k in list(node['inputs']):
        v = node['inputs'][k]
        if isinstance(v, list) and v[0] in drop_nodes:
            del node['inputs'][k]

# LoraLoaderModelOnly#145/#166
for nid, idx in (('145',0),('166',0)):
    w = nodes[int(nid)]['widgets_values']
    api[nid]['inputs'].update({'lora_name': w[0], 'strength_model': w[1]})

# UNETLoader#127
w = nodes[127]['widgets_values']
api['127']['inputs'] = {'unet_name': w[0], 'weight_dtype': w[1]}

# CLIPLoader#232 widgets ["qwen3vl_32b...","minimax","default"] — 本机无 int8_convrot CLIP, 用 nvfp4_awq
w = nodes[232]['widgets_values']
api['232']['inputs'] = {'clip_name': 'qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors', 'type': w[1], 'device': w[2]}

# VAELoader#234/235
api['234']['inputs'] = {'vae_name': nodes[234]['widgets_values'][0]}
api['235']['inputs'] = {'vae_name': nodes[235]['widgets_values'][0]}

# SamplerCustomAdvanced#125: 无 widget
# BasicGuider#126: 无 widget

out = {'prompt': api}
json.dump(api, open('/home/kai/shared/2026-10-17/h3_lossless_fusion_api.json','w'), ensure_ascii=False, indent=1)
print("API nodes:", len(api))
for k in sorted(api, key=int):
    print(f"  {k}: {api[k]['class_type']}")