"""Tests offline de comfyui_build_vfx_spritesheet_workflow (estructura del dict, sin GPU)."""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ml.comfyui_build_vfx_spritesheet_workflow import (  # noqa: E402
    comfyui_build_vfx_spritesheet_workflow,
)


def _by_class(wf, cls):
    return [n for n in wf.values() if n["class_type"] == cls]


def _id_of(wf, cls):
    return next(nid for nid, n in wf.items() if n["class_type"] == cls)


def test_golden_animatediff_loop():
    wf = comfyui_build_vfx_spritesheet_workflow(
        "burning flame, on pure black background", num_frames=16
    )
    assert len(_by_class(wf, "ADE_AnimateDiffLoaderGen1")) == 1
    assert len(_by_class(wf, "ADE_LoopedUniformContextOptions")) == 1
    ade = _by_class(wf, "ADE_AnimateDiffLoaderGen1")[0]
    assert ade["inputs"]["model_name"] == "mm_sd_v15_v2.ckpt"
    ctx = _by_class(wf, "ADE_LoopedUniformContextOptions")[0]
    assert ctx["inputs"]["closed_loop"] is True
    # batch_size = num_frames.
    latent = _by_class(wf, "EmptyLatentImage")[0]
    assert latent["inputs"]["batch_size"] == 16
    # KSampler toma el MODEL con motion.
    ade_id = _id_of(wf, "ADE_AnimateDiffLoaderGen1")
    ks = next(n for n in wf.values() if n["class_type"] == "KSampler")
    assert ks["inputs"]["model"] == [ade_id, 0]
    # context_options del loader apunta al nodo de contexto.
    ctx_id = _id_of(wf, "ADE_LoopedUniformContextOptions")
    assert ade["inputs"]["context_options"] == [ctx_id, 0]


def test_edge_open_loop():
    wf = comfyui_build_vfx_spritesheet_workflow("smoke", closed_loop=False)
    ctx = _by_class(wf, "ADE_LoopedUniformContextOptions")[0]
    assert ctx["inputs"]["closed_loop"] is False


def test_edge_lora_fx():
    wf = comfyui_build_vfx_spritesheet_workflow(
        "explosion", lora="SD15_detail_tweaker.safetensors", lora_strength=1.0
    )
    loras = _by_class(wf, "LoraLoader")
    assert len(loras) == 1
    # El ADE loader toma el MODEL de la LoRA (no del checkpoint crudo).
    lora_id = _id_of(wf, "LoraLoader")
    ade = _by_class(wf, "ADE_AnimateDiffLoaderGen1")[0]
    assert ade["inputs"]["model"] == [lora_id, 0]


def test_edge_animatelcm_params():
    wf = comfyui_build_vfx_spritesheet_workflow(
        "portal", sampler_name="lcm", steps=6, cfg=2.0
    )
    ks = next(n for n in wf.values() if n["class_type"] == "KSampler")
    assert ks["inputs"]["sampler_name"] == "lcm"
    assert ks["inputs"]["steps"] == 6 and ks["inputs"]["cfg"] == 2.0


def test_error_empty_prompt():
    try:
        comfyui_build_vfx_spritesheet_workflow("")
        assert False
    except ValueError as e:
        assert "prompt" in str(e)


def test_error_frames_lt_context():
    try:
        comfyui_build_vfx_spritesheet_workflow("flame", num_frames=8, context_length=16)
        assert False
    except ValueError as e:
        assert "num_frames" in str(e)


def test_determinism():
    a = comfyui_build_vfx_spritesheet_workflow("flame", seed=0)
    b = comfyui_build_vfx_spritesheet_workflow("flame", seed=0)
    assert a == b
