"""Construye el workflow ComfyUI AnimateDiff loop para frames de VFX en API format.

Genera N frames de un efecto 2D (humo, fuego, magia, portal) en bucle seamless
sobre fondo NEGRO (report 0140). Los frames son el insumo de:
  - comfyui_matting_luma_to_alpha  -> luminance-as-alpha (brillante=opaco, negro=transparente)
  - comfyui_build_grid / spritesheet -> montaje del spritesheet RGBA + JSON sidecar
Esos dos pasos van en un pipeline, NO en este workflow (el builder solo arma el grafo).

Vía recomendada: AnimateDiff SD1.5 con closed_loop (unica que cierra el ciclo
frame N -> 0 sin costura). Custom node ComfyUI-AnimateDiff-Evolved + motion module
mm_sd_v15_v2.ckpt YA instalados (verificado en /object_info).

Compone comfyui_build_txt2img_workflow (base) + comfyui_inject_lora (LoRA FX opcional)
e inserta la rama AnimateDiff. class_types/inputs verificados contra /object_info:
  - ADE_AnimateDiffLoaderGen1: model(MODEL), model_name(enum ['mm_sd_v15_v2.ckpt']),
    beta_schedule(enum, 'sqrt_linear (AnimateDiff)'), context_options(opcional). -> MODEL.
    OJO: el input es 'model_name', NO 'motion_model' (el plan 0140 lo nombraba mal).
  - ADE_LoopedUniformContextOptions: context_length, context_stride, context_overlap,
    closed_loop(BOOLEAN). -> CONTEXT_OPTS.
  - El KSampler se repunta al MODEL con motion; EmptyLatentImage.batch_size = num_frames.

Generar sobre fondo negro es deliberado (insumo de luma-as-alpha): el prompt deberia
incluir "on pure black background". context_length <= num_frames (ventana <= batch).

Funcion pura: sin red, sin I/O. Determinista. No valida contra el server (si el
node ADE_* o el motion module faltaran, comfyui_submit_workflow daria HTTP 400).
"""
from __future__ import annotations

import copy
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))


def _is_link(v) -> bool:
    return isinstance(v, list) and len(v) == 2 and isinstance(v[0], str) and isinstance(v[1], int)


