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

Dos modos:
- method="model": upscale con modelo ESRGAN (UpscaleModelLoader +
  ImageUpscaleWithModel). Calidad alta; necesita un modelo en
  models/upscale_models/ (ej. "4x-UltraSharp.pth").
- method="latent": reescalado en espacio de pixel con ImageScaleBy (x2, sin
  modelo ni checkpoint). Upscale rapido y barato.

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


def comfyui_build_upscale_workflow(
    image: str,
    *,
    model_name: str = "4x-UltraSharp.pth",
    method: str = "model",
) -> dict:
    """Construye el dict de un workflow de upscale para una imagen cargada.

    Args:
        image: nombre del archivo de imagen dentro de la carpeta input/ del
            servidor ComfyUI (lo que carga el nodo LoadImage).
        model_name: nombre del modelo de upscale en models/upscale_models/
            (ej. "4x-UltraSharp.pth"). Solo se usa con method="model".
            keyword-only.
        method: "model" (ESRGAN via UpscaleModelLoader + ImageUpscaleWithModel)
            o "latent" (reescalado de pixel x2 con ImageScaleBy, sin modelo).
            keyword-only.

    Returns:
        dict en API format listo para comfyui_submit_workflow.

    Raises:
        ValueError: si method no es "model" ni "latent".
    """
    if method not in ("model", "latent"):
        raise ValueError(
            f"comfyui_build_upscale_workflow: method invalido {method!r}; "
            "usa 'model' o 'latent'."
        )
    load = {
        "10": {"class_type": "LoadImage", "inputs": {"image": image}},
    }
    if method == "model":
        return {
            **load,
            "12": {
                "class_type": "UpscaleModelLoader",
                "inputs": {"model_name": model_name},
            },
            "13": {
                "class_type": "ImageUpscaleWithModel",
                "inputs": {"upscale_model": ["12", 0], "image": ["10", 0]},
            },
            "9": {
                "class_type": "SaveImage",
                "inputs": {
                    "filename_prefix": "comfy_upscale",
                    "images": ["13", 0],
                },
            },
        }
    # method == "latent": reescalado de pixel x2 sin modelo
    return {
        **load,
        "13": {
            "class_type": "ImageScaleBy",
            "inputs": {
                "upscale_method": "lanczos",
                "scale_by": 2.0,
                "image": ["10", 0],
            },
        },
        "9": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "comfy_upscale", "images": ["13", 0]},
        },
    }


if __name__ == "__main__":
    import json

    print(json.dumps(comfyui_build_upscale_workflow("example.png"), indent=2))
    print(json.dumps(comfyui_build_upscale_workflow("example.png", method="latent"), indent=2))
