"""Convert a ComfyUI UI workflow JSON to the API prompt format.

Best-effort: handles bypassed nodes (mode=4), frontend-only Note/Label/Reroute,
INT seed control_after_generate hidden widgets, and basic link-to-widget mapping.
Not a perfect replica of the frontend's graphToPrompt, but enough to run a real
workflow via POST /prompt for E2E smoke testing.
"""

from __future__ import annotations

import json
import sys
import urllib.request
from typing import Any, Dict, List, Optional, Tuple

FRONTEND_ONLY_TYPES = {
    "Note",
    "MarkdownNote",
    "Label (rgthree)",
    "Reroute",
    "PrimitiveNode",  # rarely — frontend injects value into consumer
    "Fast Groups Bypasser (rgthree)",
    "Fast Groups Muter (rgthree)",
    "Image Comparer (rgthree)",
    "ozW_BrandingNode",
    "PreviewAny",
    # KJNodes virtual variables — resolved via set_by_name in resolve_link.
    "SetNode",
    "GetNode",
    # rgthree anywhere node — frontend-only dispatcher
    "Anything Everywhere",
    "mxSlider",
}

# Nodes with a custom frontend widget layout where widgets_values position does
# NOT cleanly map to INPUT_TYPES declaration order. For those, OMIT all widget
# inputs in the API call and let the node fall back to its declared defaults.
IGNORE_WIDGETS_TYPES = {
    "ozW_SynthesizeAuthenticMetadata",
}

# Nodes whose widgets_values contain a dynamic list of LoRA entries (each a
# dict with on/lora/strength). The frontend serializes these as `lora_1`,
# `lora_2`, ... inputs in the API prompt via FlexibleOptionalInputType. Our
# converter must replicate that or the LoRAs are silently dropped — the node
# runs with no loras and produces a bypass identity model patch (i.e. the
# LoRA appears to not work at all).
DYNAMIC_LORA_NODE_TYPES = {
    "Power Lora Loader (rgthree)",
}

# Widget entries that are UI-only (videopreview state, previews, etc.) and
# should never be sent to the API — they are dict values with specific keys
# that identify them as frontend state.
def _is_ui_widget_state(w):
    if not isinstance(w, dict):
        return False
    ui_keys = ("hidden", "paused", "params", "videopreview")
    return any(k in w for k in ui_keys)

# Primitive widget types — these ARE allowed to consume widgets_values entries.
# Anything else (IMAGE, MODEL, CLIP, CONDITIONING, LATENT, MASK, VAE, SEGS, ...)
# is a connection-only input — never consumes from widgets_values.
PRIMITIVE_TYPES = {"INT", "FLOAT", "STRING", "BOOLEAN", "COMBO"}


def fetch_object_info(server: str = "http://localhost:8188") -> Dict[str, Any]:
    with urllib.request.urlopen(f"{server}/object_info", timeout=30) as r:
        return json.load(r)


SERVER_INJECTED_HIDDEN_TYPES = {
    "UNIQUE_ID",
    "EXTRA_PNGINFO",
    "PROMPT",
    "DYNPROMPT",
    "AUTH_TOKEN_COMFY_ORG",
    "API_KEY_COMFY_ORG",
}


def ordered_input_names(spec: Dict[str, Any]) -> List[Tuple[str, Any]]:
    """Return [(input_name, input_config), ...] in declaration order.

    Includes `hidden` inputs whose type is NOT server-injected (e.g.,
    `prompt_batch_data` of type STRING that Oz_RealityPromptGenerator uses).
    These are stored as widget values in the UI workflow and must be forwarded
    via the API prompt; otherwise they fall back to defaults at execution time.
    """
    inputs = spec.get("input", {}) or {}
    result: List[Tuple[str, Any]] = []
    for section in ("required", "optional"):
        for name, cfg in (inputs.get(section) or {}).items():
            result.append((name, cfg))
    for name, cfg in (inputs.get("hidden") or {}).items():
        # Skip inputs that ComfyUI's server injects automatically.
        if isinstance(cfg, str) and cfg in SERVER_INJECTED_HIDDEN_TYPES:
            continue
        if isinstance(cfg, (list, tuple)) and cfg and cfg[0] in SERVER_INJECTED_HIDDEN_TYPES:
            continue
        result.append((name, cfg))
    return result