def comfyui_build_vfx_spritesheet_workflow(
    prompt: str,
    *,
    checkpoint: str = "IMG_dreamshaper_8.safetensors",
    motion_model: str = "mm_sd_v15_v2.ckpt",
    beta_schedule: str = "sqrt_linear (AnimateDiff)",
    lora: str | None = None,
    lora_strength: float = 1.1,
    negative: str = "low quality, watermark, text, background detail",
    width: int = 512,
    height: int = 512,
    num_frames: int = 16,
    context_length: int = 16,
    context_stride: int = 1,
    context_overlap: int = 4,
    closed_loop: bool = True,
    steps: int = 20,
    cfg: float = 7.5,
    sampler_name: str = "euler",
    scheduler: str = "normal",
    seed: int = 0,
    filename_prefix: str = "vfx_loop",
) -> dict:
    """Construye el dict (API format) del workflow AnimateDiff loop de VFX.

    Args:
        prompt: prompt del efecto. Deberia incluir "on pure black background"
            (insumo de luma-as-alpha). No puede estar vacio.
        checkpoint: checkpoint SD1.5 (default 'IMG_dreamshaper_8.safetensors'; AnimateDiff
            SD1.5 cabe en 8GB, SDXL-video no). keyword-only.
        motion_model: motion module en models/animatediff_models
            (default 'mm_sd_v15_v2.ckpt'). Se asigna al input 'model_name' del loader.
        beta_schedule: schedule del motion (default 'sqrt_linear (AnimateDiff)').
        lora: LoRA de FX opcional (humo/fuego/explosion) en models/loras. None = sin LoRA.
        lora_strength: fuerza del LoRA FX sobre model y clip (recomendado 1.1).
        negative: prompt negativo.
        width, height: resolucion por frame (512x512 cabe en 8GB; bajar si OOM).
        num_frames: nº de frames del batch (EmptyLatentImage.batch_size). Debe ser
            >= context_length.
        context_length: tamano de la ventana de contexto temporal (<= num_frames).
        context_stride, context_overlap: paso y solape de las ventanas (overlap
            alto = transicion mas suave a mas coste).
        closed_loop: si True (recomendado) el ultimo frame enlaza con el primero
            (loop seamless). False = secuencia abierta.
        steps, cfg, sampler_name, scheduler, seed, filename_prefix: parametros de
            generacion. Para AnimateLCM usar (sampler 'lcm', steps 6, cfg 2.0).

    Returns:
        dict en API format listo para comfyui_submit_workflow: txt2img base
        (+ LoRA FX opcional) con ADE_LoopedUniformContextOptions +
        ADE_AnimateDiffLoaderGen1 inyectados, el KSampler repuntado al MODEL con
        motion y EmptyLatentImage.batch_size = num_frames. SaveImage escribe los
        N PNG del batch a output/.

    Raises:
        ValueError: si prompt esta vacio o num_frames < context_length (la ventana
            no puede ser mayor que el batch).
    """
    from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow

    if not prompt or not prompt.strip():
        raise ValueError("comfyui_build_vfx_spritesheet_workflow: 'prompt' no puede estar vacio")
    if num_frames < context_length:
        raise ValueError(
            "comfyui_build_vfx_spritesheet_workflow: num_frames "
            f"({num_frames}) debe ser >= context_length ({context_length})"
        )

    wf = comfyui_build_txt2img_workflow(
        checkpoint,
        prompt,
        negative,
        steps=steps,
        cfg=cfg,
        width=width,
        height=height,
        seed=seed,
        sampler_name=sampler_name,
        scheduler=scheduler,
        filename_prefix=filename_prefix,
    )

    if lora:
        from ml.comfyui_inject_lora import comfyui_inject_lora

        wf = comfyui_inject_lora(
            wf, lora, strength_model=lora_strength, strength_clip=lora_strength
        )

    wf = copy.deepcopy(wf)

    # batch_size del latente = nº de frames del loop.
    latent_id = next(
        (nid for nid, n in wf.items() if n.get("class_type") == "EmptyLatentImage"), None
    )
    if latent_id is None:
        raise ValueError(
            "comfyui_build_vfx_spritesheet_workflow: no se encontro EmptyLatentImage en la base"
        )
    wf[latent_id]["inputs"]["batch_size"] = num_frames

    ksampler_id = next(
        (nid for nid, n in wf.items() if str(n.get("class_type", "")).endswith("KSampler")), None
    )
    if ksampler_id is None:
        raise ValueError("comfyui_build_vfx_spritesheet_workflow: no se encontro KSampler en la base")
    model_src = list(wf[ksampler_id]["inputs"]["model"])

    numeric = [int(k) for k in wf.keys() if str(k).isdigit()]
    base = (max(numeric) + 1) if numeric else len(wf) + 1
    ctx_id, ade_id = str(base), str(base + 1)

    wf[ctx_id] = {
        "class_type": "ADE_LoopedUniformContextOptions",
        "inputs": {
            "context_length": context_length,
            "context_stride": context_stride,
            "context_overlap": context_overlap,
            "closed_loop": closed_loop,
        },
    }
    wf[ade_id] = {
        "class_type": "ADE_AnimateDiffLoaderGen1",
        "inputs": {
            "model": model_src,
            "model_name": motion_model,
            "beta_schedule": beta_schedule,
            "context_options": [ctx_id, 0],
        },
    }
    wf[ksampler_id]["inputs"]["model"] = [ade_id, 0]

    return wf


if __name__ == "__main__":
    import json

    wf = comfyui_build_vfx_spritesheet_workflow(
        "burning campfire flame, glowing, vfx, on pure black background",
        lora=None,
        num_frames=16,
        closed_loop=True,
        seed=0,
    )
    print(json.dumps({"nodes": list(wf), "classes": sorted({n["class_type"] for n in wf.values()})}, indent=2))
