"""
ComfyUI runtime for Runware nodes.

One node per curated model, generated from the schema bundle (`build_node`),
every node running through the Python SDK over REST. Node logic lives once in
the base class; generated classes are thin data. Plus a generic untyped node
for models newer than the installed pack.

Coverage: image (IMAGE tensor in/out), video/audio/3D (saved file → path),
text (STRING). Nested params (lora, controlNet, providerSettings) ride in via an
`advanced_json` input.

torch / numpy / PIL / folder_paths come from the ComfyUI runtime; `runware` and
`pillow` from requirements.txt. API key from RUNWARE_API_KEY.
"""

from __future__ import annotations

import asyncio
import base64
import importlib.util
import io
import json
import os
import urllib.request
import uuid
from typing import Any


# ComfyUI capability probes (resolved once at import). Audio/video outputs use the
# richer native types when available and degrade to a saved file path otherwise. Audio
# decodes via soundfile (libsndfile, no FFmpeg) or torchaudio, whichever is installed.
_HAS_AUDIO = any(importlib.util.find_spec(m) is not None for m in ("soundfile", "torchaudio"))
_VideoFromFile: Any = None
if importlib.util.find_spec("comfy_api") is not None:
    # `comfy_api.input` is a back-compat shim that no longer re-exports VideoFromFile; the class
    # lives in `comfy_api.latest` on current ComfyUI. Try both so video keeps its native VIDEO type
    # across versions instead of silently degrading to a file path.
    for _mod in ("comfy_api.latest", "comfy_api.input"):
        try:
            _VideoFromFile = importlib.import_module(_mod).VideoFromFile  # type: ignore[attr-defined]
            break
        except Exception:  # noqa: BLE001
            _VideoFromFile = None
_HAS_VIDEO = _VideoFromFile is not None


# ----------------------------------------------------------------------- client identity


def _package_version() -> str:
    """Version from the shipped pyproject.toml (no tomllib import needed)."""
    pyproject = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pyproject.toml")
    try:
        with open(pyproject, encoding="utf-8") as fh:
            for line in fh:
                stripped = line.strip()
                if stripped.startswith("version") and "=" in stripped:
                    return stripped.split("=", 1)[1].strip().strip("'\"")
    except OSError:
        pass
    return "0.0.0"


# Passed to the SDK as its User-Agent prefix so Runware can attribute ComfyUI
# traffic: `runware-comfyui/<ver> runware-python/<sdk> (python/…) schemas/…`.
USER_AGENT_PREFIX = f"runware-comfyui/{_package_version()}"


_UI_API_KEY: str | None = None  # set from the ComfyUI Settings panel via /runware/set_key


def resolve_api_key() -> str | None:
    """RUNWARE_API_KEY env var, else the ComfyUI Settings key, else `runware auth login`."""
    key = os.environ.get("RUNWARE_API_KEY")
    if key:
        return key.strip()
    if _UI_API_KEY:
        return _UI_API_KEY
    try:
        with open(os.path.expanduser("~/.runware/config.yaml")) as fh:
            for line in fh:
                lower = line.strip().lower()
                if ":" in line and any(lower.startswith(k) for k in ("api_key", "apikey", "key", "token")):
                    return line.split(":", 1)[1].strip().strip("'\"")
    except OSError:
        pass
    return None


async def run_request(request: dict[str, Any]) -> list[dict[str, Any]]:
    from runware import Runware  # noqa: PLC0415

    api_key = resolve_api_key()
    if not api_key:
        raise RuntimeError("No Runware API key. Set RUNWARE_API_KEY or run `runware auth login`.")
    async with Runware(api_key=api_key, transport="rest", user_agent_prefix=USER_AGENT_PREFIX) as client:
        return await client.run(request)


def _clean_error(exc: BaseException) -> str:
    """A short, readable one-liner for a failed request, hiding the SDK/asyncio stack. Surfaces
    the Runware error's parameter / code / HTTP status when present so the message is actionable."""
    msg = str(getattr(exc, "message", None) or exc).strip() or exc.__class__.__name__
    tail: list[str] = []
    param = getattr(exc, "parameter", None)
    code = getattr(exc, "code", None)
    status = getattr(exc, "status_code", None)
    if param:
        tail.append(f"parameter '{param}'")
    if code and str(code).lower() not in msg.lower():
        tail.append(f"code {code}")
    if status:
        tail.append(f"HTTP {status}")
    return "Runware: " + msg + (f"  [{', '.join(tail)}]" if tail else "")


