"""Construye un workflow ComfyUI de pixel-art (SDXL + LoRA SDXL_pixel-art) en API format.

Fase 1 del pipeline pixel-art (ver report 0135): genera el *look* pixel-art con
SDXL base + la LoRA `SDXL_pixel-art` (nerijs), opcionalmente acelerada con la
LCM-LoRA SDXL para iterar en 8 steps. El resultado todavia tiene pixeles de
tamano irregular y cientos de colores: el pixel-perfect (Fase 2) lo hace
`comfyui_pixelize_image` (downscale nearest + cuantizacion), NO este workflow.

Compone funciones existentes del registry, no reescribe el grafo:
  - comfyui_build_txt2img_workflow  -> base SDXL txt2img
  - comfyui_inject_multi_lora       -> encadena SDXL_pixel-art (+ SDXL_lcm-lora)

class_types/params verificados contra /object_info del servidor (8GB lowvram):
CheckpointLoaderSimple, LoraLoader, CLIPTextEncode, EmptyLatentImage,
KSampler (sampler 'lcm', scheduler 'sgm_uniform' presentes), VAEDecode, SaveImage.
LoRAs presentes en models/loras: SDXL_pixel-art.safetensors, SDXL_lcm-lora.safetensors.

Funcion pura: sin red, sin I/O. Determinista para los mismos argumentos.
"""
from __future__ import annotations

import copy
import os
import sys

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

# Defaults canonicos de la ficha HF nerijs/pixel-art-xl (report 0135).
_LCM_DEFAULTS = {"steps": 8, "cfg": 1.5, "sampler_name": "lcm", "scheduler": "sgm_uniform"}
_PLAIN_DEFAULTS = {"steps": 25, "cfg": 7.0, "sampler_name": "euler", "scheduler": "normal"}


def _inject_rembg(workflow: dict, model: str) -> dict:
    """Inserta 'Image Rembg (Remove Background)' (transparency=True) entre VAEDecode y SaveImage.

    Mismo helper que comfyui_build_item_icon_workflow / comfyui_build_sprite_sheet_workflow:
    el nodo recorta la silueta del sujeto dejando alpha, y se repunta SaveImage.images a
    la salida del Rembg para que el PNG salga con fondo transparente. No muta el dict de
    entrada (copia profunda).
    """
    wf = copy.deepcopy(workflow)
    vaedecode_id = next(
        (nid for nid, n in wf.items() if n.get("class_type") == "VAEDecode"), None
    )
    save_id = next((nid for nid, n in wf.items() if n.get("class_type") == "SaveImage"), None)
    if vaedecode_id is None or save_id is None:
        raise ValueError(
            "comfyui_build_pixelart_workflow: no se encontro VAEDecode/SaveImage para Rembg"
        )
    numeric = [int(k) for k in wf.keys() if str(k).isdigit()]
    rembg_id = str((max(numeric) + 1) if numeric else len(wf) + 1)
    wf[rembg_id] = {
        "class_type": "Image Rembg (Remove Background)",
        "inputs": {
            "images": [vaedecode_id, 0],
            "transparency": True,
            "model": model,
            "post_processing": False,
            "only_mask": False,
            "alpha_matting": False,
            "alpha_matting_foreground_threshold": 240,
            "alpha_matting_background_threshold": 10,
            "alpha_matting_erode_size": 10,
            "background_color": "none",
        },
    }
    wf[save_id]["inputs"]["images"] = [rembg_id, 0]
    return wf


