fn_registry/python/functions/ml/assemble_animated_sprite.py

"""Ensambla frames PNG RGBA en un sprite sheet horizontal + una animacion en loop.

Funcion impura: lee N frames de disco (los frames ya pixelizados de un walk cycle,
por ejemplo) y escribe DOS entregables:

  1. Un sprite sheet horizontal (1 fila x N columnas) PNG RGBA, con la transparencia
     de cada frame intacta.
  2. Una animacion en bucle (WEBP lossless o GIF animado) que reproduce los frames.

Es la pieza de ensamblado final de cualquier animacion de sprite: convierte una lista
de frames sueltos en algo que se ve animado (la .webp/.gif) y algo que un motor de
juego puede trocear (el sheet). Solo depende de PIL (Pillow), presente en el venv del
registry. No lanza excepciones: cualquier problema se reporta en el campo "error".
"""
from __future__ import annotations

import os


def assemble_animated_sprite(
    frame_paths: list,
    out_dir: str,
    *,
    name: str = "anim",
    fps: int = 8,
    fmt: str = "webp",
    loop: bool = True,
    spritesheet: bool = True,
    pad: int = 0,
) -> dict:
    """Monta un sprite sheet horizontal y una animacion en loop a partir de N frames.

    Carga cada ruta de ``frame_paths`` como RGBA. Los frames que falten o no abran se
    SALTAN (se anota un aviso en ``error``, no se aborta): se anima con los que haya.
    El tamano se normaliza al del primer frame valido; los frames de tamano distinto se
    reescalan con NEAREST a ese tamano (preserva el pixel-art duro, sin interpolacion).

    Args:
        frame_paths: lista de rutas a PNG RGBA, en orden de reproduccion.
        out_dir: directorio de salida; se crea si no existe.
        name: nombre base de los ficheros generados (``<name>_sheet.png`` y
            ``<name>.<ext>``). keyword-only.
        fps: frames por segundo de la animacion; duration_ms = round(1000/max(1,fps)).
            keyword-only.
        fmt: formato de la animacion, "webp" (recomendado) o "gif". keyword-only.
        loop: si True la animacion se repite indefinidamente; si False se reproduce una
            sola vez. keyword-only.
        spritesheet: si True genera tambien el sprite sheet horizontal PNG RGBA.
            keyword-only.
        pad: pixeles de separacion transparente entre columnas del sheet (default 0).
            keyword-only.

    Returns:
        dict con:
        - ok (bool): True si se produjo al menos la animacion con >=1 frame valido.
        - spritesheet_path (str): ruta del PNG del sheet ("" si spritesheet=False o fallo).
        - animation_path (str): ruta de la animacion WEBP/GIF ("" si fallo).
        - n_frames (int): numero de frames validos efectivamente usados.
        - frame_size ([w, h]): tamano del frame normalizado.
        - fmt (str): formato real de la animacion ("webp" o "gif").
        - error (str): avisos y/o mensaje de error; "" si todo fue limpio.
    """
    out = {
        "ok": False,
        "spritesheet_path": "",
        "animation_path": "",
        "n_frames": 0,
        "frame_size": [0, 0],
        "fmt": "",
        "error": "",
    }
    warnings: list = []

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

    if not frame_paths:
        out["error"] = "frame_paths vacio: no hay nada que ensamblar"
        return out

    fmt = str(fmt).lower().strip()
    if fmt not in ("webp", "gif"):
        out["error"] = f"fmt invalido {fmt!r}: usa 'webp' o 'gif'"
        return out
    out["fmt"] = fmt

    # --- Cargar y normalizar frames (saltando los invalidos) ---
    frames: list = []
    target = None  # (w, h) del primer frame valido
    for path in frame_paths:
        if not os.path.isfile(path):
            warnings.append(f"falta: {path}")
            continue
        try:
            with Image.open(path) as src:
                im = src.convert("RGBA")
        except (OSError, ValueError) as exc:
            warnings.append(f"no abre {path}: {exc}")
            continue
        if target is None:
            target = (im.width, im.height)
        elif (im.width, im.height) != target:
            im = im.resize(target, Image.NEAREST)
        frames.append(im)

    if not frames:
        out["error"] = "; ".join(["0 frames validos"] + warnings)
        return out

    w, h = target
    out["frame_size"] = [w, h]
    out["n_frames"] = len(frames)
    n = len(frames)

    try:
        os.makedirs(out_dir, exist_ok=True)
    except OSError as exc:
        out["error"] = "; ".join([f"no se pudo crear out_dir {out_dir!r}: {exc}"] + warnings)
        return out

    # --- Sprite sheet horizontal (1 fila x N columnas), RGBA transparente ---
    if spritesheet:
        pad = max(0, int(pad))
        sheet_w = n * w + (n - 1) * pad if n > 0 else 0
        sheet = Image.new("RGBA", (sheet_w, h), (0, 0, 0, 0))
        for i, im in enumerate(frames):
            x = i * (w + pad)
            # Tercer arg = mascara alpha del propio frame: respeta su transparencia.
            sheet.paste(im, (x, 0), im)
        sheet_path = os.path.join(out_dir, f"{name}_sheet.png")
        try:
            sheet.save(sheet_path, format="PNG")
            out["spritesheet_path"] = sheet_path
        except OSError as exc:
            warnings.append(f"sheet no guardado: {exc}")

    # --- Animacion en loop (WEBP lossless o GIF con alpha binario) ---
    duration_ms = round(1000 / max(1, int(fps)))
    loop_count = 0 if loop else 1  # 0 = infinito
    ext = fmt
    anim_path = os.path.join(out_dir, f"{name}.{ext}")

    try:
        if fmt == "webp":
            frames[0].save(
                anim_path,
                save_all=True,
                append_images=frames[1:],
                duration=duration_ms,
                loop=loop_count,
                format="WEBP",
                lossless=True,   # no ensucia el pixel-art
                disposal=2,      # limpia entre frames -> transparencia correcta
            )
        else:  # gif
            pal_frames = [_rgba_to_p_transparent(im) for im in frames]
            pal_frames[0].save(
                anim_path,
                save_all=True,
                append_images=pal_frames[1:],
                duration=duration_ms,
                loop=loop_count,
                format="GIF",
                transparency=255,  # indice reservado para el pixel transparente
                disposal=2,
            )
        out["animation_path"] = anim_path
        out["ok"] = True
    except (OSError, ValueError) as exc:
        warnings.append(f"animacion no guardada: {exc}")
        out["ok"] = False

    out["error"] = "; ".join(warnings)
    return out