async def run_request_clean(request: dict[str, Any]) -> list[dict[str, Any]]:
    """Run a request, but on failure raise a clean single-line error instead of the deep
    SDK/asyncio traceback ComfyUI would otherwise dump. Set RUNWARE_DEBUG for the full stack."""
    try:
        return await run_request(request)
    except Exception as exc:  # noqa: BLE001
        if os.environ.get("RUNWARE_DEBUG"):
            import traceback  # noqa: PLC0415
            traceback.print_exc()
        raise RuntimeError(_clean_error(exc)) from None


# ---------------------------------------------------------------------------- imaging


def _imaging() -> tuple[Any, Any, Any]:
    import numpy as np  # noqa: PLC0415
    import torch  # noqa: PLC0415
    from PIL import Image  # noqa: PLC0415
    return torch, np, Image


def tensor_to_data_uris(image: Any) -> list[str]:
    """ComfyUI IMAGE tensor [B,H,W,C] float 0-1 → list of PNG data URIs."""
    _, np, Image = _imaging()
    arr = (image.clamp(0, 1).cpu().numpy() * 255.0).round().astype("uint8")
    uris: list[str] = []
    for frame in arr:
        buf = io.BytesIO()
        Image.fromarray(frame).save(buf, format="PNG")
        uris.append("data:image/png;base64," + base64.b64encode(buf.getvalue()).decode())
    return uris


def urls_to_image(urls: list[str]) -> Any:
    torch, np, Image = _imaging()
    if not urls:
        return torch.zeros((1, 512, 512, 3))
    frames = []
    for url in urls:
        data = urllib.request.urlopen(url).read()  # noqa: S310
        pil = Image.open(io.BytesIO(data)).convert("RGB")
        frames.append(np.asarray(pil).astype("float32") / 255.0)
    return torch.from_numpy(np.stack(frames))


def mask_to_data_uri(mask: Any) -> str:
    """ComfyUI MASK [B,H,W] (or [H,W]) float 0-1 → grayscale PNG data URI."""
    _, np, Image = _imaging()
    arr = mask.clamp(0, 1).cpu().numpy()
    if arr.ndim == 3:
        arr = arr[0]
    arr = (arr * 255.0).round().astype("uint8")
    buf = io.BytesIO()
    Image.fromarray(arr, mode="L").save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


# ------------------------------------------------------------------------- file saving


def _output_dir() -> str:
    try:
        import folder_paths  # noqa: PLC0415  (provided by ComfyUI)
        return folder_paths.get_output_directory()
    except Exception:  # noqa: BLE001
        import tempfile  # noqa: PLC0415
        return tempfile.gettempdir()


def _download(url: str, suffix: str) -> str:
    path = os.path.join(_output_dir(), f"runware_{uuid.uuid4().hex[:8]}{suffix}")
    with open(path, "wb") as fh:
        fh.write(urllib.request.urlopen(url).read())  # noqa: S310
    return path


def _media_url(result: dict[str, Any]) -> str | None:
    for key in ("videoURL", "audioURL", "imageURL", "URL", "url"):
        if result.get(key):
            return result[key]
    outs = result.get("outputs")
    if isinstance(outs, dict):
        files = outs.get("files")
        if isinstance(files, list) and files and isinstance(files[0], dict):
            return files[0].get("url")
    return None


