"""Construye un workflow ComfyUI con ControlNet en API format (nodos numerados).

ControlNet condiciona la generacion con una imagen de control (canny, depth,
pose, scribble, ...). Cadena de nodos: CheckpointLoaderSimple + EmptyLatentImage
+ LoadImage (imagen de control) + ControlNetLoader -> ControlNetApply (inyecta
el control sobre el condicionamiento positivo) -> KSampler -> VAEDecode ->
SaveImage. Los CLIPTextEncode codifican el prompt positivo y el negativo.

API format: cada clave es un node_id (string); cada nodo tiene class_type +
inputs. Las conexiones entre nodos son listas [node_id, output_index]. Es el
formato que acepta POST /prompt, distinto del formato de la UI (graph con links).

Funcion pura: sin red, sin I/O. Determinista para los mismos argumentos.
"""


def comfyui_build_controlnet_workflow(
    ckpt_name: str,
    control_image: str,
    cn_name: str,
    positive: str,
    negative: str = "",
    *,
    strength: float = 1.0,
    steps: int = 20,
    cfg: float = 7.0,
    seed: int = 0,
    width: int = 512,
    height: int = 512,
) -> dict:
    """Construye el dict de un workflow txt2img guiado por ControlNet.

    Args:
        ckpt_name: nombre del checkpoint tal como lo ve el servidor ComfyUI
            (ej. "IMG_dreamshaper_8.safetensors"). Debe estar entre los que devuelve
            comfyui_object_info para CheckpointLoaderSimple.
        control_image: nombre del archivo de la imagen de control dentro de la
            carpeta input/ del servidor ComfyUI (lo carga el nodo LoadImage).
            Suele ser un mapa preprocesado (canny/depth/openpose).
        cn_name: nombre del modelo ControlNet en models/controlnet/ tal como lo
            lista comfyui_object_info para ControlNetLoader (control_net_name).
        positive: prompt positivo (lo que se quiere ver en la imagen).
        negative: prompt negativo (lo que se quiere evitar). Por defecto "".
        strength: fuerza con la que el ControlNet condiciona la generacion
            (0.0 = nula, 1.0 = plena). keyword-only.
        steps: pasos de sampling del KSampler. keyword-only.
        cfg: classifier-free guidance scale. keyword-only.
        seed: semilla del KSampler. 0 es determinista; cambiar para variar.
            keyword-only.
        width: ancho del latente/imagen en px (multiplo de 8). keyword-only.
        height: alto del latente/imagen en px (multiplo de 8). keyword-only.

    Returns:
        dict en API format listo para comfyui_submit_workflow. Las claves son
        node_ids y cada valor tiene class_type + inputs.
    """
    return {
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": ckpt_name},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1},
        },
        "10": {
            "class_type": "LoadImage",
            "inputs": {"image": control_image},
        },
        "12": {
            "class_type": "ControlNetLoader",
            "inputs": {"control_net_name": cn_name},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": positive, "clip": ["4", 1]},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative, "clip": ["4", 1]},
        },
        "13": {
            "class_type": "ControlNetApply",
            "inputs": {
                "conditioning": ["6", 0],
                "control_net": ["12", 0],
                "image": ["10", 0],
                "strength": strength,
            },
        },
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed,
                "steps": steps,
                "cfg": cfg,
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1.0,
                "model": ["4", 0],
                "positive": ["13", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "comfy_controlnet", "images": ["8", 0]},
        },
    }


if __name__ == "__main__":
    import json

    wf = comfyui_build_controlnet_workflow(
        ckpt_name="IMG_dreamshaper_8.safetensors",
        control_image="pose_canny.png",
        cn_name="control_v11p_sd15_canny.pth",
        positive="a knight in shining armor, dramatic lighting",
        negative="blurry, low quality",
        strength=0.8,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
