"""Monta un grid / contact-sheet PIL de N imagenes para comparacion visual.

Funcion impura: lee N imagenes de disco y escribe un PNG de salida. Usa PIL
(Pillow), presente en el venv del registry.

El compañero natural de comfyui_batch_generate: ese encola N variantes de un
workflow (una por seed) pero no junta los resultados. Esta funcion toma las N
imagenes ya descargadas (p.ej. con comfyui_fetch_output_image) y las dispone en
una rejilla regular para compararlas de un vistazo. Cada celda conserva el aspect
ratio (thumbnail centrado sobre fondo oscuro). Opcionalmente rotula cada celda.
"""
import math
import os


def comfyui_build_grid(
    image_paths: list,
    *,
    cols: int | None = None,
    cell: int = 512,
    out_path: str | None = None,
    labels: list | None = None,
) -> dict:
    """Compone una rejilla de imagenes y la guarda como PNG.

    Args:
        image_paths: lista de rutas a las imagenes (PNG/JPG/...) a montar, en
            orden de lectura (izquierda->derecha, arriba->abajo).
        cols: numero de columnas; si None se usa ceil(sqrt(N)) para una rejilla
            casi cuadrada. keyword-only.
        cell: lado en pixeles de cada celda cuadrada; cada imagen se reduce para
            caber dentro conservando su proporcion. keyword-only.
        out_path: ruta del PNG de salida; si None se escribe "comfy_grid.png" en
            el directorio de la primera imagen. keyword-only.
        labels: rotulos opcionales, uno por imagen (mismo orden); si se pasan, se
            reserva una franja bajo cada celda y se dibuja el texto. keyword-only.

    Returns:
        dict con:
        - ok (bool): True si el grid se monto y guardo.
        - out_path (str): ruta del PNG generado.
        - rows (int): filas de la rejilla.
        - cols (int): columnas de la rejilla.
        - error (str): mensaje de error; cadena vacia si todo OK.
    """
    out = {"ok": False, "out_path": "", "rows": 0, "cols": 0, "error": ""}

    if not image_paths:
        out["error"] = "image_paths vacio: nada que montar"
        return out

    try:
        from PIL import Image, ImageDraw
    except ImportError:
        out["error"] = "PIL (Pillow) no esta instalado en este interprete"
        return out

    missing = [p for p in image_paths if not os.path.isfile(p)]
    if missing:
        out["error"] = f"no existen {len(missing)} rutas: {missing[:5]}"
        return out

    n = len(image_paths)
    cols = int(cols) if cols and cols > 0 else max(1, math.ceil(math.sqrt(n)))
    rows = math.ceil(n / cols)
    cell = max(16, int(cell))
    label_h = 22 if labels else 0
    bg = (24, 24, 28)
    fg = (232, 232, 236)

    canvas = Image.new("RGB", (cols * cell, rows * (cell + label_h)), bg)
    draw = ImageDraw.Draw(canvas) if labels else None

    try:
        for i, path in enumerate(image_paths):
            with Image.open(path) as src:
                im = src.convert("RGB")
                im.thumbnail((cell, cell))
                r, c = divmod(i, cols)
                x = c * cell + (cell - im.width) // 2
                y = r * (cell + label_h) + (cell - im.height) // 2
                canvas.paste(im, (x, y))
            if draw is not None and i < len(labels):
                tx = c * cell + 4
                ty = r * (cell + label_h) + cell + 3
                draw.text((tx, ty), str(labels[i]), fill=fg)
    except OSError as exc:
        out["error"] = f"no se pudo leer/decodificar una imagen: {exc}"
        return out

    if out_path is None:
        out_path = os.path.join(os.path.dirname(os.path.abspath(image_paths[0])),
                                "comfy_grid.png")
    try:
        os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
        canvas.save(out_path)
    except OSError as exc:
        out["error"] = f"no se pudo escribir {out_path!r}: {exc}"
        return out

    out.update(ok=True, out_path=out_path, rows=rows, cols=cols)
    return out


if __name__ == "__main__":
    import json
    import sys

    paths = sys.argv[1:]
    if not paths:
        print("uso: comfyui_build_grid.py <img1> <img2> ...", file=sys.stderr)
        sys.exit(2)
    res = comfyui_build_grid(paths, out_path="/tmp/comfy_grid.png")
    print(json.dumps(res, indent=2))