"""Construye el workflow ComfyUI de un sprite MULTI-DIRECCIONAL (vistas 3D consistentes).

A partir de UNA imagen frontal de un personaje (fondo limpio), construye el dict (API
format) de un workflow que genera N vistas direccionales del MISMO personaje rotando en
3D: las 8 direcciones N/NE/E/SE/S/SW/W/NW de un shooter 8-way / juego top-down /
isométrico, o las 4 direcciones N/E/S/W de un RPG clásico. El valor de este builder es la
CONSISTENCIA ROTACIONAL: es el mismo personaje girado en 3D, no re-dibujado por dirección.

Es el puente 2.5D que faltaba en el catálogo `gamedev-2d`. Los builders 2D de ese grupo
(`enemy_creature`, `topdown_sprite`, `sprite_sheet`) producen UNA vista; `sprite_sheet`
re-poza con OpenPose 2D (re-dibuja la silueta para cada pose, así que la identidad del
personaje deriva entre ángulos). Aquí, en cambio, un modelo de difusión 3D (SV3D turntable
u Stable Zero123 órbita) rota la figura alrededor de su eje vertical, de modo que el casco,
el arma y la paleta son los mismos en cada dirección. Es la diferencia entre "dibujar al
goblin 8 veces" (inconsistente) y "fotografiar al goblin desde 8 ángulos" (consistente).

Dos modelos, ambos confirmados en /object_info del server 8GB:

  - SV3D (model="sv3d", default): el nodo `SV3D_Conditioning` produce una ÓRBITA de
    `video_frames` vistas equiespaciadas en 360° en una sola pasada (reutiliza la
    maquinaria img2vid de Stable Video Diffusion: `ImageOnlyCheckpointLoader` →
    `SV3D_Conditioning` → `VideoLinearCFGGuidance` → `KSampler` → `VAEDecode` → `SaveImage`,
    los N frames). Mejor consistencia angular; checkpoint `3D_sv3d_p.safetensors` (~9.4 GB en
    disco, ~2-4 GB en VRAM porque es un modelo de vídeo en lowvram). Frame 0 = la vista de
    entrada; el frame i corresponde al azimuth i·360/N.
  - Stable Zero123 (model="zero123"): el nodo `StableZero123_Conditioning_Batched` genera
    un BATCH de N vistas a azimuths equiespaciados (control directo del ángulo por vista).
    Útil como fallback si SV3D no cabe en VRAM: cada vista es una imagen independiente, así
    que se puede bajar `directions` o `size` con más holgura. Checkpoint
    `3D_stable_zero123.ckpt`. width/height nativos 256.

Función PURA: sin red, sin I/O. Solo construye el dict del grafo (los class_types y sus
inputs verificados contra /object_info). El coste GPU está al enviar con
`comfyui_submit_workflow`. La imagen de entrada debe existir ya en el `input/` del servidor
(subir antes con POST /upload/image). Hermana impura: `comfyui_generate_views_from_image`
(orquesta /object_info + submit + wait + fetch para reconstrucción 3D multi-vista, 4
cardinales); este builder es el equivalente PURO orientado a sprite direccional gamedev,
igual que `comfyui_build_img2vid_workflow` es el builder puro frente al orquestador.
"""
from __future__ import annotations

# Modelos soportados -> checkpoint por defecto que el nodo carga vía ImageOnlyCheckpointLoader.
_MODEL_CKPT = {
    "sv3d": "3D_sv3d_p.safetensors",
    "zero123": "3D_stable_zero123.ckpt",
}

# Resolución nativa por modelo cuando el caller no fija `size`. SV3D entrena a 576²;
# Stable Zero123 a 256². Usar la nativa evita degradación y picos de VRAM innecesarios.
_MODEL_NATIVE_SIZE = {"sv3d": 576, "zero123": 256}

