"""Construye un workflow ComfyUI de "hires fix" de 2 pasadas en API format.

El hires fix clasico genera una imagen pequena nitida y luego la amplia
*re-difundiendola* (no solo escalando pixeles), de modo que el modelo anade
detalle coherente a la resolucion alta. Este builder lo implementa con
UltimateSDUpscale (custom node), que hace la segunda pasada por TILES con un
modelo de upscale (ESRGAN/Remacri) + un sampler con `denoise` parcial:

    Pasada 1 (base):   CheckpointLoaderSimple -> CLIPTextEncode(+/-) +
                       EmptyLatentImage(first_pass) -> KSampler -> VAEDecode
    Pasada 2 (detalle): UpscaleModelLoader(Remacri) +
                        UltimateSDUpscale(image, model, +/-, vae, upscale_model,
                        upscale_by, denoise<1, tiled) -> SaveImage

UltimateSDUpscale ES la segunda pasada de muestreo: re-samplea cada tile con el
checkpoint (de ahi que reciba `model`, `positive`, `negative`, `vae`), por eso el
grafo tiene UN KSampler explicito (la base) + el UltimateSDUpscale (el detalle).
Distinto de `comfyui_build_upscale_workflow`, que es ESRGAN puro SIN re-difusion.

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

from __future__ import annotations


def comfyui_build_hires_fix_workflow(
    ckpt_name: str,
    positive: str,
    negative: str = "",
    *,
    first_pass: tuple[int, int] = (768, 768),
    upscale_by: float = 1.5,
    denoise: float = 0.4,
    steps: int = 20,
    cfg: float = 7.0,
    seed: int = 0,
    upscale_model: str = "4x_foolhardy_Remacri.pth",
    sampler_name: str = "euler",
    scheduler: str = "normal",
    tile_width: int = 512,
    tile_height: int = 512,
    filename_prefix: str = "hires",
) -> dict:
    """Construye el dict de un workflow hires-fix (base + UltimateSDUpscale).

    Args:
        ckpt_name: checkpoint tal como lo ve el servidor (CheckpointLoaderSimple).
        positive: prompt positivo (se usa en la base y en la re-difusion tiled).
        negative: prompt negativo. Por defecto "".
        first_pass: (ancho, alto) en px de la pasada base (latente pequeno y
            rapido). Por defecto (768, 768). keyword-only.
        upscale_by: factor de ampliacion de UltimateSDUpscale sobre la imagen base
            (1.5 -> 768 pasa a 1152). keyword-only.
        denoise: fuerza de re-difusion de la segunda pasada (0.4 por defecto).
            <1 para conservar la composicion base y solo anadir detalle; 1.0 la
            re-generaria entera. keyword-only.
        steps: pasos de sampling (ambas pasadas). keyword-only.
        cfg: classifier-free guidance (ambas pasadas). keyword-only.
        seed: semilla de la pasada base (UltimateSDUpscale usa la misma).
            keyword-only.
        upscale_model: modelo de upscale en models/upscale_models/ que usa
            UltimateSDUpscale para escalar antes de re-difundir (ej.
            "4x_foolhardy_Remacri.pth"). keyword-only.
        sampler_name: sampler (ambas pasadas). keyword-only.
        scheduler: scheduler (ambas pasadas). keyword-only.
        tile_width: ancho de tile de UltimateSDUpscale (px). Tiles mas pequenos =
            menos VRAM, mas costuras. keyword-only.
        tile_height: alto de tile de UltimateSDUpscale (px). keyword-only.
        filename_prefix: prefijo del PNG final que escribe SaveImage. keyword-only.

    Returns:
        dict en API format listo para `comfyui_submit_workflow`. node_ids:
        "4" CheckpointLoaderSimple, "5" EmptyLatentImage, "6"/"7" CLIPTextEncode,
        "3" KSampler (base), "8" VAEDecode, "11" UpscaleModelLoader,
        "12" UltimateSDUpscale, "9" SaveImage.
    """
    w, h = first_pass
    return {
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": ckpt_name},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": w, "height": h, "batch_size": 1},
        },
        "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": 1.0,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0],
            },
        },
        "8": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "11": {
            "class_type": "UpscaleModelLoader",
            "inputs": {"model_name": upscale_model},
        },
        "12": {
            "class_type": "UltimateSDUpscale",
            "inputs": {
                "image": ["8", 0],
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "vae": ["4", 2],
                "upscale_model": ["11", 0],
                "upscale_by": upscale_by,
                "seed": seed,
                "steps": steps,
                "cfg": cfg,
                "sampler_name": sampler_name,
                "scheduler": scheduler,
                "denoise": denoise,
                "mode_type": "Linear",
                "tile_width": tile_width,
                "tile_height": tile_height,
                "mask_blur": 8,
                "tile_padding": 32,
                "seam_fix_mode": "None",
                "seam_fix_denoise": 1.0,
                "seam_fix_width": 64,
                "seam_fix_mask_blur": 8,
                "seam_fix_padding": 16,
                "force_uniform_tiles": True,
                "tiled_decode": False,
                "batch_size": 1,
            },
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": filename_prefix, "images": ["12", 0]},
        },
    }


if __name__ == "__main__":
    import json

    wf = comfyui_build_hires_fix_workflow(
        ckpt_name="IMG_dreamshaper_8.safetensors",
        positive="a fox in a forest, intricate detail, sharp focus",
        negative="blurry, low quality",
        first_pass=(768, 768),
        upscale_by=1.5,
        denoise=0.4,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
