"""Construye un workflow ComfyUI img2img en API format (dict de nodos numerados).

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_img2img_workflow(
    ckpt_name: str,
    init_image: str,
    positive: str,
    negative: str = "",
    *,
    denoise: float = 0.6,
    steps: int = 20,
    cfg: float = 7.0,
    seed: int = 0,
    sampler_name: str = "euler",
    scheduler: str = "normal",
) -> dict:
    """Construye el dict de un workflow img2img para SD1.5 / SDXL.

    Cadena de nodos: CheckpointLoaderSimple + LoadImage -> VAEEncode ->
    KSampler (con denoise < 1.0 para conservar la imagen base) -> VAEDecode ->
    SaveImage. CLIPTextEncode codifica el prompt positivo y el negativo.

    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.
        init_image: nombre del archivo de imagen base dentro de la carpeta
            input/ del servidor ComfyUI (lo que carga el nodo LoadImage).
        positive: prompt positivo (lo que se quiere ver en la imagen).
        negative: prompt negativo (lo que se quiere evitar). Por defecto "".
        denoise: fuerza de denoising del KSampler (0.0 = identica a la base,
            1.0 = ignora la base). Tipico 0.4-0.7 para img2img. keyword-only.
        steps: pasos de sampling del KSampler. keyword-only.
        cfg: classifier-free guidance scale. keyword-only.
        seed: semilla del KSampler. keyword-only.
        sampler_name: nombre del sampler (ej. "euler", "dpmpp_2m"). keyword-only.
        scheduler: scheduler del sampler (ej. "normal", "karras"). 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},
        },
        "10": {
            "class_type": "LoadImage",
            "inputs": {"image": init_image},
        },
        "11": {
            "class_type": "VAEEncode",
            "inputs": {"pixels": ["10", 0], "vae": ["4", 2]},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": positive, "clip": ["4", 1]},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative, "clip": ["4", 1]},
        },
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed,
                "steps": steps,
                "cfg": cfg,
                "sampler_name": sampler_name,
                "scheduler": scheduler,
                "denoise": denoise,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["11", 0],
            },
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "comfy_img2img", "images": ["8", 0]},
        },
    }


if __name__ == "__main__":
    import json

    wf = comfyui_build_img2img_workflow(
        ckpt_name="IMG_dreamshaper_8.safetensors",
        init_image="example.png",
        positive="a cozy cabin in the woods, golden hour, sharp focus",
        negative="blurry, low quality",
        denoise=0.6,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
