import uuid
from typing import Any

from comfy_server.models import Input


def build_fal_workflow(input: Input) -> dict[str, Any]:
    load_image_nodes = ["LoadImage"]

    # Start with fal_inputs from the proxy (for video/audio uploads)
    fal_inputs = dict(input.fal_inputs) if input.fal_inputs else {}
    fal_inputs_dev_info = dict(input.fal_inputs_dev_info) if input.fal_inputs_dev_info else {}

    # Also scan prompt for image URLs
    for node_id, node_data in input.prompt.items():
        if node_data.get("class_type") in load_image_nodes:
            image_uri = node_data.get("inputs", {}).get("image", "")
            if (
                image_uri
                and image_uri.startswith("https://")
                or image_uri.startswith("data:image")
            ):
                field = str(uuid.uuid4())
                fal_inputs_dev_info[field] = {
                    "key": [node_id, "inputs", "image"],
                    "class_type": "LoadImage",
                }
                fal_inputs[field] = image_uri

    return {
        "prompt": input.prompt,
        "fal_inputs": fal_inputs,
        "fal_inputs_dev_info": fal_inputs_dev_info,
    }
