"""Construye un workflow ComfyUI txt2img + IPAdapter en API format (dict de nodos).

Parte de comfyui_build_txt2img_workflow y le injerta la rama IPAdapter del custom
node ComfyUI_IPAdapter_plus (cubiq):

- mode='style': IPAdapterUnifiedLoader + IPAdapter. La imagen de referencia
  transfiere estilo/composicion al resultado (image prompt clasico).
- mode='faceid': IPAdapterUnifiedLoaderFaceID + IPAdapterFaceID. Usa insightface
  para extraer el embedding de la cara de la referencia y el .bin FaceID + su
  LoRA para imponer un **rostro consistente** en el personaje generado.

En ambos casos la salida MODEL de la rama IPAdapter se repunta al KSampler, de
modo que el sampler genera ya condicionado por la imagen de referencia.

Funcion pura: sin red, sin I/O. Determinista para los mismos argumentos. Los
class_type/inputs estan verificados contra /object_info del servidor (IPAdapter
plus): IPAdapterUnifiedLoader(model,preset)->[MODEL,IPADAPTER],
IPAdapter(model,ipadapter,image,weight,start_at,end_at,weight_type)->[MODEL],
IPAdapterUnifiedLoaderFaceID(model,preset,lora_strength,provider)->[MODEL,IPADAPTER],
IPAdapterFaceID(model,ipadapter,image,weight,weight_faceidv2,weight_type,
combine_embeds,start_at,end_at,embeds_scaling)->[MODEL,IMAGE].
"""
from __future__ import annotations

import os
import sys

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

# Presets por defecto del IPAdapterUnifiedLoader(FaceID) segun el modo.
_DEFAULT_PRESET = {
    "style": "STANDARD (medium strength)",
    "faceid": "FACEID PLUS V2",
}
# weight_type por defecto: el nodo IPAdapter usa 'standard', el FaceID usa 'linear'.
_DEFAULT_WEIGHT_TYPE = {
    "style": "standard",
    "faceid": "linear",
}