def _rgba_to_p_transparent(im, alpha_threshold: int = 128):
    """Convierte un frame RGBA a modo P reservando el indice 255 como transparente.

    GIF solo soporta 1 bit de alpha: cada pixel es opaco o totalmente transparente.
    Los pixeles con alpha < alpha_threshold se mapean al indice 255 (transparente);
    el resto se cuantiza a 255 colores (indices 0..254).
    """
    from PIL import Image

    alpha = im.getchannel("A")
    # Cuantiza el RGB a 255 colores -> indices 0..254 libres, 255 para transparencia.
    p = im.convert("RGB").quantize(colors=255, method=Image.Quantize.MEDIANCUT)
    # Mascara de los pixeles "transparentes" (alpha por debajo del umbral).
    mask = alpha.point(lambda a: 255 if a < alpha_threshold else 0)
    p.paste(255, (0, 0), mask)
    return p


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

    from PIL import Image as _Image, ImageDraw as _ImageDraw

    # --- Genera 4 frames de prueba: un cuadrado de color que se mueve de izquierda a
    #     derecha sobre un lienzo RGBA transparente de 32x32. ---
    tmp = tempfile.mkdtemp(prefix="assemble_sprite_demo_")
    demo_frames: list = []
    box = 10
    for i in range(4):
        frame = _Image.new("RGBA", (32, 32), (0, 0, 0, 0))  # fondo transparente
        d = _ImageDraw.Draw(frame)
        x0 = 1 + i * 6  # se desplaza hacia la derecha cada frame
        d.rectangle([x0, 11, x0 + box, 11 + box], fill=(40, 180, 230, 255))
        fpath = os.path.join(tmp, f"frame_{i:02d}.png")
        frame.save(fpath)
        demo_frames.append(fpath)

    result = assemble_animated_sprite(
        demo_frames, tmp, name="walk_demo", fps=8, fmt="webp"
    )
    print(json.dumps(result, indent=2))