"""comfyui_walk_cycle_oneshot — personaje andando -> sprite pixel-art animado (GIF/WEBP).

Pipeline one-shot (issue 0087) que materializa el caso 1 de la investigacion de
animacion de sprites (report 0217): un ciclo de caminar de personaje se anima
**frame-by-frame dirigido por pose**, NO con modelos de video. Por cada fase del
paso se genera el MISMO personaje cambiando solo la pose (ControlNet OpenPose +
seed fija + IPAdapter opcional), se recorta el fondo a alpha (Rembg) y se pixeliza
a un grid duro de `size` x `size` RGBA. Los N frames se ensamblan en un sprite
sheet horizontal + una animacion en loop (WEBP/GIF).

Dos metodos para las poses, seleccionables con `pose_method`:
  - "openpose" (preferido): se dibujan N esqueletos OpenPose del walk cycle
    (render_openpose_walk_skeletons) y se alimentan al ControlNet OpenPose, un
    frame por pose. El timing del paso es exactamente el dibujado.
  - "prompt" (fallback): sin esqueletos, la fase del paso se describe en el prompt
    ("left leg forward", "passing pose", ...) con seed fija. A 32x32 el detalle de
    pose se simplifica, asi que un walk de 4 frames reconocible basta. Se usa
    cuando el ControlNet no da poses utilizables o el caller lo pide.
  - "auto" (default): intenta "openpose"; si el render de esqueletos falla, cae a
    "prompt" y lo refleja en pose_method_used.

La identidad consistente entre frames se ancla con **seed fija** (mismo personaje,
misma semilla) + prompt base fijo, y opcionalmente con `ref_image` (IPAdapter-FaceID
en comfyui_build_sprite_sheet_workflow).

Compone funciones del registry, no reescribe su logica:
    render_openpose_walk_skeletons_py_ml    (esqueletos OpenPose del walk cycle)
    comfyui_build_sprite_sheet_workflow_py_ml(1 frame: identidad + pose + alpha)
    comfyui_submit_workflow_py_ml            (POST /prompt)
    comfyui_wait_result_py_ml                (poll /history)
    comfyui_fetch_output_image_py_ml         (GET /view -> disco)
    comfyui_pixelize_sprite_png_py_ml        (PNG alta-res -> NxN RGBA pixel-art)
    assemble_animated_sprite_py_ml           (frames -> sprite sheet + WEBP/GIF loop)

Pipeline impuro: red (HTTP a ComfyUI), GPU (generacion), escritura en disco.
No-throw: cualquier fallo se captura y viaja en el dict de estado (campo error).
Si un frame concreto falla se salta y la animacion se monta con los que haya.
"""

from __future__ import annotations

import os
import sys

# Importa las funciones del registry (mismo arbol python/functions).
_FUNCTIONS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _FUNCTIONS_ROOT not in sys.path:
    sys.path.insert(0, _FUNCTIONS_ROOT)

from ml.assemble_animated_sprite import assemble_animated_sprite
from ml.comfyui_build_sprite_sheet_workflow import comfyui_build_sprite_sheet_workflow
from ml.comfyui_fetch_output_image import comfyui_fetch_output_image
from ml.comfyui_pixelize_sprite_png import comfyui_pixelize_sprite_png
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_wait_result import comfyui_wait_result
from ml.render_openpose_walk_skeletons import render_openpose_walk_skeletons

# Descriptores de fase del paso para el modo "prompt" (y como apoyo). El ciclo
# canonico de 4 fases de un walk lateral: contacto pierna izquierda, paso (piernas
# juntas, cuerpo arriba), contacto pierna derecha, paso. Para N != 4 se reparten
# ciclicamente sobre estas cuatro.
_WALK_PHASES = [
    "walking, left leg forward, mid stride, dynamic walk pose",
    "walking, legs together passing position, standing tall",
    "walking, right leg forward, mid stride, dynamic walk pose",
    "walking, legs together passing position, standing tall",
]