def _audio_output(url: str) -> Any:
    """Decode a returned audio file into a ComfyUI AUDIO dict {waveform, sample_rate}.
    Prefers soundfile (libsndfile, no FFmpeg dependency) and falls back to torchaudio."""
    import torch  # noqa: PLC0415
    path = _download(url, _SUFFIX.get("audioInference", ".mp3"))
    try:  # soundfile: self-contained libsndfile, decodes WAV/FLAC/OGG (and MP3 on libsndfile >= 1.1)
        import numpy as np  # noqa: PLC0415
        import soundfile as sf  # noqa: PLC0415
        data, sample_rate = sf.read(path, dtype="float32", always_2d=True)  # (frames, channels)
        return {"waveform": torch.from_numpy(np.ascontiguousarray(data.T)).unsqueeze(0),
                "sample_rate": int(sample_rate)}
    except Exception:  # noqa: BLE001
        pass
    try:
        import torchaudio  # noqa: PLC0415
        waveform, sample_rate = torchaudio.load(path)
        return {"waveform": waveform.unsqueeze(0), "sample_rate": int(sample_rate)}
    except Exception as exc:  # noqa: BLE001
        print(f"[Runware] audio decode failed ({exc}); install `soundfile`. File saved at {path}")
        return {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}


def _video_output(url: str) -> Any:
    """Wrap a returned video file as a ComfyUI VIDEO object, or fall back to its path."""
    path = _download(url, _SUFFIX.get("videoInference", ".mp4"))
    if _VideoFromFile is not None:
        try:
            return _VideoFromFile(path)
        except Exception as exc:  # noqa: BLE001
            print(f"[Runware] video wrap failed ({exc}); file saved at {path}")
    return path


def _set_path(obj: dict[str, Any], path: list[str], value: Any) -> None:
    cur = obj
    for key in path[:-1]:
        cur = cur.setdefault(key, {})
    cur[path[-1]] = value


def _total_cost(results: list[dict[str, Any]]) -> float:
    """Sum the API cost across results (present when includeCost was requested)."""
    try:
        return sum(float(r.get("cost") or 0) for r in results if isinstance(r, dict))
    except (TypeError, ValueError):
        return 0.0


def _socket_fields_from_inputs(input_types: dict[str, Any]) -> dict[str, set[str]]:
    """Per array-feature socket, the item fields this model accepts, read from the `rw_socket_fields`
    the generator placed on the socket in INPUT_TYPES."""
    out: dict[str, set[str]] = {}
    for scope in ("required", "optional"):
        for name, spec in (input_types.get(scope) or {}).items():
            if isinstance(spec, (list, tuple)) and len(spec) == 2 and isinstance(spec[1], dict):
                fields = spec[1].get("rw_socket_fields")
                if fields:
                    out[name] = set(fields)
    return out


def _strip_unsupported_fields(request: dict[str, Any], socket_fields: dict[str, set[str]]) -> None:
    """Drop fields a model does not declare from its feature values. A shared builder is the union
    of every model's fields, so it can produce more than a given model accepts: array features
    (ControlNet/IP-Adapter/LoRA) strip per item, single-object features (Speech) strip the object."""
    for slot, allowed in socket_fields.items():
        value = request.get(slot)
        if isinstance(value, list):
            request[slot] = [
                {k: v for k, v in item.items() if k in allowed} if isinstance(item, dict) else item
                for item in value
            ]
        elif isinstance(value, dict):
            request[slot] = {k: v for k, v in value.items() if k in allowed}


def _run_info(results: list[dict[str, Any]]) -> str:
    """A short on-node summary of a run: total cost and whether the safety check flagged
    anything. Each field appears only when the response carried it (cost needs includeCost;
    NSFWContent is present only when the model ran a content check)."""
    parts: list[str] = []
    cost = _total_cost(results)
    if cost:
        parts.append("$" + f"{cost:.6f}".rstrip("0").rstrip("."))
    flags = [r.get("NSFWContent") for r in results if isinstance(r, dict) and "NSFWContent" in r]
    if flags:
        parts.append("NSFW: " + ("yes" if any(flags) else "no"))
    return "  ·  ".join(parts)  # middot separates cost from the NSFW flag


_SIZE_TRIAD = ("width", "height", "resolution")