def comfyui_build_pixelart_workflow(
    positive: str,
    negative: str = "blurry, jpeg artifacts, gradient, smooth shading, anti-aliasing",
    *,
    ckpt_name: str = "IMG_juggernaut_xl_v11.safetensors",
    pixel_lora: str = "SDXL_pixel-art.safetensors",
    lora_strength: float = 1.2,
    use_lcm: bool = True,
    lcm_lora: str = "SDXL_lcm-lora.safetensors",
    lcm_strength: float = 1.0,
    steps: int | None = None,
    cfg: float | None = None,
    width: int = 1024,
    height: int = 1024,
    seed: int = 0,
    sampler_name: str | None = None,
    scheduler: str | None = None,
    transparent: bool = True,
    rembg_model: str = "u2net",
    filename_prefix: str = "pixelart",
) -> dict:
    """Construye el dict (API format) del workflow pixel-art SDXL + LoRA.

    Args:
        positive: prompt positivo. Trucos para empujar el grid: incluir 'pixel',
            'isometric view', '32x32 style'. No puede estar vacio.
        negative: prompt negativo (por defecto evita blur/gradientes/anti-alias,
            que estropean el look pixel).
        ckpt_name: checkpoint SDXL base (default 'IMG_juggernaut_xl_v11.safetensors',
            el SDXL instalado; SDXL_pixel-art es una LoRA SDXL). keyword-only.
        pixel_lora: archivo de la LoRA de estilo pixel-art en models/loras.
        lora_strength: fuerza de SDXL_pixel-art sobre model y clip (recomendado 1.2).
            Se clampa a [0.0, 2.0].
        use_lcm: si True encadena la LCM-LoRA SDXL y usa los defaults rapidos
            (8 steps, cfg 1.5, sampler 'lcm', scheduler 'sgm_uniform'); si False
            usa los defaults SDXL normales (25 steps, cfg 7, 'euler'/'normal').
        lcm_lora: archivo de la LCM-LoRA SDXL en models/loras (solo si use_lcm).
        lcm_strength: fuerza de la LCM-LoRA sobre model y clip (recomendado 1.0).
        steps, cfg, sampler_name, scheduler: si se pasan, sobreescriben el default
            del modo (LCM vs normal). None = usar el default del modo.
        width, height: resolucion base (1024x1024 SDXL; luego downscale x8 -> 128
            en la Fase 2 con comfyui_pixelize_image).
        seed: semilla del KSampler.
        transparent: si True (default) inyecta 'Image Rembg' tras el VAEDecode y el
            PNG sale con alpha (fondo recortado) — lo habitual para sprites de sujeto
            (personajes, criaturas, objetos). Si False deja la imagen opaca sobre
            fondo plano, para tiles/texturas/fondos que no quieren transparencia.
            keyword-only.
        rembg_model: modelo Rembg ('u2net' general, 'isnet-anime' para anime). Solo
            se usa si transparent=True. keyword-only.
        filename_prefix: prefijo del PNG en output/.

    Returns:
        dict en API format listo para comfyui_submit_workflow, con el
        CheckpointLoaderSimple, 1 LoraLoader (SDXL_pixel-art) o 2 (SDXL_pixel-art +
        SDXL_lcm-lora si use_lcm), KSampler con los params del modo, un nodo
        'Image Rembg' antes del SaveImage si transparent, y SaveImage.

    Raises:
        ValueError: si positive esta vacio, o si la base no tiene VAEDecode/SaveImage
            donde inyectar el Rembg (propagado por el helper, solo si transparent).
    """
    from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
    from ml.comfyui_inject_multi_lora import comfyui_inject_multi_lora

    if not positive or not positive.strip():
        raise ValueError("comfyui_build_pixelart_workflow: 'positive' no puede estar vacio")

    lora_strength = max(0.0, min(2.0, float(lora_strength)))
    lcm_strength = max(0.0, min(2.0, float(lcm_strength)))

    defaults = _LCM_DEFAULTS if use_lcm else _PLAIN_DEFAULTS
    eff_steps = defaults["steps"] if steps is None else int(steps)
    eff_cfg = defaults["cfg"] if cfg is None else float(cfg)
    eff_sampler = defaults["sampler_name"] if sampler_name is None else sampler_name
    eff_scheduler = defaults["scheduler"] if scheduler is None else scheduler

    base = comfyui_build_txt2img_workflow(
        ckpt_name,
        positive,
        negative,
        steps=eff_steps,
        cfg=eff_cfg,
        width=width,
        height=height,
        seed=seed,
        sampler_name=eff_sampler,
        scheduler=eff_scheduler,
        filename_prefix=filename_prefix,
    )

    loras = [
        {"name": pixel_lora, "strength_model": lora_strength, "strength_clip": lora_strength},
    ]
    if use_lcm:
        loras.append(
            {"name": lcm_lora, "strength_model": lcm_strength, "strength_clip": lcm_strength}
        )

    wf = comfyui_inject_multi_lora(base, loras)

    if transparent:
        wf = _inject_rembg(wf, rembg_model)

    return wf


if __name__ == "__main__":
    import json

    wf = comfyui_build_pixelart_workflow(
        "isometric tiny house, pixel, 32x32 style, vibrant colors",
        use_lcm=True,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