def _first_type(cfg: Any) -> str:
    """Extract the type string from an INPUT_TYPES value."""
    if isinstance(cfg, list) and cfg:
        t = cfg[0]
        if isinstance(t, str):
            return t
        if isinstance(t, list):
            return "COMBO"
    if isinstance(cfg, tuple) and cfg:
        return str(cfg[0])
    return "UNKNOWN"


def convert(
    ui_wf: Dict[str, Any],
    object_info: Dict[str, Any],
    verbose: bool = False,
) -> Dict[str, Any]:
    api: Dict[str, Any] = {}
    links_by_id: Dict[int, List[Any]] = {L[0]: L for L in ui_wf.get("links", [])}
    nodes_by_id: Dict[int, Dict[str, Any]] = {n["id"]: n for n in ui_wf["nodes"]}

    # Build SetNode name → source link mapping for GetNode resolution.
    # KJNodes / cg-use-everywhere SetNode stashes its input under a name;
    # GetNode of the same name reads it back. These are virtual on the server —
    # the frontend rewrites connections to bypass them entirely. We do the
    # same here.
    set_by_name: Dict[str, int] = {}  # name -> SetNode's incoming link id
    for n in ui_wf.get("nodes", []):
        if n.get("type") != "SetNode":
            continue
        wv = n.get("widgets_values") or []
        name = wv[0] if wv and isinstance(wv[0], str) else None
        if not name:
            continue
        for inp in n.get("inputs") or []:
            if inp.get("link") is not None:
                set_by_name[name] = inp["link"]
                break

    # Build broadcast map from all "Anything Everywhere" nodes. These are
    # rgthree/cg-use-everywhere pseudo-nodes that broadcast a value of a given
    # type to every consumer of that type in the graph. The frontend resolves
    # these at queue time; we must do the same for headless conversion.
    broadcast_by_type: Dict[str, Tuple[str, int]] = {}
    for n in ui_wf.get("nodes", []):
        if n.get("type") not in ("Anything Everywhere", "Anything Everywhere?",
                                  "Anything Everywhere3", "Seed Everywhere",
                                  "Prompts Everywhere"):
            continue
        for inp in n.get("inputs") or []:
            lid = inp.get("link")
            if lid is None:
                continue
            itype = inp.get("type")
            if not itype or itype == "*":
                continue
            row = links_by_id.get(lid)
            if not row:
                continue
            src_id, src_slot = row[1], row[2]
            # Walk through Reroute/GetNode/SetNode like resolve_link does
            # (simple version — just take the direct source for now)
            broadcast_by_type[itype] = (str(src_id), src_slot)

    skipped_bypassed: List[int] = []
    skipped_frontend: List[int] = []
    skipped_unknown: List[int] = []

    def resolve_link(
        link_id: int, visited: Optional[set] = None
    ) -> Optional[Tuple[str, int]]:
        """Walk through Reroute and bypassed (mode=4) nodes to find a real source.

        For a bypassed node, we pass through its input of the SAME type at the
        same slot index — rgthree-style bypass semantics.
        """
        if visited is None:
            visited = set()
        if link_id in visited:
            return None
        visited.add(link_id)
        row = links_by_id.get(link_id)
        if not row:
            return None
        _lid, src_id, src_slot, _tgt_id, _tgt_slot, src_type = row
        src_node = nodes_by_id.get(src_id)
        if not src_node:
            return None

        # Reroute: one input, one output, walk upstream
        if src_node.get("type") == "Reroute":
            src_inputs = src_node.get("inputs") or []
            if src_inputs and src_inputs[0].get("link") is not None:
                return resolve_link(src_inputs[0]["link"], visited)
            return None

        # GetNode: virtual KJNodes variable — look up the matching SetNode
        # by name and recursively resolve that SetNode's stored source link.
        if src_node.get("type") == "GetNode":
            wv = src_node.get("widgets_values") or []
            name = wv[0] if wv and isinstance(wv[0], str) else None
            if name and name in set_by_name:
                return resolve_link(set_by_name[name], visited)
            return None

        # SetNode: pass-through to whatever is stored (its first input link).
        # Some SetNodes act as both stash AND passthrough — downstream edges
        # can go directly from SetNode's output.
        if src_node.get("type") == "SetNode":
            for inp in src_node.get("inputs") or []:
                if inp.get("link") is not None:
                    return resolve_link(inp["link"], visited)
            return None

        # Bypassed node: pass-through via same-type input (prefer same slot index)
        if src_node.get("mode") == 4:
            src_inputs = src_node.get("inputs") or []
            candidates = [
                (i, inp)
                for i, inp in enumerate(src_inputs)
                if inp.get("type") == src_type and inp.get("link") is not None
            ]
            if not candidates:
                return None
            # Prefer the input at the same index as the output slot, else first
            match = next(
                (inp for i, inp in candidates if i == src_slot),
                candidates[0][1],
            )
            return resolve_link(match["link"], visited)

        return (str(src_id), src_slot)

    for node in ui_wf["nodes"]:
        ntype = node.get("type", "")
        nid = str(node["id"])

        if node.get("mode") == 4:
            skipped_bypassed.append(node["id"])
            continue
        if ntype in FRONTEND_ONLY_TYPES:
            skipped_frontend.append(node["id"])
            continue

        spec = object_info.get(ntype)
        if not spec:
            skipped_unknown.append(node["id"])
            if verbose:
                print(f"WARN: unknown type {ntype} id={nid}")
            continue

        ordered = ordered_input_names(spec)

        # For nodes with broken widget→inputs mapping, only forward connection
        # inputs and let the node use declared defaults for every widget.
        ignore_widgets = ntype in IGNORE_WIDGETS_TYPES

        # Fresh per-node dicts — MUST be reset here (not below) because the
        # dict-typed widgets_values branch also consumes api_inputs and would
        # otherwise inherit stale state from the previous iteration.
        api_inputs: Dict[str, Any] = {}
        connection_by_name: Dict[str, int] = {}
        for inp in node.get("inputs") or []:
            if inp.get("link") is not None:
                connection_by_name[inp.get("name")] = inp["link"]

        # Dict-typed widgets_values (new ComfyUI schema) — map by NAME directly
        # into api_inputs, skipping UI-only state entries like videopreview.
        # Used by VHS_LoadVideo and some other nodes that moved to dict form.
        wv_raw = node.get("widgets_values")
        if isinstance(wv_raw, dict):
            if not ignore_widgets:
                for wk, wv in wv_raw.items():
                    if _is_ui_widget_state(wv):
                        continue
                    # Only forward names that the node declares as inputs
                    if any(name == wk for name, _ in ordered):
                        # Skip if the input is already a connection
                        if wk not in {n_.get("name") for n_ in (node.get("inputs") or []) if n_.get("link") is not None}:
                            api_inputs[wk] = wv
            # Still process connections + broadcast for unconnected inputs
            for name, cfg in ordered:
                if name in connection_by_name:
                    resolved = resolve_link(connection_by_name[name])
                    if resolved is not None:
                        api_inputs[name] = list(resolved)
                elif name not in api_inputs:
                    type_str = _first_type(cfg)
                    if (type_str not in PRIMITIVE_TYPES
                            and type_str in broadcast_by_type):
                        api_inputs[name] = list(broadcast_by_type[type_str])
            if ntype in DYNAMIC_LORA_NODE_TYPES:
                lora_idx = 1
                for w in wv_raw.values():
                    if (isinstance(w, dict) and "on" in w and "lora" in w
                            and "strength" in w):
                        api_inputs[f"lora_{lora_idx}"] = {
                            "on": w.get("on", True),
                            "lora": w.get("lora", ""),
                            "strength": w.get("strength", 1),
                            "strengthTwo": w.get("strengthTwo"),
                        }
                        lora_idx += 1
            api[nid] = {"class_type": ntype, "inputs": api_inputs}
            continue

        # connection_by_name already built above before dict branch

        wv = list(node.get("widgets_values") or [])
        wi = 0  # widget index

        # Drop UI-only state entries (videopreview dict) from the positional
        # list so we don't accidentally map them onto a real input slot.
        wv = [w for w in wv if not _is_ui_widget_state(w)]

        for name, cfg in ordered:
            type_str = _first_type(cfg)
            is_primitive = type_str in PRIMITIVE_TYPES

            connected = name in connection_by_name

            if connected:
                resolved = resolve_link(connection_by_name[name])
                if resolved is not None:
                    api_inputs[name] = list(resolved)
                # else: missing upstream — omit
            elif not is_primitive and type_str in broadcast_by_type:
                # Anything Everywhere broadcast: inject the shared source for
                # this type if the input is unconnected.
                api_inputs[name] = list(broadcast_by_type[type_str])

            # Primitive inputs have a SLOT in widgets_values even when converted
            # to a connection (ComfyUI keeps the stale value). So ALWAYS advance
            # the widget pointer for primitive inputs.
            if is_primitive and wi < len(wv):
                if not connected and not ignore_widgets:
                    api_inputs[name] = wv[wi]
                wi += 1
                # Handle control_after_generate hidden widget for seeds
                opts = cfg[1] if isinstance(cfg, list) and len(cfg) > 1 else {}
                if isinstance(opts, dict) and (
                    opts.get("control_after_generate") or name == "seed"
                ):
                    if wi < len(wv) and isinstance(wv[wi], str) and wv[wi] in (
                        "fixed", "increment", "decrement", "randomize"
                    ):
                        wi += 1

        # Ignore hidden inputs (unique_id, extra_pnginfo) — server injects them

        # Dynamic LoRA entries for nodes like rgthree Power Lora Loader.
        # These are stored as dict widgets at arbitrary positions in
        # widgets_values; extract each `{on, lora, strength}` and expose as
        # lora_1, lora_2, ... so the node sees them in its **kwargs.
        if ntype in DYNAMIC_LORA_NODE_TYPES:
            lora_idx = 1
            for w in node.get("widgets_values") or []:
                if (isinstance(w, dict) and "on" in w and "lora" in w
                        and "strength" in w):
                    api_inputs[f"lora_{lora_idx}"] = {
                        "on": w.get("on", True),
                        "lora": w.get("lora", ""),
                        "strength": w.get("strength", 1),
                        "strengthTwo": w.get("strengthTwo"),
                    }
                    lora_idx += 1

        api[nid] = {"class_type": ntype, "inputs": api_inputs}

    if verbose:
        print(f"bypassed skipped: {len(skipped_bypassed)}")
        print(f"frontend skipped: {len(skipped_frontend)}")
        print(f"unknown skipped: {len(skipped_unknown)}")
        print(f"api nodes: {len(api)}")
    return api


def main() -> None:
    if len(sys.argv) < 2:
        print("usage: ui_to_api.py <workflow.json> [output.json]")
        sys.exit(2)
    with open(sys.argv[1], encoding="utf-8") as f:
        ui = json.load(f)
    info = fetch_object_info()
    api = convert(ui, info, verbose=True)
    out = sys.argv[2] if len(sys.argv) > 2 else "api_prompt.json"
    with open(out, "w", encoding="utf-8") as f:
        json.dump(api, f, ensure_ascii=False, indent=2)
    print(f"wrote {out}: {len(api)} nodes")


if __name__ == "__main__":
    main()