"""Tests del mixer comfyui_compose_capabilities (funcion pura).

Cubre: base intacto sin capacidades, combinaciones (solo loras; loras+facedetailer;
ipadapter+lora; hires), error paths (controlnet/ipadapter incompatibles), pureza,
conexiones validas en todas, y que activar una capacidad cambia el grafo.
"""

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_compose_capabilities import comfyui_compose_capabilities
from _comfyui_wf_assert import assert_api_format, class_types, node_by_ct


def _base():
    return comfyui_build_txt2img_workflow("IMG_dreamshaper_8.safetensors", "POS", "NEG")


def _count_ct(wf, ct):
    return sum(1 for n in wf.values() if n["class_type"] == ct)


def test_sin_capacidades_base_intacto():
    base = _base()
    out = comfyui_compose_capabilities(base)
    assert_api_format(out)
    # Mismos class_types y mismo numero de nodos que el base.
    assert class_types(out) == class_types(base)
    assert set(out) == set(base)


def test_solo_loras_encadena():
    base = _base()
    out = comfyui_compose_capabilities(
        base,
        loras=[
            {"name": "SD15_3d_render_redmond.safetensors", "strength_model": 0.9},
            {"name": "SD15_detail_tweaker.safetensors", "strength_model": 0.5},
        ],
    )
    assert_api_format(out)
    assert _count_ct(out, "LoraLoader") == 2
    # El KSampler cuelga del ultimo LoraLoader.
    ks = node_by_ct(out, "KSampler")
    assert out[ks["inputs"]["model"][0]]["class_type"] == "LoraLoader"


def test_loras_mas_facedetailer():
    out = comfyui_compose_capabilities(
        _base(),
        loras=[{"name": "SD15_detail_tweaker.safetensors", "strength_model": 0.5}],
        facedetailer={"denoise": 0.45},
    )
    assert_api_format(out)
    assert _count_ct(out, "LoraLoader") == 1
    assert _count_ct(out, "FaceDetailer") == 1
    # Tras facedetailer queda un unico SaveImage (el del detailer).
    assert _count_ct(out, "SaveImage") == 1
    save = node_by_ct(out, "SaveImage")
    assert out[save["inputs"]["images"][0]]["class_type"] == "FaceDetailer"


def test_ipadapter_mas_lora_toma_model_del_lora():
    out = comfyui_compose_capabilities(
        _base(),
        loras=[{"name": "a.safetensors"}],
        ipadapter={"ref_image": "ref.png", "mode": "style", "weight": 0.8},
    )
    assert_api_format(out)
    lora_id = next(nid for nid, n in out.items() if n["class_type"] == "LoraLoader")
    loader = node_by_ct(out, "IPAdapterUnifiedLoader")
    assert loader["inputs"]["model"] == [lora_id, 0]
    # KSampler.model cuelga del IPAdapter.
    ks = node_by_ct(out, "KSampler")
    assert out[ks["inputs"]["model"][0]]["class_type"] == "IPAdapter"


def test_hires_anade_upscale():
    out = comfyui_compose_capabilities(_base(), hires={})
    assert_api_format(out)
    assert "UltimateSDUpscale" in class_types(out)
    assert "UpscaleModelLoader" in class_types(out)


def test_facedetailer_detecta_checkpoint_y_prompt():
    # Sin pasar ckpt_name ni positive, se detectan del workflow.
    out = comfyui_compose_capabilities(_base(), facedetailer={})
    fd = node_by_ct(out, "FaceDetailer")
    # El FaceDetailer usa el checkpoint del base (reutilizado).
    assert "FaceDetailer" in class_types(out)
    # El CLIPTextEncode positivo del detailer lleva el texto del base.
    pos_text = out[fd["inputs"]["positive"][0]]["inputs"]["text"]
    assert pos_text == "POS"


def test_controlnet_sin_imagen_propaga_valueerror():
    with pytest.raises(ValueError):
        comfyui_compose_capabilities(_base(), controlnet={"cn_name": "m.pth"})


def test_ipadapter_sin_ref_propaga_valueerror():
    with pytest.raises(ValueError):
        comfyui_compose_capabilities(_base(), ipadapter={"mode": "style"})


def test_no_muta_base():
    base = _base()
    snapshot = set(base)
    comfyui_compose_capabilities(
        base,
        loras=[{"name": "a.safetensors"}],
        facedetailer={},
    )
    assert set(base) == snapshot


def test_activar_capacidad_cambia_grafo():
    base = _base()
    plain = comfyui_compose_capabilities(base)
    with_lora = comfyui_compose_capabilities(base, loras=[{"name": "a.safetensors"}])
    with_fd = comfyui_compose_capabilities(base, facedetailer={})
    # Cada activacion introduce class_types nuevos respecto al base.
    assert "LoraLoader" not in class_types(plain)
    assert "LoraLoader" in class_types(with_lora)
    assert "FaceDetailer" not in class_types(plain)
    assert "FaceDetailer" in class_types(with_fd)
    assert len(set(with_lora)) > len(set(plain))


def test_combinacion_controlnet_ipadapter_lora_valida():
    # ControlNet + IPAdapter + LoRA juntos producen api format valido.
    out = comfyui_compose_capabilities(
        _base(),
        loras=[{"name": "a.safetensors"}],
        controlnet={"control_image": "ctrl.png", "cn_name": "cn.pth", "strength": 0.6},
        ipadapter={"ref_image": "ref.png", "mode": "style"},
    )
    assert_api_format(out)
    cts = class_types(out)
    assert {"LoraLoader", "ControlNetApply", "IPAdapter"} <= cts