def comfyui_build_ipadapter_workflow(
    prompt: str,
    ref_image: str,
    *,
    base_checkpoint: str,
    mode: str = "style",
    weight: float = 0.8,
    negative: str = "",
    preset: str | None = None,
    weight_type: str | None = None,
    start_at: float = 0.0,
    end_at: float = 1.0,
    weight_faceidv2: float = 1.0,
    lora_strength: float = 0.6,
    combine_embeds: str = "concat",
    embeds_scaling: str = "V only",
    provider: str = "CPU",
    steps: int = 20,
    cfg: float = 7.0,
    width: int = 512,
    height: int = 512,
    seed: int = 0,
    sampler_name: str = "euler",
    scheduler: str = "normal",
    filename_prefix: str = "ipadapter",
) -> dict:
    """Construye un workflow txt2img condicionado por una imagen de referencia.

    Args:
        prompt: prompt positivo (texto del resultado deseado).
        ref_image: nombre del archivo de imagen de referencia en el directorio
            input/ del servidor ComfyUI (lo carga un nodo LoadImage). En faceid
            debe contener una cara nitida; en style es la imagen de estilo.
        base_checkpoint: checkpoint SD1.5/SDXL (debe casar con los modelos
            IPAdapter: usa modelos SD1.5 con checkpoints SD1.5). keyword-only.
        mode: 'style' (transfiere estilo/composicion) o 'faceid' (rostro
            consistente via insightface + FaceID). keyword-only.
        weight: peso de la influencia IPAdapter (0..1+). 0.8 es un buen punto de
            partida; sube para mas parecido, baja para mas libertad del prompt.
        negative: prompt negativo.
        preset: preset del UnifiedLoader. Si None usa el default del modo
            ('STANDARD (medium strength)' para style, 'FACEID PLUS V2' para faceid).
        weight_type: tipo de ponderacion del nodo IPAdapter/FaceID. Si None usa el
            default del modo ('standard' para style, 'linear' para faceid).
        start_at: fraccion del sampling donde empieza a aplicar IPAdapter (0..1).
        end_at: fraccion del sampling donde deja de aplicar (0..1).
        weight_faceidv2: peso del embedding FaceID v2 (solo mode='faceid').
        lora_strength: fuerza de la LoRA FaceID que carga el UnifiedLoaderFaceID
            (solo mode='faceid').
        combine_embeds: como combinar embeddings si hay varias caras
            ('concat'|'add'|'subtract'|'average'|'norm average'). Solo faceid.
        embeds_scaling: escalado de embeddings ('V only'|'K+V'|...). Solo faceid.
        provider: backend de insightface ('CPU'|'CUDA'|...). CPU por defecto para
            no competir por VRAM con el modelo de difusion. Solo faceid.
        steps, cfg, width, height, seed, sampler_name, scheduler, filename_prefix:
            parametros de generacion que se pasan a comfyui_build_txt2img_workflow.

    Returns:
        dict en API format listo para comfyui_submit_workflow, con la base
        txt2img + LoadImage + la rama IPAdapter del modo elegido, y el KSampler
        repuntado a la salida MODEL de esa rama.

    Raises:
        ValueError: si mode no es 'style' ni 'faceid', si ref_image esta vacio, o
            si no se puede localizar el checkpoint/KSampler en la base.
    """
    from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow

    if mode not in ("style", "faceid"):
        raise ValueError(
            f"comfyui_build_ipadapter_workflow: mode debe ser 'style' o 'faceid', no {mode!r}"
        )
    if not ref_image:
        raise ValueError("comfyui_build_ipadapter_workflow: ref_image no puede estar vacio")

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

    ckpt = next(
        (nid for nid, n in wf.items() if str(n.get("class_type", "")).startswith("CheckpointLoader")),
        None,
    )
    ksampler = next(
        (nid for nid, n in wf.items() if str(n.get("class_type", "")).endswith("KSampler")),
        None,
    )
    if ckpt is None or ksampler is None:
        raise ValueError(
            "comfyui_build_ipadapter_workflow: no se encontro CheckpointLoader/KSampler en la base"
        )

    numeric = [int(k) for k in wf.keys() if str(k).isdigit()]
    base_id = (max(numeric) + 1) if numeric else len(wf) + 1
    load_id = str(base_id)
    loader_id = str(base_id + 1)
    apply_id = str(base_id + 2)

    used_preset = preset if preset is not None else _DEFAULT_PRESET[mode]
    used_wtype = weight_type if weight_type is not None else _DEFAULT_WEIGHT_TYPE[mode]

    # Carga la imagen de referencia (slot 0 = IMAGE).
    wf[load_id] = {
        "class_type": "LoadImage",
        "inputs": {"image": ref_image},
    }

    if mode == "style":
        wf[loader_id] = {
            "class_type": "IPAdapterUnifiedLoader",
            "inputs": {"model": [ckpt, 0], "preset": used_preset},
        }
        wf[apply_id] = {
            "class_type": "IPAdapter",
            "inputs": {
                "model": [loader_id, 0],
                "ipadapter": [loader_id, 1],
                "image": [load_id, 0],
                "weight": weight,
                "start_at": start_at,
                "end_at": end_at,
                "weight_type": used_wtype,
            },
        }
    else:  # faceid
        wf[loader_id] = {
            "class_type": "IPAdapterUnifiedLoaderFaceID",
            "inputs": {
                "model": [ckpt, 0],
                "preset": used_preset,
                "lora_strength": lora_strength,
                "provider": provider,
            },
        }
        wf[apply_id] = {
            "class_type": "IPAdapterFaceID",
            "inputs": {
                "model": [loader_id, 0],
                "ipadapter": [loader_id, 1],
                "image": [load_id, 0],
                "weight": weight,
                "weight_faceidv2": weight_faceidv2,
                "weight_type": used_wtype,
                "combine_embeds": combine_embeds,
                "start_at": start_at,
                "end_at": end_at,
                "embeds_scaling": embeds_scaling,
            },
        }

    # Repunta el KSampler para que tome el MODEL condicionado por IPAdapter.
    wf[ksampler]["inputs"]["model"] = [apply_id, 0]
    return wf


if __name__ == "__main__":
    import json

    wf_style = comfyui_build_ipadapter_workflow(
        "a fantasy castle on a hill, oil painting",
        "example.png",
        base_checkpoint="IMG_dreamshaper_8.safetensors",
        mode="style",
        weight=0.8,
    )
    wf_face = comfyui_build_ipadapter_workflow(
        "portrait of a knight in armor, cinematic",
        "showcase_char.png",
        base_checkpoint="IMG_dreamshaper_8.safetensors",
        mode="faceid",
        weight=0.9,
    )
    print(json.dumps({"style_nodes": list(wf_style), "faceid_nodes": list(wf_face)}, indent=2))
