"""Construye un workflow ComfyUI SDXL base+refiner en API format.

SDXL genera en dos etapas: un checkpoint base produce el latente con la mayor
parte de los pasos y un checkpoint refiner termina los ultimos pasos para pulir
el detalle. Se encadenan dos KSamplerAdvanced compartiendo el numero total de
pasos: el base arranca el ruido y devuelve el latente con ruido sobrante
(return_with_leftover_noise=enable, no decodifica), y el refiner lo recoge
(add_noise=disable) y lo lleva al final.

Cadena de nodos: CheckpointLoaderSimple base + CheckpointLoaderSimple refiner +
EmptyLatentImage + 4 CLIPTextEncode (positivo/negativo por cada CLIP) ->
KSamplerAdvanced base -> KSamplerAdvanced refiner -> VAEDecode -> SaveImage.

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_sdxl_refiner_workflow(
    base_ckpt: str,
    refiner_ckpt: str,
    positive: str,
    negative: str = "",
    *,
    base_steps: int = 20,
    refiner_steps: int = 5,
    cfg: float = 7.0,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
) -> dict:
    """Construye el dict de un workflow SDXL base+refiner (dos KSamplerAdvanced).

    Args:
        base_ckpt: nombre del checkpoint base SDXL tal como lo ve el servidor
            ComfyUI (ej. "sd_xl_base_1.0.safetensors"). En CheckpointLoaderSimple.
        refiner_ckpt: nombre del checkpoint refiner SDXL
            (ej. "sd_xl_refiner_1.0.safetensors").
        positive: prompt positivo (lo que se quiere ver en la imagen). Se usa
            tanto para el CLIP del base como para el del refiner.
        negative: prompt negativo (lo que se quiere evitar). Por defecto "".
        base_steps: pasos que ejecuta el sampler base (del 0 a base_steps).
            keyword-only.
        refiner_steps: pasos que ejecuta el refiner (de base_steps al total).
            El total de pasos es base_steps + refiner_steps. keyword-only.
        cfg: classifier-free guidance scale (compartido). keyword-only.
        seed: semilla de ruido (compartida por ambos samplers). keyword-only.
        width: ancho del latente/imagen en px (SDXL nativo 1024). keyword-only.
        height: alto del latente/imagen en px (SDXL nativo 1024). keyword-only.

    Returns:
        dict en API format listo para comfyui_submit_workflow. Las claves son
        node_ids y cada valor tiene class_type + inputs.
    """
    total_steps = base_steps + refiner_steps
    return {
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": base_ckpt},
        },
        "14": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": refiner_ckpt},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": positive, "clip": ["4", 1]},
        },
        "7": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative, "clip": ["4", 1]},
        },
        "16": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": positive, "clip": ["14", 1]},
        },
        "17": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative, "clip": ["14", 1]},
        },
        "3": {
            "class_type": "KSamplerAdvanced",
            "inputs": {
                "add_noise": "enable",
                "noise_seed": seed,
                "steps": total_steps,
                "cfg": cfg,
                "sampler_name": "euler",
                "scheduler": "normal",
                "start_at_step": 0,
                "end_at_step": base_steps,
                "return_with_leftover_noise": "enable",
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "15": {
            "class_type": "KSamplerAdvanced",
            "inputs": {
                "add_noise": "disable",
                "noise_seed": seed,
                "steps": total_steps,
                "cfg": cfg,
                "sampler_name": "euler",
                "scheduler": "normal",
                "start_at_step": base_steps,
                "end_at_step": 10000,
                "return_with_leftover_noise": "disable",
                "model": ["14", 0],
                "positive": ["16", 0],
                "negative": ["17", 0],
                "latent_image": ["3", 0],
            },
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["15", 0], "vae": ["14", 2]},
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "comfy_sdxl", "images": ["8", 0]},
        },
    }


if __name__ == "__main__":
    import json

    wf = comfyui_build_sdxl_refiner_workflow(
        base_ckpt="sd_xl_base_1.0.safetensors",
        refiner_ckpt="sd_xl_refiner_1.0.safetensors",
        positive="a majestic lion on a cliff at sunset, ultra detailed",
        negative="blurry, low quality",
        base_steps=20,
        refiner_steps=5,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
