"""
Generate `workflow.json` for Workflow1111.

Hand-writing ~3000 lines of graph JSON is how wiring bugs get in, so the graph
is generated and then *verified* — see `verify()` at the bottom. The single
most important guarantee: an `fn` operator's input ports are derived from the
bound function's own signature via `inspect`, so port order can never drift
away from the Python argument order (the executor passes `fn` args
positionally, in port order).

    python apps/05_workflow1111/build_workflow.py

Re-run after editing `nodes.py`. Node positions come from `layout.json` (the
hand-arranged, overlap-checked layout), so rebuilding preserves the canvas
arrangement instead of resetting it. To adopt a new arrangement, drag nodes in
the canvas and re-snapshot `layout.json` from the saved `workflow.json`.
"""

import inspect
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from gradio.workflow import _is_injected_param  # noqa: E402

import nodes as N  # noqa: E402


def fn_params(func):
    """A bound function's real input parameters.

    Skips gradio's injected parameters (`OAuthToken`, `OAuthProfile`,
    `Request`) — gradio supplies those itself, so they must not become ports.
    """
    hints = getattr(func, "__annotations__", {})
    try:
        from typing import get_type_hints
        hints = get_type_hints(func)
    except Exception:
        pass
    return [p for p in inspect.signature(func).parameters
            if not _is_injected_param(hints.get(p))]

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "workflow.json")
LAYOUT = os.path.join(HERE, "layout.json")

# Sample images ship in the Space repo and are referenced by their public Hub
# URL. That is the one default shape that satisfies everything at once: the
# canvas only renders a reference default that carries a `url` key (it strips
# `path` from graph defaults for safety), while `call_model`/`call_space` need
# something a remote provider can actually fetch — which a relative
# `/gradio_api/file=` URL is not, but an absolute https one is. The same URLs
# therefore work identically when running locally and on the Space.
SPACE_ID = os.environ.get("WORKFLOW1111_SPACE", "ysharma/Workflow1111")
SAMPLE_BASE = f"https://huggingface.co/spaces/{SPACE_ID}/resolve/main/samples"


def sample(filename):
    return {"url": f"{SAMPLE_BASE}/{filename}"}

# Verified working on HF Inference Providers — see the probe results recorded
# in the README. Swapping these is the main "model checkpoint" knob.
T2I_MODEL = "black-forest-labs/FLUX.1-schnell"
T2I_QUALITY_MODEL = "black-forest-labs/FLUX.1-dev"
EDIT_MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
LLM_MODEL = "Qwen/Qwen3-4B-Instruct-2507"
VLM_MODEL = "Qwen/Qwen2.5-VL-72B-Instruct"
DETECT_MODEL = "facebook/detr-resnet-50"
CLASSIFY_MODEL = "google/vit-base-patch16-224"

RMBG_SPACE = "briaai/BRIA-RMBG-2.0"
UPSCALE_SPACE = "gokaygokay/AuraSR-v2"

references, operators, subjects, edges = [], [], [], []

COL = [60, 420, 800, 1180, 1560, 1940]  # x positions by pipeline stage


def _size(n_in, n_out, width, base=64, per_port=30):
    return width, base + per_port * max(n_in, n_out, 1)


def ref(node_id, label, port_type, x, y, default=None, width=250):
    """A reference node — a free input. These become the API parameters."""
    references.append({
        "id": node_id, "role": "reference", "label": label, "asset_type": port_type,
        "inputs": [{"id": "in", "label": label, "type": port_type}],
        "outputs": [{"id": "out", "label": label, "type": port_type}],
        "data": {} if default is None else {"out": default},
        "x": x, "y": y, "width": width,
        "height": 200 if port_type in ("image", "audio", "video") else 96,
    })
    return node_id


