"""Tests de estructura, orden y pureza para comfyui_inject_multi_lora (funcion pura)."""

import os
import sys

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

import pytest

from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
from ml.comfyui_inject_multi_lora import comfyui_inject_multi_lora
from _comfyui_wf_assert import assert_api_format, class_types


def _lora_nodes(wf):
    """Mapa lora_name -> (node_id, inputs) de los LoraLoader del workflow."""
    return {
        n["inputs"]["lora_name"]: (nid, n["inputs"])
        for nid, n in wf.items()
        if n["class_type"] == "LoraLoader"
    }


def test_encadena_n_loras():
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    inj = comfyui_inject_multi_lora(
        base,
        [
            {"name": "a.safetensors", "strength_model": 0.9},
            {"name": "b.safetensors", "strength_model": 0.5},
            {"name": "c.safetensors", "strength_model": 0.3},
        ],
    )
    assert_api_format(inj)
    loras = _lora_nodes(inj)
    assert set(loras) == {"a.safetensors", "b.safetensors", "c.safetensors"}


def test_orden_y_cadena_validos():
    # Cadena esperada: checkpoint -> a -> b -> KSampler/CLIPTextEncode.
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    inj = comfyui_inject_multi_lora(
        base,
        [{"name": "a.safetensors"}, {"name": "b.safetensors"}],
    )
    loras = _lora_nodes(inj)
    a_id, a_in = loras["a.safetensors"]
    b_id, b_in = loras["b.safetensors"]
    ckpt = next(nid for nid, n in inj.items() if n["class_type"] == "CheckpointLoaderSimple")

    # 'a' (primer elemento) toma el MODEL del checkpoint.
    assert a_in["model"] == [ckpt, 0]
    # 'b' (segundo elemento) toma el MODEL de 'a' (salida slot 0).
    assert b_in["model"] == [a_id, 0]
    # El KSampler queda al final de la cadena: toma el MODEL de 'b'.
    ks = next(n for n in inj.values() if n["class_type"] == "KSampler")
    assert ks["inputs"]["model"] == [b_id, 0]
    # El CLIP tambien se encadena hasta 'b' (salida slot 1).
    cte = next(n for n in inj.values() if n["class_type"] == "CLIPTextEncode")
    assert cte["inputs"]["clip"] == [b_id, 1]


def test_respeta_pesos_por_posicion():
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    inj = comfyui_inject_multi_lora(
        base,
        [
            {"name": "a.safetensors", "strength_model": 0.9, "strength_clip": 0.8},
            {"name": "b.safetensors", "strength_model": 0.4},
        ],
    )
    loras = _lora_nodes(inj)
    assert loras["a.safetensors"][1]["strength_model"] == 0.9
    assert loras["a.safetensors"][1]["strength_clip"] == 0.8
    assert loras["b.safetensors"][1]["strength_model"] == 0.4
    # default strength_clip = 1.0 cuando no se especifica
    assert loras["b.safetensors"][1]["strength_clip"] == 1.0


def test_no_muta_la_entrada():
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    claves_antes = set(base)
    _ = comfyui_inject_multi_lora(base, [{"name": "a.safetensors"}, {"name": "b.safetensors"}])
    assert "LoraLoader" not in class_types(base)
    assert set(base) == claves_antes


def test_lista_vacia_devuelve_copia_sin_loras():
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    inj = comfyui_inject_multi_lora(base, [])
    assert "LoraLoader" not in class_types(inj)
    assert class_types(inj) == class_types(base)


def test_lora_sin_name_lanza_valueerror():
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    with pytest.raises(ValueError):
        comfyui_inject_multi_lora(base, [{"strength_model": 0.5}])
    with pytest.raises(ValueError):
        comfyui_inject_multi_lora(base, ["a.safetensors"])  # no es dict


def test_determinista():
    base = comfyui_build_txt2img_workflow("ck.safetensors", "POS", "NEG")
    spec = [{"name": "a.safetensors"}, {"name": "b.safetensors"}]
    assert comfyui_inject_multi_lora(base, spec) == comfyui_inject_multi_lora(base, spec)