def _apply_param_plan(target: dict[str, Any], kwargs: dict[str, Any], plan: dict[str, dict[str, Any]]) -> None:
    """Map widget inputs onto a request/item dict per the generated param plan."""
    size_choices: list[tuple[Any, dict[str, Any]]] = []
    for name, value in kwargs.items():
        if value is None:
            continue
        p = plan.get(name, {"path": [name], "kind": "scalar"})
        path, kind = p["path"], p["kind"]
        if p.get("gated_by") and not kwargs.get(p["gated_by"]):
            continue  # its group gate (e.g. `safety`) is off → leave the field to the model
        if kind in ("skip", "group_gate"):  # a gate's value widget, or a group toggle: never sent
            continue
        if kind == "size_choice":  # applied last: authoritative over the size triad
            size_choices.append((value, p))
            continue
        if kind == "scalar":
            if p.get("omit_empty") and isinstance(value, str) and not value.strip():
                continue  # empty optional string → let the model default
            if "omit_value" in p and value == p["omit_value"]:
                continue  # the "(default)" sentinel → omit
            _set_path(target, path, value)
        elif kind == "gate":  # a reveal toggle: send the paired value only when on
            if value:
                paired = kwargs.get(p["value_key"])
                if paired is not None:
                    _set_path(target, path, paired)
        elif kind == "bool_tri":  # (default) / enabled / disabled
            if value == "enabled":
                _set_path(target, path, True)
            elif value == "disabled":
                _set_path(target, path, False)
        elif kind == "feature":  # a builder's typed output (a list or object), placed as-is
            _set_path(target, path, value)
        elif kind == "media_string":
            if isinstance(value, str) and value.strip():
                _set_path(target, path, value.strip())
        elif kind == "media_array":
            if isinstance(value, str) and value.strip():
                _set_path(target, path, [line.strip() for line in value.splitlines() if line.strip()])
        elif kind == "image_scalar":
            _set_path(target, path, tensor_to_data_uris(value)[0])
        elif kind == "mask_scalar":
            _set_path(target, path, mask_to_data_uri(value))
        elif kind == "image_array":
            _set_path(target, path, tensor_to_data_uris(value))
        elif kind == "image_objects":
            type_value = p.get("type")
            _set_path(target, path, [
                {"image": u, **({"type": type_value} if type_value else {})}
                for u in tensor_to_data_uris(value)
            ])

    # A `size` dropdown folds the mutually-exclusive resolution/width/height into one
    # choice. Apply it after everything else and let it own the whole triad: a concrete
    # size clears the others so we never send a conflicting pair. An empty choice ("use
    # width and height") leaves the width/height widgets that were set above untouched.
    for value, p in size_choices:
        chosen = (p.get("options") or {}).get(value)
        if not isinstance(chosen, dict):
            continue
        concrete = {k: v for k, v in chosen.items() if not k.startswith("__")}
        if chosen.get("__none__") or concrete:  # explicit "None", or a concrete size: own the triad
            for k in _SIZE_TRIAD:
                target.pop(k, None)
            target.update(concrete)
        # an empty choice ("Custom") leaves the width/height widgets set above untouched


# --------------------------------------------------------------------------- base node


# Output classification by taskType (the `inference` list spans more than the 5
# modalities: it also carries image-producing operations and caption/vectorize).
_IMAGE_TASKS = {"imageInference", "upscale", "removeBackground", "imageMasking", "controlNetPreprocess"}
_TEXT_TASKS = {"textInference", "caption", "promptEnhance"}
_AUDIO_TASKS = {"audioInference"}
_VIDEO_TASKS = {"videoInference"}
_SUFFIX = {"videoInference": ".mp4", "audioInference": ".mp3", "3dInference": ".glb", "vectorize": ".svg"}
_LABEL = {
    "imageInference": "Image", "upscale": "Image", "removeBackground": "Image",
    "imageMasking": "Image", "controlNetPreprocess": "Image",
    "videoInference": "Video", "audioInference": "Audio", "3dInference": "3D",
    "textInference": "Text", "caption": "Text", "promptEnhance": "Text",
    "vectorize": "Vector", "training": "Training",
}
_FILE_RET = (("STRING",), ("file_path",))
_RET = {
    "image": (("IMAGE",), ("image",)),
    "text": (("STRING",), ("text",)),
    "audio": (("AUDIO",), ("audio",)) if _HAS_AUDIO else _FILE_RET,
    "video": (("VIDEO",), ("video",)) if _HAS_VIDEO else _FILE_RET,
    "file": _FILE_RET,
}