def fn(node_id, fn_name, x, y, *, label=None, types=None, data=None,
       required=(), outputs=None, width=290):
    """An `fn` operator. Input ports are generated from the bound function's
    signature, so the positional call order is correct by construction."""
    func = N.BIND[fn_name]
    params = fn_params(func)
    types = types or {}
    data = data or {}

    inputs = [{
        "id": f"in_{p}",
        "label": p,
        "type": types.get(p, "text"),
        **({"required": True} if p in required else {}),
    } for p in params]

    outs = outputs or [("out_0", "output", "text")]
    out_ports = [{"id": oid, "label": olabel, "type": otype, "output_index": i}
                 for i, (oid, olabel, otype) in enumerate(outs)]

    w, h = _size(len(inputs), len(out_ports), width)
    operators.append({
        "id": node_id, "role": "operator", "kind": "fn", "fn": fn_name,
        "label": label or fn_name,
        "inputs": inputs, "outputs": out_ports,
        "data": {f"in_{k}": v for k, v in data.items()},
        "x": x, "y": y, "width": w, "height": h,
    })
    return node_id


def model(node_id, model_id, endpoint, pipeline_tag, x, y, *, label=None,
          inputs=(), outputs=None, data=None, width=290):
    """A `model` operator (HF Inference Providers).

    Input port **ids** are forwarded verbatim as keyword arguments to
    `InferenceClient.<endpoint>()`, which is what gives txt2img its real
    negative-prompt / steps / CFG / seed / size controls.
    """
    in_ports = [{"id": pid, "label": plabel, "type": ptype,
                 **({"required": True} if req else {})}
                for pid, plabel, ptype, req in inputs]
    outs = outputs or [("out_0", "Image", "image")]
    out_ports = [{"id": oid, "label": olabel, "type": otype, "output_index": i}
                 for i, (oid, olabel, otype) in enumerate(outs)]

    w, h = _size(len(in_ports), len(out_ports), width)
    operators.append({
        "id": node_id, "role": "operator", "kind": "model",
        "model_id": model_id, "pipeline_tag": pipeline_tag, "endpoint": endpoint,
        "label": label or model_id.split("/")[-1],
        "inputs": in_ports, "outputs": out_ports, "data": data or {},
        "x": x, "y": y, "width": w, "height": h,
    })
    return node_id


def space(node_id, space_id, endpoint, x, y, *, label=None, inputs=(),
          outputs=None, data=None, width=290):
    """A `space` operator. Inputs are passed **positionally**, in port order."""
    in_ports = [{"id": pid, "label": plabel, "type": ptype,
                 **({"required": True} if req else {})}
                for pid, plabel, ptype, req in inputs]
    out_ports = [{"id": oid, "label": olabel, "type": otype, "output_index": idx}
                 for oid, olabel, otype, idx in (outputs or [])]

    w, h = _size(len(in_ports), len(out_ports), width)
    operators.append({
        "id": node_id, "role": "operator", "kind": "space",
        "space_id": space_id, "endpoint": endpoint,
        "label": label or space_id.split("/")[-1],
        "inputs": in_ports, "outputs": out_ports, "data": data or {},
        "x": x, "y": y, "width": w, "height": h,
    })
    return node_id


def out(node_id, label, port_type, x, y, width=300):
    """A subject node — a workflow output, and an API endpoint result."""
    subjects.append({
        "id": node_id, "role": "subject", "label": label, "asset_type": port_type,
        "inputs": [{"id": "in", "label": label, "type": port_type}],
        "outputs": [{"id": "out", "label": label, "type": port_type}],
        "data": {},
        "x": x, "y": y, "width": width,
        "height": 260 if port_type == "image" else 190,
    })
    return node_id


def link(a, b):
    """Wire "node.port" → "node.port"."""
    fnode, fport = a.split(".")
    tnode, tport = b.split(".")
    edges.append({
        "id": f"e{len(edges) + 1}",
        "from_node_id": fnode, "from_port_id": fport,
        "to_node_id": tnode, "to_port_id": tport,
        "type": None,  # filled in by verify() from the source port
    })


# ═════════════════════════════════════════════════════════════════════════════
# 1 · txt2img — the flagship pipeline
# ═════════════════════════════════════════════════════════════════════════════
Y = 60
ref("ref_prompt", "Prompt", "text", COL[0], Y,
    "a red fox standing in a snowy pine forest, looking at the camera")
ref("ref_negative", "Negative prompt", "text", COL[0], Y + 120, "")
ref("ref_style", "Style preset", "text", COL[0], Y + 240, "Cinematic")

fn("op_style", "apply_style", COL[1], Y, label="① Prompt builder",
   required=("prompt",),
   data={"extra_tags": "", "quality_boost": True},
   outputs=[("out_0", "prompt", "text")])