# CFG por defecto por modelo cuando el caller no lo fija. SV3D usa guía de vídeo baja
# interpolada (VideoLinearCFGGuidance de min_cfg a cfg); Zero123 usa una CFG algo mayor.
_MODEL_DEFAULT_CFG = {"sv3d": 2.5, "zero123": 4.0}

# Nombres de dirección por nº de direcciones, en el orden del orbit empezando por la vista
# frontal (frame 0 = personaje de frente, mirando a cámara). Convención gamedev top-down:
# "front" = el personaje mirando hacia ABAJO/al jugador (south). Cada paso avanza el azimuth
# 360/N grados. El orden es solo una etiqueta de conveniencia para el caller al ensamblar el
# grid direccional; el grafo no depende de él.
_DIRECTION_NAMES = {
    4: ["S", "E", "N", "W"],
    8: ["S", "SE", "E", "NE", "N", "NW", "W", "SW"],
}


def directional_sprite_view_order(directions: int) -> list:
    """Lista ordenada de nombres de dirección que produce el orbit, frame 0..N-1.

    El frame i del workflow (i-ésima imagen del SaveImage, en orden de azimuth creciente
    desde la vista frontal) corresponde a la dirección devuelta en la posición i. Para
    4 u 8 direcciones devuelve los nombres cardinales canónicos; para cualquier otro N
    devuelve etiquetas genéricas por azimuth (``az0``, ``az45``, ...). Pura.

    Args:
        directions: número de direcciones del orbit (>= 1).

    Returns:
        lista de `directions` nombres de dirección, alineada por índice con los frames.
    """
    directions = int(directions)
    if directions in _DIRECTION_NAMES:
        return list(_DIRECTION_NAMES[directions])
    return [f"az{round(i * 360 / directions)}" for i in range(directions)]