def _output_kind(task_type: str) -> str:
    if task_type in _IMAGE_TASKS:
        return "image"
    if task_type in _TEXT_TASKS:
        return "text"
    if task_type in _AUDIO_TASKS:
        return "audio"
    if task_type in _VIDEO_TASKS:
        return "video"
    return "file"


class _RunwareNode:
    MODEL = ""
    TASK_TYPE = ""
    PARAM_PLAN: dict[str, dict[str, Any]] = {}
    _INPUT_TYPES: dict[str, Any] = {}
    _SOCKET_FIELDS: dict[str, set[str]] = {}  # per array-feature socket, item fields this model accepts

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "execute"
    CATEGORY = "Runware"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return cls._INPUT_TYPES

    async def execute(self, **kwargs: Any) -> Any:
        request = await asyncio.to_thread(self._build_request, kwargs)
        results = await run_request_clean(request)
        if os.environ.get("RUNWARE_DEBUG"):
            keys = sorted({k for r in results if isinstance(r, dict) for k in r})
            nsfw = [r.get("NSFWContent") for r in results if isinstance(r, dict)]
            print(f"[Runware] {self.TASK_TYPE} sent safety={request.get('safety')} | "
                  f"result keys={keys} | NSFWContent={nsfw}")
        outputs = await asyncio.to_thread(self._outputs, results)
        info = _run_info(results)  # cost + NSFW flag, shown on the node
        if info:
            print(f"[Runware] {self.TASK_TYPE}: {info}")
            return {"ui": {"runware_info": [info]}, "result": outputs}
        keys = sorted({k for r in results if isinstance(r, dict) for k in r})
        print(f"[Runware] {self.TASK_TYPE}: no cost/NSFW in response; result keys = {keys}")
        return outputs

    def _build_request(self, kwargs: dict[str, Any]) -> dict[str, Any]:
        request: dict[str, Any] = {"model": self.MODEL, "taskType": self.TASK_TYPE, "includeCost": True}
        advanced = kwargs.pop("advanced_json", "") or ""

        # Widgets, media, images, and connected feature sockets all flow through the plan.
        # Params governed by a validation rule are off-by-default toggles in the node, so a
        # default request never breaks a rule and nothing is dropped silently.
        _apply_param_plan(request, kwargs, self.PARAM_PLAN)

        # A ControlNet/IP-Adapter builder is shared across models, so it can carry item fields
        # this model does not accept (architectures expand the base object). Drop those, so a
        # builder wired into a model that only has the base fields sends a request it accepts.
        _strip_unsupported_fields(request, self._SOCKET_FIELDS)

        # Raw advanced JSON is the user's escape hatch and wins on conflict.
        if advanced.strip():
            try:
                data = json.loads(advanced)
                if isinstance(data, dict):
                    request.update(data)
            except json.JSONDecodeError:
                pass
        return request

    def _outputs(self, results: list[dict[str, Any]]) -> tuple[Any]:
        results = [r for r in results if isinstance(r, dict)]
        kind = _output_kind(self.TASK_TYPE)
        if kind == "image":
            return (urls_to_image([r["imageURL"] for r in results if r.get("imageURL")]),)
        if kind == "text":
            return ("\n".join(r.get("text", "") for r in results if r.get("text")),)
        if kind == "audio" and _HAS_AUDIO:
            for r in results:
                url = _media_url(r)
                if url:
                    return (_audio_output(url),)
        if kind == "video" and _HAS_VIDEO:
            for r in results:
                url = _media_url(r)
                if url:
                    return (_video_output(url),)
        for r in results:  # 3d / vector / other, and audio/video without the native type
            url = _media_url(r)
            if url:
                return (_download(url, _SUFFIX.get(self.TASK_TYPE, ".bin")),)
        return ("",)


# ----------------------------------------------------------------------- generic node