fn("op_negative", "build_negative", COL[1], Y + 190, label="① Negative builder",
   types={"use_base": "boolean", "safety_filter": "boolean"},
   data={"negative": "", "use_base": True, "safety_filter": True},
   outputs=[("out_0", "negative", "text")])

fn("op_sampler", "sampler_settings", COL[1], Y + 380, label="② Sampler",
   types={"steps": "number", "cfg_scale": "number", "seed": "number",
          "width": "number", "height": "number"},
   data={"steps": 4, "cfg_scale": 1.0, "seed": -1,
         "aspect": "1:1 Square", "width": 1024, "height": 1024},
   outputs=[("out_steps", "steps", "number"), ("out_cfg", "cfg", "number"),
            ("out_seed", "seed", "number"), ("out_width", "width", "number"),
            ("out_height", "height", "number")])

# txt2img is an `fn` node, not a `model` node, on purpose: the canvas rewrites
# a model node's ports to the endpoint's canonical schema (just `prompt` for
# text_to_image), which silently discarded the negative prompt, steps, CFG,
# seed and size. `fn` ports are left alone, so the control surface survives.
fn("op_txt2img", "txt2img", COL[2], Y, label="③ txt2img · FLUX.1-schnell",
   types={"steps": "number", "cfg_scale": "number", "seed": "number",
          "width": "number", "height": "number"},
   required=("prompt",),
   data={"negative_prompt": "", "steps": 4, "cfg_scale": 1.0, "seed": -1,
         "width": 1024, "height": 1024, "model_id": T2I_MODEL},
   outputs=[("out_0", "image", "image")])

fn("op_geninfo", "generation_info", COL[2], Y + 300, label="④ Generation params",
   types={"steps": "number", "cfg_scale": "number", "seed": "number",
          "width": "number", "height": "number"},
   data={"model_id": T2I_MODEL},
   outputs=[("out_0", "parameters", "text")])

fn("op_post", "postprocess", COL[3], Y, label="⑤ Post-processing",
   types={"image": "image", "upscale": "number", "sharpen": "number",
          "saturation": "number", "contrast": "number", "brightness": "number",
          "vignette": "number", "grain": "number", "border": "number"},
   required=("image",),
   data={"upscale": 1.0, "upscale_method": "Lanczos", "sharpen": 0.35,
         "saturation": 1.05, "contrast": 1.02, "brightness": 1.0,
         "vignette": 0.12, "grain": 0.04, "border": 0.0, "watermark": ""},
   outputs=[("out_0", "image", "image")])

out("sub_image", "🖼 Image", "image", COL[4], Y)
out("sub_params", "📋 Generation parameters", "text", COL[4], Y + 300)

link("ref_prompt.out", "op_style.in_prompt")
link("ref_style.out", "op_style.in_style")
link("ref_style.out", "op_negative.in_style")
link("ref_negative.out", "op_negative.in_negative")

link("op_style.out_0", "op_txt2img.in_prompt")
link("op_negative.out_0", "op_txt2img.in_negative_prompt")
link("op_sampler.out_steps", "op_txt2img.in_steps")
link("op_sampler.out_cfg", "op_txt2img.in_cfg_scale")
link("op_sampler.out_seed", "op_txt2img.in_seed")
link("op_sampler.out_width", "op_txt2img.in_width")
link("op_sampler.out_height", "op_txt2img.in_height")

link("op_style.out_0", "op_geninfo.in_prompt")
link("op_negative.out_0", "op_geninfo.in_negative")
link("op_sampler.out_steps", "op_geninfo.in_steps")
link("op_sampler.out_cfg", "op_geninfo.in_cfg_scale")
link("op_sampler.out_seed", "op_geninfo.in_seed")
link("op_sampler.out_width", "op_geninfo.in_width")
link("op_sampler.out_height", "op_geninfo.in_height")

link("op_txt2img.out_0", "op_post.in_image")
link("op_geninfo.out_0", "op_post.in_embed_info")
link("op_post.out_0", "sub_image.in")
link("op_geninfo.out_0", "sub_params.in")