def _phase_for(index: int, total: int) -> str:
    """Devuelve el descriptor de fase del paso para el frame `index` de `total`.

    Mapea el frame al ciclo de 4 fases base de forma uniforme, de modo que el
    primer y el ultimo frame cierren un loop continuo (el frame `total` volveria a
    la fase del frame 0).
    """
    pos = (index / max(1, total)) * len(_WALK_PHASES)
    return _WALK_PHASES[int(pos) % len(_WALK_PHASES)]


def comfyui_walk_cycle_oneshot(
    character: str,
    *,
    frames: int = 4,
    size: int = 32,
    colors: int = 16,
    fps: int = 8,
    checkpoint: str = "IMG_dreamshaper_8.safetensors",
    ref_image: str | None = None,
    server: str = "127.0.0.1:8188",
    dest_dir: str = "~/ComfyUI/output",
    seed: int = 0,
    pose_method: str = "auto",
    comfy_input_dir: str = "~/ComfyUI/input",
    controlnet_strength: float = 0.7,
    controlnet_end: float = 0.8,
    facing: str = "right",
    engine: str = "pixeloe",
    palette=None,
    fmt: str = "webp",
    negative: str = "blurry, lowres, extra limbs, deformed, multiple characters",
    width: int = 512,
    height: int = 768,
    steps: int = 24,
    cfg: float = 7.0,
    comfy_python: str | None = None,
    wait_timeout: float = 300.0,
    filename_prefix: str = "walk_cycle",
    keep_frames: bool = True,
    **gen_kwargs,
) -> dict:
    """Genera una animacion de personaje andando en pixel-art, end-to-end.

    Args:
        character: prompt del personaje ("pixel art knight, full body, side view").
            No puede estar vacio. La identidad se mantiene entre frames con seed fija.
        frames: numero de frames del ciclo (>=2, recomendado 4-8). keyword-only.
        size: lado del grid pixel-art final por frame en pixeles (32 sprites
            pequenos, 64 personajes con mas detalle). keyword-only.
        colors: numero de colores de la paleta libre por frame (cuantizacion
            MEDIANCUT) cuando palette es None. keyword-only.
        fps: cadencia de la animacion en frames por segundo. keyword-only.
        checkpoint: checkpoint SD1.5 (ControlNet OpenPose + IPAdapter-FaceID solo
            instalados en SD1.5; default 'IMG_dreamshaper_8.safetensors'). keyword-only.
        ref_image: imagen de cara de referencia en el input/ del servidor para
            IPAdapter-FaceID (segunda ancla de identidad). None usa solo seed +
            prompt. keyword-only.
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        dest_dir: directorio donde guardar frames + sprite sheet + animacion.
            keyword-only.
        seed: semilla FIJA del KSampler para todos los frames (identidad estable).
            keyword-only.
        pose_method: "openpose" (esqueletos OpenPose -> ControlNet, control exacto),
            "prompt" (fase del paso descrita en el prompt, sin esqueletos), o "auto"
            (intenta openpose, cae a prompt si el render falla). keyword-only.
        comfy_input_dir: directorio input/ del servidor ComfyUI donde se escriben
            los esqueletos para que el ControlNet los lea (server local).
            keyword-only.
        controlnet_strength: fuerza del ControlNet OpenPose (0.7 da buen control de
            la pose sin aplastar el estilo). keyword-only.
        controlnet_end: fraccion de pasos en que el OpenPose deja de aplicarse
            (end<1.0 libera los ultimos pasos para pelo/ropa). keyword-only.
        facing: direccion a la que mira el personaje en los esqueletos ("right" o
            "left"). keyword-only.
        engine: motor de downscale del pixelizado ("pixeloe" contrast-aware o
            "nearest"). keyword-only.
        palette: None (paleta libre a `colors`), nombre builtin ("pico-8", "nes",
            "game-boy") o lista de hex. keyword-only.
        fmt: formato de la animacion ("webp" recomendado para alpha, o "gif").
            keyword-only.
        negative: prompt negativo de generacion. keyword-only.
        width, height: resolucion de generacion por frame (512x768 vertical encuadra
            cuerpo entero). keyword-only.
        steps, cfg: parametros del KSampler. keyword-only.
        comfy_python: interprete con pixeloe para el pixelizado (None autodetecta
            ~/ComfyUI/.venv). keyword-only.
        wait_timeout: segundos maximos esperando cada frame al server. keyword-only.
        filename_prefix: prefijo de los archivos de salida. keyword-only.
        keep_frames: si True conserva los PNG de cada frame pixelizado; si False los
            borra tras montar la animacion. keyword-only.
        **gen_kwargs: params extra para comfyui_build_sprite_sheet_workflow
            (sampler_name, scheduler, char_lora, lora_strength, weight, ...).

    Returns:
        dict con:
        - ok (bool): True si se produjo la animacion con >=1 frame.
        - frames (list[str]): rutas de los PNG pixelizados (size x size RGBA).
        - spritesheet_path (str): ruta del sprite sheet horizontal.
        - animation_path (str): ruta de la animacion WEBP/GIF en loop.
        - size (int): lado real de cada frame pixelizado.
        - n_frames (int): numero de frames producidos (puede ser < `frames` si
          alguno fallo y se salto).
        - seed (int): semilla usada.
        - pose_method_used (str): "openpose" o "prompt" (refleja el fallback).
        - skipped (list[int]): indices de frames que fallaron y se saltaron.
        - error (str): mensaje de error; vacio si todo OK.
    """
    out = {
        "ok": False, "frames": [], "spritesheet_path": "", "animation_path": "",
        "size": int(size), "n_frames": 0, "seed": int(seed),
        "pose_method_used": "", "skipped": [], "error": "",
    }

    if not character or not character.strip():
        out["error"] = "character vacio"
        return out
    if int(frames) < 2:
        out["error"] = f"frames debe ser >= 2, recibido {frames!r}"
        return out
    if int(size) < 1:
        out["error"] = f"size debe ser >= 1, recibido {size!r}"
        return out
    if pose_method not in ("auto", "openpose", "prompt"):
        out["error"] = f"pose_method invalido: {pose_method!r}"
        return out

    n = int(frames)
    dest = os.path.expanduser(dest_dir)
    frames_dir = os.path.join(dest, f"{filename_prefix}_frames")
    try:
        os.makedirs(frames_dir, exist_ok=True)
    except OSError as exc:
        out["error"] = f"no se pudo crear dest_dir {frames_dir!r}: {exc}"
        return out

    # --- Fase 0: resolver el metodo de poses + esqueletos OpenPose si aplica. ---
    skeleton_names: list[str | None] = [None] * n
    method = "prompt" if pose_method == "prompt" else "openpose"
    if method == "openpose":
        input_dir = os.path.expanduser(comfy_input_dir)
        sk = render_openpose_walk_skeletons(
            input_dir, frames=n, width=int(width), height=int(height),
            facing=facing, filename_prefix=f"{filename_prefix}_pose",
        )
        if sk.get("ok") and sk.get("skeleton_paths"):
            paths = sk["skeleton_paths"]
            # El builder espera el basename relativo al input/ del servidor.
            skeleton_names = [os.path.basename(p) for p in paths][:n]
            # Si se generaron menos esqueletos que frames, rellena con None.
            while len(skeleton_names) < n:
                skeleton_names.append(None)
        else:
            # Fallback limpio: no hay esqueletos -> modo prompt.
            method = "prompt"
            out["error"] = (
                f"render de esqueletos OpenPose fallo ({sk.get('error')}); "
                f"fallback a pose por prompt"
            )
            if pose_method == "openpose":
                # El caller forzo openpose y no hay esqueletos: aun asi seguimos en
                # prompt para no abortar, pero queda anotado en el error.
                pass
    out["pose_method_used"] = method

    # --- Fase 1..N: generar + pixelizar cada frame del ciclo. ---
    pixel_frames: list[str] = []
    for i in range(n):
        pose_skeleton = skeleton_names[i] if method == "openpose" else None
        # En modo openpose la pose la fija el esqueleto: no se mete la fase en el
        # prompt para no competir con el ControlNet. En modo prompt, la fase guia.
        if method == "prompt":
            subject_i = f"{character}, {_phase_for(i, n)}"
        else:
            subject_i = f"{character}, walking"

        try:
            wf = comfyui_build_sprite_sheet_workflow(
                subject_i,
                ref_image=ref_image,
                pose_skeleton=pose_skeleton,
                ckpt_name=checkpoint,
                controlnet_strength=controlnet_strength,
                controlnet_end=controlnet_end,
                transparent=True,
                negative=negative,
                width=int(width),
                height=int(height),
                steps=int(steps),
                cfg=float(cfg),
                seed=int(seed),  # FIJA: identidad consistente entre frames.
                filename_prefix=f"{filename_prefix}_f{i}",
                **gen_kwargs,
            )
        except (ValueError, TypeError) as exc:
            out["skipped"].append(i)
            if not out["error"]:
                out["error"] = f"build workflow frame {i} fallo: {exc}"
            continue

        # submit -> wait -> fetch (alta resolucion RGBA). Cualquier fallo de red/
        # GPU salta este frame y sigue (error path del DoD).
        try:
            sub = comfyui_submit_workflow(wf, server=server)
            prompt_id = sub["prompt_id"]
            outputs = comfyui_wait_result(prompt_id, server=server, timeout=wait_timeout)
        except (RuntimeError, KeyError, OSError, TimeoutError) as exc:
            out["skipped"].append(i)
            if not out["error"]:
                out["error"] = f"frame {i} fallo en submit/wait: {exc}"
            continue

        img = None
        for node_out in outputs.values():
            images = node_out.get("images") if isinstance(node_out, dict) else None
            if images:
                img = images[0]
                break
        if img is None:
            out["skipped"].append(i)
            continue

        fetched = comfyui_fetch_output_image(
            img["filename"], subfolder=img.get("subfolder", ""),
            type_=img.get("type", "output"), server=server, dest_dir=frames_dir,
        )
        if not fetched.get("ok"):
            out["skipped"].append(i)
            if not out["error"]:
                out["error"] = f"frame {i} fetch fallo: {fetched.get('error')}"
            continue

        # Pixelizar el PNG alta-res a un grid duro size x size RGBA.
        frame_px = os.path.join(frames_dir, f"{filename_prefix}_px_{i:02d}.png")
        px = comfyui_pixelize_sprite_png(
            fetched["path"], frame_px, size=int(size), colors=int(colors),
            engine=engine, palette=palette, transparent=True, autocrop=True,
            comfy_python=comfy_python,
        )
        if not px.get("ok"):
            out["skipped"].append(i)
            if not out["error"]:
                out["error"] = f"frame {i} pixelizado fallo: {px.get('error')}"
            continue

        pixel_frames.append(frame_px)

    if not pixel_frames:
        if not out["error"]:
            out["error"] = "ningun frame se genero correctamente"
        return out

    # --- Ensamblado: sprite sheet horizontal + animacion en loop. ---
    asm = assemble_animated_sprite(
        pixel_frames, dest, name=filename_prefix, fps=int(fps), fmt=fmt,
        loop=True, spritesheet=True,
    )
    if not asm.get("ok"):
        out["frames"] = pixel_frames
        out["n_frames"] = len(pixel_frames)
        if not out["error"]:
            out["error"] = f"ensamblado fallo: {asm.get('error')}"
        return out

    out["frames"] = pixel_frames
    out["spritesheet_path"] = asm.get("spritesheet_path", "")
    out["animation_path"] = asm.get("animation_path", "")
    out["n_frames"] = len(pixel_frames)
    # El size real lo confirma el primer frame ensamblado.
    fs = asm.get("frame_size") or [int(size), int(size)]
    out["size"] = int(fs[0])

    if not keep_frames:
        for fp in pixel_frames:
            try:
                os.remove(fp)
            except OSError:
                pass

    out["ok"] = True
    # El error de fallback de poses (si lo hubo) es informativo, no invalida el ok.
    return out


if __name__ == "__main__":
    import json

    res = comfyui_walk_cycle_oneshot(
        "pixel art knight, full body, side view",
        frames=4, size=32, colors=16, fps=8, seed=42,
        dest_dir="/tmp/comfy_walk_cycle",
    )
    print(json.dumps(res, indent=2))