class RunwareCustom:
    # Generic escape hatch: any model + taskType, any modality. It makes no assumption about what the
    # model takes or returns; it sends a raw request and hands back the raw response as JSON. Feed
    # media in as UUIDs/URLs in the JSON (see Runware Upload Image), and pull results out of the JSON
    # with Runware Get, then wire those to a loader / Preview / Save node.
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("result_json",)
    FUNCTION = "execute"
    CATEGORY = "Runware/Custom"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return {
            "required": {
                "model": ("STRING", {"default": "runware:100@1", "tooltip": "Model AIR"}),
                "taskType": ("STRING", {"default": "imageInference"}),
            },
            "optional": {
                "request_json": ("STRING", {"multiline": True, "default": "{}",
                                            "tooltip": "The full request body as JSON, merged over model/taskType. "
                                                       "Put any media inputs here as UUIDs/URLs (see Runware Upload Image)."}),
            },
        }

    async def execute(self, model: str, taskType: str, request_json: str = "{}") -> Any:
        request: dict[str, Any] = {"model": model, "taskType": taskType, "includeCost": True}
        try:
            data = json.loads(request_json or "{}")
            if isinstance(data, dict):
                request.update(data)
        except json.JSONDecodeError:
            pass

        results = await run_request_clean(request)
        results = [r for r in results if isinstance(r, dict)]
        outputs = (json.dumps(results),)
        info = _run_info(results)
        return {"ui": {"runware_info": [info]}, "result": outputs} if info else outputs


# ----------------------------------------------------------- generated param builders


class _RunwareBuilder:
    """Builds one feature (LoRA, ControlNet, PuLID, ...) and outputs its own typed
    value: a list for stackable features (chained), an object for single ones.
    Generated from the feature's schema; wired into the model node's typed socket."""

    SLOT = ""  # the feature's socket name; for arrays, also the chain input
    IS_ARRAY = False
    REQUIRED: list[str] = []  # fields the feature can't be built without
    PARAM_PLAN: dict[str, dict[str, Any]] = {}
    _INPUT_TYPES: dict[str, Any] = {}

    RETURN_TYPES = ("RUNWARE_PARAMS",)
    RETURN_NAMES = ("params",)
    FUNCTION = "build"
    CATEGORY = "Runware/Params"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return cls._INPUT_TYPES

    def build(self, **kwargs: Any) -> tuple[Any]:
        chain = kwargs.pop(self.SLOT, None) if self.IS_ARRAY else None
        advanced = kwargs.pop("advanced_json", "") or ""
        item: dict[str, Any] = {}
        _apply_param_plan(item, kwargs, self.PARAM_PLAN)
        if advanced.strip():
            try:
                data = json.loads(advanced)
                if isinstance(data, dict):
                    item.update(data)
            except json.JSONDecodeError:
                pass
        # An unfilled required field (unconnected image, empty model, ...) would make the
        # API reject the whole request. Drop the incomplete feature instead of sending it.
        complete = all(item.get(f) not in (None, "", [], {}) for f in self.REQUIRED)
        if self.IS_ARRAY:
            items = list(chain) if isinstance(chain, list) else []
            if complete:
                items.append(item)
            return (items or None,)
        return (item if complete else None,)


# ----------------------------------------------------------------------- upload utility


async def run_upload(image_uri: str) -> str:
    from runware import Runware  # noqa: PLC0415

    api_key = resolve_api_key()
    if not api_key:
        raise RuntimeError("No Runware API key. Set RUNWARE_API_KEY or run `runware auth login`.")
    async with Runware(api_key=api_key, transport="rest", user_agent_prefix=USER_AGENT_PREFIX) as client:
        for r in await client.media_storage({"operation": "upload", "media": image_uri}):
            if isinstance(r, dict) and r.get("mediaUUID"):
                return str(r["mediaUUID"])
    return ""