# ═════════════════════════════════════════════════════════════════════════════
# 2 · Hires fix — upscale the txt2img result, then refine it with img2img
# ═════════════════════════════════════════════════════════════════════════════
Y = 700
ref("ref_hires_instruction", "Hires refine instruction", "text", COL[2], Y,
    "enhance fine detail and micro-texture, keep the composition identical")

fn("op_hires_prep", "prep_image", COL[3], Y + 130, label="⑥ Hires prep",
   types={"image": "image", "max_side": "number", "strip_alpha": "boolean"},
   required=("image",),
   data={"max_side": 1024, "mode": "Fit", "strip_alpha": True},
   outputs=[("out_0", "image", "image")])

model("op_hires", EDIT_MODEL, "image_to_image", "image-to-image", COL[4], Y + 130,
      label="⑦ Hires fix · FLUX.1-Kontext",
      inputs=[("image", "image", "image", True),
              ("prompt", "prompt", "text", True)])

out("sub_hires", "✨ Hires image", "image", COL[5], Y + 130)

link("op_post.out_0", "op_hires_prep.in_image")
link("op_hires_prep.out_0", "op_hires.image")
link("ref_hires_instruction.out", "op_hires.prompt")
link("op_hires.out_0", "sub_hires.in")

# ═════════════════════════════════════════════════════════════════════════════
# 3 · img2img — edit an uploaded image by instruction
# ═════════════════════════════════════════════════════════════════════════════
Y = 1080
ref("ref_init_image", "Init image", "image", COL[0], Y, sample("init_image.jpg"))
ref("ref_edit_instruction", "Edit instruction", "text", COL[0], Y + 240,
    "make it a snowy winter scene at golden hour")

fn("op_i2i_prep", "prep_image", COL[1], Y, label="① Prepare init image",
   types={"image": "image", "max_side": "number", "strip_alpha": "boolean"},
   required=("image",),
   data={"max_side": 1024, "mode": "Fit", "strip_alpha": True},
   outputs=[("out_0", "image", "image")])

model("op_i2i", EDIT_MODEL, "image_to_image", "image-to-image", COL[2], Y,
      label="② img2img · FLUX.1-Kontext",
      inputs=[("image", "image", "image", True),
              ("prompt", "prompt", "text", True)])

fn("op_i2i_post", "postprocess", COL[3], Y, label="③ Post-processing",
   types={"image": "image", "upscale": "number", "sharpen": "number",
          "saturation": "number", "contrast": "number", "brightness": "number",
          "vignette": "number", "grain": "number", "border": "number"},
   required=("image",),
   data={"upscale": 1.0, "upscale_method": "Lanczos", "sharpen": 0.3,
         "saturation": 1.0, "contrast": 1.0, "brightness": 1.0,
         "vignette": 0.0, "grain": 0.0, "border": 0.0, "watermark": ""},
   outputs=[("out_0", "image", "image")])

out("sub_i2i", "🎨 Edited image", "image", COL[4], Y)

link("ref_init_image.out", "op_i2i_prep.in_image")
link("op_i2i_prep.out_0", "op_i2i.image")
link("ref_edit_instruction.out", "op_i2i.prompt")
link("op_i2i.out_0", "op_i2i_post.in_image")
link("op_i2i_post.out_0", "sub_i2i.in")

# ═════════════════════════════════════════════════════════════════════════════
# 4 · Prompt magic — an LLM writes the prompt for you
# ═════════════════════════════════════════════════════════════════════════════
Y = 1450
ref("ref_idea", "Rough idea", "text", COL[0], Y, "a lighthouse in a storm")

fn("op_magic", "magic_instruction", COL[1], Y, label="① Build instruction",
   required=("idea",),
   data={"target_style": "Cinematic", "verbosity": "Detailed"},
   outputs=[("out_0", "instruction", "text")])

fn("op_llm", "chat_llm", COL[2], Y, label="② Prompt LLM · Qwen3-4B",
   types={"max_tokens": "number"}, required=("prompt",),
   data={"model_id": LLM_MODEL, "max_tokens": 512},
   outputs=[("out_0", "Text", "text")])

fn("op_clean_magic", "clean_prompt", COL[3], Y, label="③ Tidy up",
   types={"max_tags": "number"}, required=("raw",),
   data={"max_tags": 40},
   outputs=[("out_0", "prompt", "text")])