def comfyui_build_directional_sprite_workflow(
    input_image: str,
    *,
    directions: int = 8,
    model: str = "sv3d",
    elevation: float = 0.0,
    size: int | None = None,
    orbit_frames: int | None = None,
    seed: int = 0,
    ckpt: str | None = None,
    steps: int = 20,
    cfg: float | None = None,
    min_cfg: float = 1.0,
    sampler_name: str = "euler",
    scheduler: str = "karras",
    filename_prefix: str = "directional_sprite",
) -> dict:
    """Construye el dict (API format) del workflow de sprite multi-direccional.

    Recibe el sprite frontal de UN personaje (imagen en el `input/` del servidor, fondo
    limpio) y construye un workflow que produce N vistas direccionales CONSISTENTES (el
    mismo personaje rotado en 3D), una por dirección de movimiento del juego.

    Args:
        input_image: nombre del archivo de imagen frontal del personaje dentro del `input/`
            del servidor ComfyUI (lo carga el nodo LoadImage). Debe existir ya (subir con
            POST /upload/image o copiar a ~/ComfyUI/input/). Idealmente con fondo limpio y el
            personaje centrado y encuadrado entero, mirando a cámara. No puede estar vacío.
        directions: número de direcciones a generar. 8 (default) = 8-way N/NE/E/SE/S/SW/W/NW
            (shooter 8-direcciones, top-down/iso); 4 = N/E/S/W (RPG clásico). Cualquier N>=1
            es válido (vistas equiespaciadas a 360/N grados); 4 y 8 reciben nombres cardinales
            (ver `directional_sprite_view_order`). keyword-only.
        model: "sv3d" (default, órbita turntable, mejor consistencia angular) o "zero123"
            (batch de vistas con control de azimuth, fallback de menor VRAM). keyword-only.
        elevation: elevación de cámara en grados. 0 = órbita en el ecuador (vista lateral
            pura). Un juego top-down/isométrico suele querer algo de PICADO (mirar al
            personaje desde arriba): subir a ~15-30. keyword-only.
        size: lado en px de cada vista. None (default) usa la resolución nativa del modelo
            (576 para sv3d, 256 para zero123). Bajarlo reduce VRAM a costa de detalle.
            keyword-only.
        orbit_frames: SOLO sv3d. Número de frames del orbit que sintetiza SV3D_Conditioning.
            None (default) = `directions` (una imagen por dirección, salida directa). Subirlo
            (p. ej. 21, el nativo de SV3D) da una órbita más densa y consistente de la que el
            caller submuestrea las `directions` vistas (frame más cercano por azimuth); a más
            frames, más VRAM. Ignorado por zero123 (su batch es exactamente `directions`).
            keyword-only.
        seed: semilla del KSampler (0 = determinista; cambiar para variar la órbita).
            keyword-only.
        ckpt: nombre del checkpoint en el servidor. None (default) usa el del modelo
            (`3D_sv3d_p.safetensors` / `3D_stable_zero123.ckpt`). keyword-only.
        steps: pasos de sampling del KSampler. keyword-only.
        cfg: guidance scale. None (default) usa el del modelo (2.5 sv3d / 4.0 zero123).
            keyword-only.
        min_cfg: SOLO sv3d. CFG del primer frame para VideoLinearCFGGuidance (interpola de
            min_cfg a cfg a lo largo del orbit). keyword-only.
        sampler_name: algoritmo del KSampler. Default "euler". keyword-only.
        scheduler: scheduler del KSampler. Default "karras". keyword-only.
        filename_prefix: prefijo del archivo de salida del SaveImage. keyword-only.

    Returns:
        dict en API format listo para `comfyui_submit_workflow`. Las claves son node_ids
        (string) y cada valor tiene class_type + inputs. SV3D devuelve 6 nodos
        (ImageOnlyCheckpointLoader, LoadImage, SV3D_Conditioning, VideoLinearCFGGuidance,
        KSampler, VAEDecode, SaveImage → los N frames del orbit). Zero123 devuelve
        ImageOnlyCheckpointLoader + LoadImage + StableZero123_Conditioning_Batched + KSampler
        + VAEDecode + SaveImage (un batch de `directions` vistas). El frame i (i-ésima imagen
        del SaveImage, azimuth creciente desde la frontal) = dirección i de
        `directional_sprite_view_order(directions)`.

    Raises:
        ValueError: si input_image está vacío, model no es "sv3d"/"zero123", o directions < 1.
    """
    if not input_image or not input_image.strip():
        raise ValueError(
            "comfyui_build_directional_sprite_workflow: 'input_image' no puede estar vacío "
            "(el modelo 3D necesita la vista frontal en input/)."
        )
    if model not in _MODEL_CKPT:
        raise ValueError(
            "comfyui_build_directional_sprite_workflow: model "
            f"'{model}' inválido; usa uno de {sorted(_MODEL_CKPT)}."
        )
    directions = int(directions)
    if directions < 1:
        raise ValueError(
            "comfyui_build_directional_sprite_workflow: 'directions' debe ser >= 1 "
            f"(recibido {directions}); 4 u 8 son los valores canónicos de gamedev."
        )

    input_image = input_image.strip()
    ckpt_name = ckpt or _MODEL_CKPT[model]
    px = int(size) if size is not None else _MODEL_NATIVE_SIZE[model]
    guidance = float(cfg) if cfg is not None else _MODEL_DEFAULT_CFG[model]

    if model == "sv3d":
        return _build_sv3d(
            input_image, ckpt_name, directions, elevation, px, orbit_frames, seed,
            steps, guidance, min_cfg, sampler_name, scheduler, filename_prefix,
        )
    return _build_zero123(
        input_image, ckpt_name, directions, elevation, px, seed,
        steps, guidance, sampler_name, scheduler, filename_prefix,
    )


