#!/usr/bin/env python3
# Setup Boogu official workflows on GPU machine:
# 1) copy official templates (turbo t2i / edit) into user workflows dir
# 2) create nvfp4 variants (swap diffusion model filename)
# 3) dump subgraph node summary (for API verification prompt)
import json, os

TPL_DIR = "/home/zyw/ComfyUI/venv/lib/python3.12/site-packages/comfyui_workflow_templates_json/templates"
WF_DIR = "/home/zyw/ComfyUI/user/default/workflows"
os.makedirs(WF_DIR, exist_ok=True)

pairs = [
    ("image_boogu_image_0_1_turbo_t2i.json", "boogu_image_0_1_turbo_t2i_official.json"),
    ("image_boogu_image_0_1_edit.json", "boogu_image_0_1_edit_official.json"),
]

for src, dst in pairs:
    with open(os.path.join(TPL_DIR, src), encoding="utf-8") as f:
        text = f.read()
    with open(os.path.join(WF_DIR, dst), "w", encoding="utf-8") as f:
        f.write(text)
    print("copied:", dst, f"({len(text)} bytes)")

nvfp4_map = {
    "boogu_image_0_1_turbo_t2i_official.json": (
        "boogu_image_0_1_turbo_t2i_nvfp4.json",
        "boogu_image_turbo_fp8_scaled.safetensors",
        "boogu_image_turbo_hotfix_nvfp4.safetensors",
    ),
    "boogu_image_0_1_edit_official.json": (
        "boogu_image_0_1_edit_nvfp4.json",
        "boogu_image_edit_fp8_scaled.safetensors",
        "boogu_image_edit_nvfp4.safetensors",
    ),
}

for base, (nvfp4_name, old, new) in nvfp4_map.items():
    with open(os.path.join(WF_DIR, base), encoding="utf-8") as f:
        text = f.read()
    n = text.count(old)
    if n == 0:
        print("WARN: pattern not found in", base, "- skipping")
        continue
    text = text.replace(old, new)
    with open(os.path.join(WF_DIR, nvfp4_name), "w", encoding="utf-8") as f:
        f.write(text)
    print("created nvfp4 variant:", nvfp4_name, f"(replaced {n}x)")

# --- dump node summary of both official templates for API verification ---
def walk(o):
    if isinstance(o, dict):
        if "type" in o and "widgets_values" in o:
            print("NODE:", o.get("type"), "|", o.get("widgets_values"))
        for v in o.values():
            walk(v)
    elif isinstance(o, list):
        for v in o:
            walk(v)

for name in ["boogu_image_0_1_turbo_t2i_official.json", "boogu_image_0_1_edit_official.json"]:
    print(f"\n===== NODE DUMP: {name} =====")
    with open(os.path.join(WF_DIR, name), encoding="utf-8") as f:
        walk(json.load(f))