out("sub_magic", "🪄 Generated prompt", "text", COL[4], Y)

link("ref_idea.out", "op_magic.in_idea")
link("op_magic.out_0", "op_llm.in_prompt")
link("op_llm.out_0", "op_clean_magic.in_raw")
link("op_clean_magic.out_0", "sub_magic.in")

# ═════════════════════════════════════════════════════════════════════════════
# 5 · Interrogate — recover a prompt (and labels) from an image
# ═════════════════════════════════════════════════════════════════════════════
Y = 1780
ref("ref_interrogate_image", "Image to interrogate", "image", COL[0], Y,
    sample("interrogate.jpg"))

fn("op_vlm", "interrogate", COL[1], Y, label="① Interrogate · Qwen2.5-VL",
   types={"image": "image", "max_tokens": "number"}, required=("image",),
   data={"instruction": "Describe this image as a Stable Diffusion prompt: "
                        "comma-separated visual tags only, covering subject, setting, "
                        "composition, lighting, colour and medium. No sentences, "
                        "no preamble.",
         "model_id": VLM_MODEL, "max_tokens": 512},
   outputs=[("out_0", "Text", "text")])

fn("op_clean_interrogate", "clean_prompt", COL[2], Y, label="② Tidy up",
   types={"max_tags": "number"}, required=("raw",),
   data={"max_tags": 45},
   outputs=[("out_0", "prompt", "text")])

# `fn`, not `model`: a `json` output port reaches the canvas as the literal
# string "[object Object]" (JS String(obj) instead of JSON.stringify), so the
# labels never survive the edge. Text ports carrying JSON do.
fn("op_classify", "classify_image", COL[1], Y + 260, label="③ Classify · ViT",
   types={"image": "image"}, required=("image",),
   data={"model_id": CLASSIFY_MODEL},
   outputs=[("out_0", "labels", "text")])

fn("op_labels", "top_labels", COL[2], Y + 260, label="④ Rank labels",
   types={"labels": "text", "top_k": "number", "min_score": "number"},
   required=("labels",),
   data={"top_k": 5, "min_score": 0.01},
   outputs=[("out_0", "table", "text"), ("out_1", "rows", "text")])

out("sub_interrogated", "🔍 Recovered prompt", "text", COL[3], Y)
out("sub_labels", "🏷 Classification", "text", COL[3], Y + 260)

link("ref_interrogate_image.out", "op_vlm.in_image")
link("op_vlm.out_0", "op_clean_interrogate.in_raw")
link("op_clean_interrogate.out_0", "sub_interrogated.in")
link("ref_interrogate_image.out", "op_classify.in_image")
link("op_classify.out_0", "op_labels.in_labels")
link("op_labels.out_0", "sub_labels.in")

# ═════════════════════════════════════════════════════════════════════════════
# 6 · Detect & mask — object detection into an inpainting mask
# ═════════════════════════════════════════════════════════════════════════════
Y = 2200
ref("ref_detect_image", "Image to analyse", "image", COL[0], Y, sample("detect.jpg"))

fn("op_detect", "detect_objects", COL[1], Y, label="① Detect · DETR",
   types={"image": "image", "min_score": "number"}, required=("image",),
   data={"model_id": DETECT_MODEL, "min_score": 0.0},
   outputs=[("out_0", "detections", "text")])

fn("op_draw", "draw_detections", COL[2], Y, label="② Annotate",
   types={"image": "image", "detections": "text", "min_score": "number",
          "show_labels": "boolean"},
   required=("image", "detections"),
   data={"min_score": 0.5, "show_labels": True},
   outputs=[("out_0", "image", "image"), ("out_1", "summary", "text")])

fn("op_mask", "mask_from_detections", COL[2], Y + 300, label="③ Build inpaint mask",
   types={"image": "image", "detections": "text", "min_score": "number",
          "feather": "number", "invert": "boolean", "preview": "boolean"},
   required=("image", "detections"),
   data={"label_filter": "", "min_score": 0.5, "feather": 8,
         "invert": False, "preview": False},
   outputs=[("out_0", "mask", "image")])

out("sub_detected", "📦 Detected objects", "image", COL[3], Y)
out("sub_detect_summary", "📝 Detection summary", "text", COL[3], Y + 300)
out("sub_mask", "🎭 Inpaint mask", "image", COL[4], Y + 300)