def _build_sv3d(input_image, ckpt_name, directions, elevation, size, orbit_frames,
                seed, steps, cfg, min_cfg, sampler_name, scheduler, filename_prefix) -> dict:
    """Workflow SV3D: 1 imagen frontal -> orbit de N vistas en una pasada.

    SV3D reutiliza la cadena img2vid de Stable Video Diffusion: el checkpoint 3D_sv3d_p se
    carga con ImageOnlyCheckpointLoader (MODEL + CLIP_VISION + VAE), SV3D_Conditioning
    produce el conditioning del orbit y la latente, se muestrea como vídeo con guía CFG
    lineal y el SaveImage emite los `video_frames` frames del turntable.
    """
    video_frames = int(orbit_frames) if orbit_frames else directions
    return {
        "1": {"class_type": "LoadImage", "inputs": {"image": input_image}},
        "2": {"class_type": "ImageOnlyCheckpointLoader", "inputs": {"ckpt_name": ckpt_name}},
        "3": {
            "class_type": "SV3D_Conditioning",
            "inputs": {
                "clip_vision": ["2", 1],
                "init_image": ["1", 0],
                "vae": ["2", 2],
                "width": size,
                "height": size,
                "video_frames": video_frames,
                "elevation": float(elevation),
            },
        },
        "4": {
            "class_type": "VideoLinearCFGGuidance",
            "inputs": {"model": ["2", 0], "min_cfg": float(min_cfg)},
        },
        "5": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(seed),
                "steps": int(steps),
                "cfg": float(cfg),
                "sampler_name": sampler_name,
                "scheduler": scheduler,
                "denoise": 1.0,
                "model": ["4", 0],
                "positive": ["3", 0],
                "negative": ["3", 1],
                "latent_image": ["3", 2],
            },
        },
        "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["2", 2]}},
        "7": {
            "class_type": "SaveImage",
            "inputs": {"images": ["6", 0], "filename_prefix": filename_prefix},
        },
    }


def _build_zero123(input_image, ckpt_name, directions, elevation, size, seed,
                   steps, cfg, sampler_name, scheduler, filename_prefix) -> dict:
    """Workflow Stable Zero123: batch de N vistas a azimuths equiespaciados.

    StableZero123_Conditioning_Batched genera `directions` vistas en una pasada, empezando
    en azimuth 0 (la frontal) y avanzando 360/N grados por vista. El KSampler las muestrea
    como un batch y el SaveImage las emite en orden de azimuth creciente.
    """
    increment = 360.0 / directions
    return {
        "1": {"class_type": "LoadImage", "inputs": {"image": input_image}},
        "2": {"class_type": "ImageOnlyCheckpointLoader", "inputs": {"ckpt_name": ckpt_name}},
        "3": {
            "class_type": "StableZero123_Conditioning_Batched",
            "inputs": {
                "clip_vision": ["2", 1],
                "init_image": ["1", 0],
                "vae": ["2", 2],
                "width": size,
                "height": size,
                "batch_size": directions,
                "elevation": float(elevation),
                "azimuth": 0.0,
                "elevation_batch_increment": 0.0,
                "azimuth_batch_increment": increment,
            },
        },
        "4": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(seed),
                "steps": int(steps),
                "cfg": float(cfg),
                "sampler_name": sampler_name,
                "scheduler": scheduler,
                "denoise": 1.0,
                "model": ["2", 0],
                "positive": ["3", 0],
                "negative": ["3", 1],
                "latent_image": ["3", 2],
            },
        },
        "5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0], "vae": ["2", 2]}},
        "6": {
            "class_type": "SaveImage",
            "inputs": {"images": ["5", 0], "filename_prefix": filename_prefix},
        },
    }


if __name__ == "__main__":
    import json

    wf = comfyui_build_directional_sprite_workflow(
        "goblin_front.png",
        directions=8,
        model="sv3d",
        elevation=15.0,
        seed=7,
    )
    print(json.dumps({
        "nodes": list(wf),
        "classes": sorted({n["class_type"] for n in wf.values()}),
        "video_frames": wf["3"]["inputs"]["video_frames"],
        "size": wf["3"]["inputs"]["width"],
        "view_order": directional_sprite_view_order(8),
    }, indent=2))