class RunwareUploadImage:
    """Upload an image once -> reusable UUID (for media inputs or the custom node)."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("image_uuid",)
    FUNCTION = "execute"
    CATEGORY = "Runware/Params"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return {"required": {"image": ("IMAGE",)}}

    async def execute(self, image: Any) -> tuple[str]:
        uris = await asyncio.to_thread(tensor_to_data_uris, image)
        return (await run_upload(uris[0]),)


class RunwareLoadImage:
    """Download an image URL (or a Runware result URL) into an IMAGE tensor."""

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "execute"
    CATEGORY = "Runware/Params"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return {"required": {"url": ("STRING", {"default": "", "tooltip": "Image URL"})}}

    async def execute(self, url: str) -> tuple[Any]:
        url = (url or "").strip()
        return (await asyncio.to_thread(urls_to_image, [url] if url else []),)


def _first_url(data: Any) -> str | None:
    """First URL/data-URI string found anywhere in a decoded response (depth-first)."""
    if isinstance(data, str):
        return data if data.startswith(("http://", "https://", "data:")) else None
    if isinstance(data, dict):
        data = list(data.values())
    if isinstance(data, list):
        for v in data:
            found = _first_url(v)
            if found:
                return found
    return None


def _json_get(data: Any, path: str) -> Any:
    """Walk a dot-path into decoded JSON; numeric segments index into lists. An empty path returns
    the first URL found (the common single-output case)."""
    segs = [p for p in (path or "").split(".") if p]
    if not segs:
        return _first_url(data)
    cur = data
    for seg in segs:
        if isinstance(cur, list):
            try:
                cur = cur[int(seg)]
            except (ValueError, IndexError):
                return None
        elif isinstance(cur, dict):
            cur = cur.get(seg)
        else:
            return None
    return cur


class RunwareGet:
    """Pull a field out of a Runware result_json by dot-path, e.g. `0.imageURL`, `0.videoURL`, or
    `0.outputs.files.0.url` (the leading index picks the task in the result list; numeric segments
    index lists). Leave path empty to grab the first URL found. Output is a STRING you feed to a
    loader such as Runware Load Image (URL), or to a Save/Preview node."""

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("value",)
    FUNCTION = "execute"
    CATEGORY = "Runware/Params"

    @classmethod
    def INPUT_TYPES(cls) -> dict[str, Any]:
        return {
            "required": {"result_json": ("STRING", {"forceInput": True})},
            "optional": {
                "path": ("STRING", {"default": "", "tooltip": "Dot-path into the response, e.g. "
                                    "0.imageURL or 0.outputs.files.0.url. Empty = first URL found."}),
            },
        }

    def execute(self, result_json: str, path: str = "") -> tuple[str]:
        try:
            data = json.loads(result_json or "null")
        except json.JSONDecodeError:
            return ("",)
        value = _json_get(data, path)
        if value is None:
            return ("",)
        return (value if isinstance(value, str) else json.dumps(value),)


# ----------------------------------------------------------------------------- factory


_SPECS_PATH = os.path.join(os.path.dirname(__file__), "nodes.json")


_EXTRA_NODES = {
    "RunwareCustom": (RunwareCustom, "Runware (custom)"),
    "RunwareGet": (RunwareGet, "Runware Get"),
    "RunwareUploadImage": (RunwareUploadImage, "Runware Upload Image"),
    "RunwareLoadImage": (RunwareLoadImage, "Runware Load Image (URL)"),
}


def _build() -> tuple[dict[str, Any], dict[str, str]]:
    classes: dict[str, Any] = {k: v[0] for k, v in _EXTRA_NODES.items()}
    names: dict[str, str] = {k: v[1] for k, v in _EXTRA_NODES.items()}

    with open(_SPECS_PATH) as fh:
        data = json.load(fh)

    for spec in data["models"]:
        task_type = spec["task_type"]
        rtypes, rnames = _RET[_output_kind(task_type)]
        key = spec["key"]
        classes[key] = type(key, (_RunwareNode,), {
            "MODEL": spec["model"],
            "TASK_TYPE": task_type,
            "PARAM_PLAN": spec["param_plan"],
            "_INPUT_TYPES": spec["input_types"],
            "_SOCKET_FIELDS": _socket_fields_from_inputs(spec["input_types"]),
            "RETURN_TYPES": rtypes,
            "RETURN_NAMES": rnames,
            "CATEGORY": f"Runware/{_LABEL.get(task_type, 'Other')}/{spec['creator']}",
        })
        names[key] = spec["name"]

    for spec in data.get("architectures", []):
        task_type = spec["task_type"]
        rtypes, rnames = _RET[_output_kind(task_type)]
        key = spec["key"]
        classes[key] = type(key, (_RunwareNode,), {
            "MODEL": spec.get("default_model", ""),  # the model widget overrides this per run
            "TASK_TYPE": task_type,
            "PARAM_PLAN": spec["param_plan"],
            "_INPUT_TYPES": spec["input_types"],
            "_SOCKET_FIELDS": _socket_fields_from_inputs(spec["input_types"]),
            "RETURN_TYPES": rtypes,
            "RETURN_NAMES": rnames,
            "CATEGORY": "Runware/Custom models",
        })
        names[key] = spec["name"]

    for b in data.get("builders", []):
        input_types = b["input_types"]
        slot, typ = b["slot"], b["type"]
        if b["is_array"]:  # stackable features chain through a same-typed input
            input_types.setdefault("optional", {})[slot] = (typ,)
        key = b["key"]
        classes[key] = type(key, (_RunwareBuilder,), {
            "SLOT": slot,
            "IS_ARRAY": b["is_array"],
            "REQUIRED": b.get("required", []),
            "PARAM_PLAN": b["param_plan"],
            "_INPUT_TYPES": input_types,
            "RETURN_TYPES": (typ,),
            "RETURN_NAMES": (slot,),
        })
        names[key] = b["name"]

    # Two different models can share a display name (e.g. Vidu Q1 ships as both a video and an image
    # model). They register fine under distinct keys, but the identical label is confusing on search,
    # so disambiguate a collision by the node's modality: "Vidu Q1 (Video)" / "Vidu Q1 (Image)".
    by_label: dict[str, list[str]] = {}
    for key, label in names.items():
        by_label.setdefault(label, []).append(key)
    for label, keys in by_label.items():
        if len(keys) < 2:
            continue
        for key in keys:
            mod = _LABEL.get(getattr(classes[key], "TASK_TYPE", ""))
            if mod:
                names[key] = f"{label} ({mod})"

    return classes, names


NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = _build()


# ------------------------------------------------------------- frontend search route


async def search_models(params: dict[str, Any]) -> list[dict[str, str]]:
    """Run a catalog model search, flattened to [{air, name}] for the picker."""
    from runware import Runware  # noqa: PLC0415

    api_key = resolve_api_key()
    if not api_key:
        raise RuntimeError("No Runware API key. Set RUNWARE_API_KEY or run `runware auth login`.")
    async with Runware(api_key=api_key, transport="rest", user_agent_prefix=USER_AGENT_PREFIX) as client:
        results = await client.model_search(params)
    models: list[dict[str, str]] = []
    for r in results:
        if not isinstance(r, dict):
            continue
        for m in r.get("results") or []:
            if isinstance(m, dict):
                air = m.get("air") or m.get("model") or m.get("id")
                if air:
                    models.append({"air": str(air), "name": str(m.get("name") or air)})
    return models


_ROUTES_REGISTERED = False


def _register_routes() -> None:
    """Register the /runware/model_search HTTP route (no-op outside ComfyUI)."""
    global _ROUTES_REGISTERED
    if _ROUTES_REGISTERED:
        return
    try:
        from server import PromptServer  # noqa: PLC0415
        from aiohttp import web  # noqa: PLC0415
    except Exception:  # noqa: BLE001
        return

    @PromptServer.instance.routes.post("/runware/model_search")  # type: ignore[misc]
    async def _model_search_route(request: Any) -> Any:
        try:
            body = await request.json()
        except Exception:  # noqa: BLE001
            body = {}
        params: dict[str, Any] = {"search": str(body.get("search") or ""), "limit": int(body.get("limit") or 40)}
        for key in ("category", "type", "conditioning", "architecture"):
            if body.get(key):
                params[key] = body[key]
        try:
            models = await search_models(params)
        except Exception as exc:  # noqa: BLE001
            return web.json_response({"error": str(exc)}, status=500)
        return web.json_response({"models": models})

    @PromptServer.instance.routes.post("/runware/set_key")  # type: ignore[misc]
    async def _set_key_route(request: Any) -> Any:
        global _UI_API_KEY
        try:
            body = await request.json()
        except Exception:  # noqa: BLE001
            body = {}
        key = str(body.get("key") or "").strip()
        _UI_API_KEY = key or None
        return web.json_response({"ok": True})

    _ROUTES_REGISTERED = True


_register_routes()