link("ref_detect_image.out", "op_detect.in_image")
link("ref_detect_image.out", "op_draw.in_image")
link("op_detect.out_0", "op_draw.in_detections")
link("ref_detect_image.out", "op_mask.in_image")
link("op_detect.out_0", "op_mask.in_detections")
link("op_draw.out_0", "sub_detected.in")
link("op_draw.out_1", "sub_detect_summary.in")
link("op_mask.out_0", "sub_mask.in")

# ═════════════════════════════════════════════════════════════════════════════
# 7 · Prompt matrix — four variants rendered in parallel into an X/Y grid
# ═════════════════════════════════════════════════════════════════════════════
Y = 2700
ref("ref_matrix_base", "Matrix base prompt", "text", COL[0], Y, "a lone tree on a hill")
ref("ref_matrix_variants", "Variants (| separated)", "text", COL[0], Y + 120,
    "at sunrise | in a thunderstorm | under the milky way | in autumn fog")

fn("op_matrix", "prompt_matrix", COL[1], Y, label="① Expand matrix",
   required=("base_prompt",),
   data={"shared_tags": "cinematic, highly detailed, dramatic lighting"},
   outputs=[("out_p1", "prompt 1", "text"), ("out_p2", "prompt 2", "text"),
            ("out_p3", "prompt 3", "text"), ("out_p4", "prompt 4", "text"),
            ("out_labels", "labels", "text")])

for i in range(4):
    fn(f"op_grid_{i + 1}", "txt2img", COL[2], Y + i * 250,
       label=f"② Render {i + 1}",
       types={"steps": "number", "cfg_scale": "number", "seed": "number",
              "width": "number", "height": "number"},
       required=("prompt",),
       data={"negative_prompt": "", "steps": 4, "cfg_scale": 1.0,
             "seed": 1000 + i * 111, "width": 768, "height": 768,
             "model_id": T2I_MODEL},
       outputs=[("out_0", "image", "image")])
    link(f"op_matrix.out_p{i + 1}", f"op_grid_{i + 1}.in_prompt")
    link(f"op_grid_{i + 1}.out_0", f"op_sheet.in_image_{i + 1}")

fn("op_sheet", "contact_sheet", COL[3], Y + 320, label="③ Contact sheet",
   types={"image_1": "image", "image_2": "image", "image_3": "image",
          "image_4": "image", "columns": "number", "gap": "number"},
   data={"columns": 2, "gap": 16, "title": "Prompt matrix"},
   outputs=[("out_0", "grid", "image")])

out("sub_grid", "🧩 X/Y grid", "image", COL[4], Y + 320)

link("ref_matrix_base.out", "op_matrix.in_base_prompt")
link("ref_matrix_variants.out", "op_matrix.in_variations")
link("op_matrix.out_labels", "op_sheet.in_labels")
link("op_sheet.out_0", "sub_grid.in")

# ═════════════════════════════════════════════════════════════════════════════
# 8 · Extras — one upload, three post-processors (two local, one remote)
# ═════════════════════════════════════════════════════════════════════════════
Y = 3560
ref("ref_extras_image", "Extras input image", "image", COL[0], Y, sample("extras.jpg"))

fn("op_extras", "extras_upscale", COL[1], Y, label="① Upscale (local, instant)",
   types={"image": "image", "factor": "number", "sharpen": "number",
          "denoise": "boolean", "restore_contrast": "boolean"},
   required=("image",),
   data={"factor": 2.0, "method": "Lanczos", "sharpen": 0.45,
         "denoise": False, "restore_contrast": True},
   outputs=[("out_0", "image", "image"), ("out_1", "report", "text")])

space("op_aurasr", UPSCALE_SPACE, "/process_image", COL[1], Y + 300,
      label="② Upscale ×4 (AuraSR GAN)",
      inputs=[("input_image", "image", "image", True)],
      outputs=[("out_0", "Upscaled", "image", 1)])

space("op_rmbg", RMBG_SPACE, "/image", COL[1], Y + 500,
      label="③ Remove background (BRIA)",
      inputs=[("image", "image", "image", True)],
      outputs=[("out_0", "Cutout", "image", 1)])

out("sub_upscaled", "🔍 Upscaled (local)", "image", COL[2], Y)
out("sub_upscale_report", "📝 Upscale report", "text", COL[3], Y)
out("sub_aurasr", "🚀 Upscaled ×4 (GAN)", "image", COL[2], Y + 300)
out("sub_cutout", "✂ Background removed", "image", COL[2], Y + 620)

link("ref_extras_image.out", "op_extras.in_image")
link("op_extras.out_0", "sub_upscaled.in")
link("op_extras.out_1", "sub_upscale_report.in")
link("ref_extras_image.out", "op_aurasr.input_image")
link("op_aurasr.out_0", "sub_aurasr.in")
link("ref_extras_image.out", "op_rmbg.image")
link("op_rmbg.out_0", "sub_cutout.in")

# ═════════════════════════════════════════════════════════════════════════════
# 9 · ControlNet-style annotator previews (local)
# ═════════════════════════════════════════════════════════════════════════════
Y = 4340
ref("ref_control_image", "Annotator input", "image", COL[0], Y, sample("control.jpg"))

fn("op_control", "controlnet_preprocess", COL[1], Y, label="Annotator",
   types={"image": "image", "low_threshold": "number", "high_threshold": "number",
          "invert": "boolean", "blur": "number"},
   required=("image",),
   data={"mode": "Canny edges", "low_threshold": 60, "high_threshold": 160,
         "invert": False, "blur": 0.0},
   outputs=[("out_0", "map", "image")])

out("sub_control", "🕸 Annotator map", "image", COL[2], Y)

link("ref_control_image.out", "op_control.in_image")
link("op_control.out_0", "sub_control.in")

# ═════════════════════════════════════════════════════════════════════════════
# 10 · PNG Info — read generation parameters back out of a file
# ═════════════════════════════════════════════════════════════════════════════
Y = 4700
ref("ref_pnginfo_image", "PNG to inspect", "image", COL[0], Y,
    sample("with_parameters.png"))

fn("op_pnginfo", "png_info", COL[1], Y, label="Read PNG metadata",
   types={"image": "image"}, required=("image",),
   outputs=[("out_0", "report", "text"), ("out_1", "fields", "text")])

out("sub_png_report", "🧾 PNG info", "text", COL[2], Y)
out("sub_png_fields", "🧮 Parsed fields", "text", COL[3], Y)

link("ref_pnginfo_image.out", "op_pnginfo.in_image")
link("op_pnginfo.out_0", "sub_png_report.in")
link("op_pnginfo.out_1", "sub_png_fields.in")


# ═════════════════════════════════════════════════════════════════════════════
# Verification — catch wiring mistakes here, not at runtime
# ═════════════════════════════════════════════════════════════════════════════

from gradio.workflow import _INFERENCE_ENDPOINT_SCHEMAS  # noqa: E402


def verify():
    problems = []
    nodes = references + operators + subjects
    by_id = {}

    for n in nodes:
        if n["id"] in by_id:
            problems.append(f"duplicate node id: {n['id']}")
        by_id[n["id"]] = n

    # fn nodes: ports must mirror the Python signature exactly (positional call)
    for n in operators:
        if n["kind"] != "fn":
            continue
        func = N.BIND.get(n["fn"])
        if func is None:
            problems.append(f"{n['id']}: fn '{n['fn']}' is not in nodes.BIND")
            continue
        params = fn_params(func)
        labels = [p["label"] for p in n["inputs"]]
        if labels != params:
            problems.append(f"{n['id']}: port order {labels} != signature {params}")
        for key in n["data"]:
            if key not in {p["id"] for p in n["inputs"]}:
                problems.append(f"{n['id']}: data key '{key}' is not an input port")

    # model nodes: their ports must match the endpoint's canonical schema
    # EXACTLY. The canvas rewrites any model node whose ports differ, silently
    # dropping extra inputs (and orphaning the edges into them) the first time
    # the graph is opened in a browser. Anything needing a richer control
    # surface than the schema allows has to be an `fn` node calling
    # InferenceClient itself — that is why `txt2img` is one.
    for n in operators:
        if n["kind"] != "model":
            continue
        schema = _INFERENCE_ENDPOINT_SCHEMAS.get(n["endpoint"])
        if schema is None:
            problems.append(f"{n['id']}: unknown endpoint '{n['endpoint']}'")
            continue
        expected = [p["id"] for p in schema["inputs"]]
        actual = [p["id"] for p in n["inputs"]]
        if actual != expected:
            problems.append(
                f"{n['id']}: model ports {actual} != {n['endpoint']} schema "
                f"{expected} — the canvas would rewrite this node")

    # edges: endpoints must exist, and types must agree
    port_type = {}
    for n in nodes:
        for p in n.get("inputs", []):
            port_type[(n["id"], p["id"], "in")] = p["type"]
        for p in n.get("outputs", []):
            port_type[(n["id"], p["id"], "out")] = p["type"]

    fed = set()
    for e in edges:
        src = (e["from_node_id"], e["from_port_id"], "out")
        dst = (e["to_node_id"], e["to_port_id"], "in")
        if src not in port_type:
            problems.append(f"edge {e['id']}: no output port {src[0]}.{src[1]}")
            continue
        if dst not in port_type:
            problems.append(f"edge {e['id']}: no input port {dst[0]}.{dst[1]}")
            continue
        if dst in fed:
            problems.append(f"edge {e['id']}: {dst[0]}.{dst[1]} has two incoming edges")
        fed.add(dst)

        stype, dtype = port_type[src], port_type[dst]
        e["type"] = stype
        compatible = stype == dtype or "text" in (stype, dtype) and {stype, dtype} <= {
            "text", "number", "boolean", "json"}
        if not compatible:
            problems.append(
                f"edge {e['id']}: type mismatch {src[0]}.{src[1]}({stype}) "
                f"→ {dst[0]}.{dst[1]}({dtype})")

    # every subject must be fed, and every required input must be satisfied
    for s in subjects:
        if (s["id"], "in", "in") not in fed:
            problems.append(f"subject {s['id']} has no incoming edge")
    for n in operators:
        for p in n["inputs"]:
            if not p.get("required"):
                continue
            if (n["id"], p["id"], "in") not in fed and p["id"] not in n["data"]:
                problems.append(
                    f"{n['id']}: required input '{p['id']}' is neither wired nor defaulted")

    # nothing may be orphaned
    touched = {e["from_node_id"] for e in edges} | {e["to_node_id"] for e in edges}
    for n in nodes:
        if n["id"] not in touched:
            problems.append(f"orphan node: {n['id']}")

    return problems


if __name__ == "__main__":
    issues = verify()
    if issues:
        print(f"REFUSING TO WRITE — {len(issues)} problem(s):")
        for p in issues:
            print("  •", p)
        sys.exit(1)

    # Apply the curated layout. Positions in `layout.json` are the hand-arranged
    # ones (dragged in the canvas, then overlap-checked); the x/y computed above
    # are only a fallback for nodes the layout doesn't know about yet.
    placed = 0
    if os.path.exists(LAYOUT):
        with open(LAYOUT, encoding="utf-8") as f:
            layout = json.load(f)
        for node in references + operators + subjects:
            pos = layout.get(node["id"])
            if pos:
                node["x"], node["y"] = pos["x"], pos["y"]
                placed += 1
        missing = [n["id"] for n in references + operators + subjects
                   if n["id"] not in layout]
        if missing:
            print(f"  note: {len(missing)} node(s) not in layout.json, using "
                  f"generated positions: {', '.join(missing[:6])}")

    graph = {
        "schema_version": "2",
        "name": "Workflow1111 · Diffusion Studio",
        "references": references,
        "operators": operators,
        "subjects": subjects,
        "edges": edges,
    }
    with open(OUT, "w", encoding="utf-8") as f:
        json.dump(graph, f, indent=2, ensure_ascii=False)

    kinds = {}
    for o in operators:
        kinds[o["kind"]] = kinds.get(o["kind"], 0) + 1
    print(f"wrote {os.path.relpath(OUT, os.getcwd())}")
    print(f"  {len(references)} references, {len(operators)} operators "
          f"({', '.join(f'{v} {k}' for k, v in sorted(kinds.items()))}), "
          f"{len(subjects)} subjects, {len(edges)} edges")
    print(f"  {placed} node positions applied from layout.json")