"""Construye el workflow ComfyUI que EDITA SOLO UNA REGION de un asset ya pintado (inpaint).

Es el cuarto vertice del eje `transform` del catalogo gamedev-2d. Los otros tres
parten de:

  - txt2img (enemy_creature, item_icon): inventan la forma desde texto en blanco.
  - img2img (asset_variant): reescriben TODO el asset conservando silueta/pose.
  - ControlNet (sprite_from_sketch): pintan un sprite desde la silueta de un boceto.

Este builder cubre el dolor que ninguno resuelve: tienes un asset terminado y quieres
cambiar SOLO una parte -- ponerle otra arma, quitarle el casco, anadirle un escudo,
reparar una zona danada -- dejando el RESTO del pixel intacto. Eso es inpaint: una
mascara marca que region se repinta (blanco = editar, negro = conservar) y el sampler
solo regenera ahi.

Diferencia clave con los hermanos:

  - vs asset_variant (img2img): variant reescribe el asset ENTERO con denoise medio;
    aqui solo cambia la region enmascarada y el resto se preserva bit a bit fuera del
    crecimiento de la mascara.
  - vs sprite_from_sketch (ControlNet): sketch parte de un dibujo de lineas que guia
    la forma de un sprite NUEVO; aqui se parte de un asset YA pintado y de una mascara
    que delimita la zona a editar.

Mecanismo (modo por defecto `vae_encode`):

    CheckpointLoaderSimple -> [LoraLoader opcional de estilo] -> KSampler
    LoadImage(asset)  -> [ImageScale opcional] ->\
    LoadImageMask(mask) ----------------------- > VAEEncodeForInpaint(grow_mask_by) -> KSampler.latent
    CLIPTextEncode(prompt de la region + estilo + "seamless blend") -> KSampler.positive
    KSampler(denoise alto) -> VAEDecode -> SaveImage

`VAEEncodeForInpaint` codifica el latente respetando la mascara y dilata sus bordes
`grow_mask_by` pixeles para que la costura entre lo viejo y lo nuevo quede difuminada
(sin bordes duros visibles). En inpaint el `denoise` suele ser alto (1.0): la region
se reescribe por completo a partir del prompt; un denoise menor conserva parte de los
pixeles originales bajo la mascara (util para reparaciones suaves).

Modo de degradacion `noise_mask` (para el error path): si el servidor NO expone
`VAEEncodeForInpaint`, el caller -- que SI puede consultar /object_info -- pide
`mode="noise_mask"` y el builder arma el camino equivalente con nodos basicos:

    LoadImage -> VAEEncode -> SetLatentNoiseMask(mask) -> KSampler.latent

opcionalmente dilatando la mascara con `GrowMask(expand=grow_mask)` antes de aplicarla,
ya que `SetLatentNoiseMask` no crece la mascara por si mismo.

Compone:
  - comfyui_build_inpaint_workflow -> base inpaint (Checkpoint/LoadImage/LoadImageMask/
    VAEEncodeForInpaint/KSampler/VAEDecode/SaveImage).
  - comfyui_inject_lora            -> LoRA de estilo opcional (coherencia con el set).

El unico codigo propio es: el prompt gamedev (region + estilo + mezcla limpia), el
repunte de `grow_mask_by` y `filename_prefix`, el escalado opcional CONSISTENTE de
imagen y mascara (escalar solo una las desalinearia), y la construccion del camino
`noise_mask` de degradacion.

class_types/inputs verificados contra /object_info del servidor (8GB lowvram):
CheckpointLoaderSimple, LoadImage, LoadImageMask, ImageScale, ImageToMask,
VAEEncode, VAEEncodeForInpaint, SetLatentNoiseMask, GrowMask, CLIPTextEncode,
KSampler, VAEDecode, SaveImage, LoraLoader. Nota: ImageScale en este servidor NO
ofrece 'lanczos'; los metodos validos son 'nearest-exact', 'bilinear', 'area',
'bicubic'.

Funcion pura: sin red, sin I/O. No muta dicts de entrada (copia profunda al insertar
nodos). NO valida que input_image/mask_image/checkpoint/lora existan en el servidor
(eso es responsabilidad del caller / comfyui_validate_workflow antes de enviar).
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__), ".."))

# Negativo por defecto pensado para edicion local de un asset: la region nueva debe
# fundirse con el resto (sin costuras ni bordes de mascara visibles), sin duplicar el
# sujeto ni meter texto/marcas. NO restringe material/color (el prompt manda en la zona).
_INPAINT_ASSET_NEGATIVE = (
    "visible seam, visible mask edge, hard border, halo, blurry, lowres, "
    "deformed, bad anatomy, extra limbs, duplicate, multiple subjects, "
    "extra objects, text, watermark, signature, logo, jpeg artifacts"
)

_VALID_MODES = ("vae_encode", "noise_mask")


def _new_id(wf: dict) -> str:
    """Devuelve un node_id numerico libre (max id existente + 1)."""
    numeric = [int(k) for k in wf.keys() if str(k).isdigit()]
    return str((max(numeric) + 1) if numeric else len(wf) + 1)


def _find_class(wf: dict, class_type: str) -> str | None:
    """Primer node_id cuyo class_type coincide exactamente, o None."""
    return next(
        (nid for nid, n in wf.items() if n.get("class_type") == class_type), None
    )


def _scale_inpaint_to_size(
    wf: dict, *, size: int, img_method: str, mask_method: str, crop: str
) -> dict:
    """Escala imagen base Y mascara al mismo size para no desalinearlas.

    Inserta un ImageScale entre el LoadImage del asset y el VAEEncodeForInpaint, y
    reemplaza el LoadImageMask por LoadImage(mask) -> ImageScale -> ImageToMask(red)
    de modo que ambos lleguen al nodo de inpaint a size x size con el mismo recorte.
    Escalar solo la imagen (dejando la mascara a su tamano original) produciria una
    mascara que ya no cubre la region correcta: por eso aqui se escalan las DOS o
    ninguna. Pura: trabaja sobre una copia profunda.
    """
    wf = copy.deepcopy(wf)
    enc_id = _find_class(wf, "VAEEncodeForInpaint")
    if enc_id is None:
        raise ValueError(
            "comfyui_build_inpaint_asset_workflow: no se encontro VAEEncodeForInpaint "
            "para escalar (size solo aplica al modo vae_encode)"
        )
    enc = wf[enc_id]["inputs"]
    img_src = list(enc.get("pixels", []))
    mask_link = list(enc.get("mask", []))
    if not img_src or not mask_link:
        raise ValueError(
            "comfyui_build_inpaint_asset_workflow: VAEEncodeForInpaint sin pixels/mask"
        )

    # Imagen base: LoadImage[0] -> ImageScale -> VAEEncodeForInpaint.pixels
    scale_img = _new_id(wf)
    wf[scale_img] = {
        "class_type": "ImageScale",
        "inputs": {
            "image": img_src,
            "upscale_method": img_method,
            "width": int(size),
            "height": int(size),
            "crop": crop,
        },
    }
    enc["pixels"] = [scale_img, 0]

    # Mascara: el nodo fuente (LoadImageMask) se convierte en LoadImage para poder
    # escalarla como imagen y volver a MASK con ImageToMask al mismo size.
    mask_node_id = mask_link[0]
    mask_node = wf[mask_node_id]
    mask_file = mask_node["inputs"].get("image", "")
    mask_node["class_type"] = "LoadImage"
    mask_node["inputs"] = {"image": mask_file}

    scale_mask = _new_id(wf)
    wf[scale_mask] = {
        "class_type": "ImageScale",
        "inputs": {
            "image": [mask_node_id, 0],
            "upscale_method": mask_method,
            "width": int(size),
            "height": int(size),
            "crop": crop,
        },
    }
    to_mask = _new_id(wf)
    wf[to_mask] = {
        "class_type": "ImageToMask",
        "inputs": {"image": [scale_mask, 0], "channel": "red"},
    }
    enc["mask"] = [to_mask, 0]
    return wf


def _to_noise_mask(wf: dict, *, grow_mask: int) -> dict:
    """Reemplaza el camino VAEEncodeForInpaint por VAEEncode + SetLatentNoiseMask.

    Degradacion para servidores sin VAEEncodeForInpaint: codifica la imagen con un
    VAEEncode normal y aplica la mascara al latente con SetLatentNoiseMask, opcional-
    mente dilatandola antes con GrowMask(expand=grow_mask) porque SetLatentNoiseMask
    no crece la mascara. Repunta KSampler.latent_image al nuevo SetLatentNoiseMask.
    Pura: copia profunda.
    """
    wf = copy.deepcopy(wf)
    enc_id = _find_class(wf, "VAEEncodeForInpaint")
    if enc_id is None:
        raise ValueError(
            "comfyui_build_inpaint_asset_workflow: base sin VAEEncodeForInpaint; "
            "no se puede degradar a noise_mask"
        )
    enc = wf[enc_id]["inputs"]
    pixels = list(enc["pixels"])
    vae = list(enc["vae"])
    mask_link = list(enc["mask"])

    # VAEEncode normal en el mismo node_id (reutiliza el id; cae el grow_mask_by).
    wf[enc_id] = {
        "class_type": "VAEEncode",
        "inputs": {"pixels": pixels, "vae": vae},
    }

    mask_out = mask_link
    if grow_mask and int(grow_mask) != 0:
        grow_id = _new_id(wf)
        wf[grow_id] = {
            "class_type": "GrowMask",
            "inputs": {
                "mask": mask_link,
                "expand": int(grow_mask),
                "tapered_corners": True,
            },
        }
        mask_out = [grow_id, 0]

    setmask = _new_id(wf)
    wf[setmask] = {
        "class_type": "SetLatentNoiseMask",
        "inputs": {"samples": [enc_id, 0], "mask": mask_out},
    }

    ks_id = _find_class(wf, "KSampler")
    if ks_id is not None:
        wf[ks_id]["inputs"]["latent_image"] = [setmask, 0]
    return wf


def comfyui_build_inpaint_asset_workflow(
    input_image: str,
    mask_image: str,
    prompt: str,
    *,
    checkpoint: str = "IMG_dreamshaper_8.safetensors",
    denoise: float = 1.0,
    style: str = "game asset",
    grow_mask: int = 6,
    size: int | None = None,
    seed: int = 0,
    lora: str | None = None,
    lora_strength: float = 1.0,
    mode: str = "vae_encode",
    negative: str | None = None,
    steps: int = 28,
    cfg: float = 7.0,
    sampler_name: str = "dpmpp_2m",
    scheduler: str = "karras",
    img_upscale_method: str = "bilinear",
    mask_upscale_method: str = "nearest-exact",
    crop: str = "disabled",
    filename_prefix: str = "inpaint_asset",
) -> dict:
    """Construye el dict (API format) de un inpaint que edita SOLO la region enmascarada.

    Args:
        input_image: nombre del archivo del asset a editar dentro de la carpeta input/
            del servidor ComfyUI (un sprite/icono YA pintado). Lo carga LoadImage.
            Subelo antes con POST /upload/image o copialo a ~/ComfyUI/input/. No puede
            estar vacio.
        mask_image: nombre del archivo de la mascara dentro de input/ del servidor.
            BLANCO = region a editar/repintar, NEGRO = region a conservar. Debe tener
            la MISMA resolucion que input_image (salvo que pases `size`, que reescala
            ambas de forma consistente). No puede estar vacio: una edicion sin mascara
            no es inpaint -- para reescribir el asset entero usa asset_variant.
        prompt: que poner en la region enmascarada ("a golden sword", "a blue shield",
            "empty background", "intact armor"). Es lo unico que se regenera; el resto
            del asset se conserva. No puede estar vacio.
        checkpoint: checkpoint del servidor. 'IMG_dreamshaper_8.safetensors' (SD1.5, holgado
            en 8GB lowvram) por defecto. keyword-only.
        denoise: fuerza de denoising del KSampler dentro de la mascara. En inpaint suele
            ser alto: 1.0 (por defecto) reescribe la region por completo; <1.0 conserva
            parte de los pixeles originales (reparacion suave). Se clampa a [0.0, 1.0].
            keyword-only.
        style: descriptor de estilo que mantiene la region coherente con el asset y el
            set ("game asset", "dark fantasy creature", "pixel art"). Pasa el MISMO
            style + checkpoint + (lora) que el resto del set. keyword-only.
        grow_mask: pixeles que se dilata el borde de la mascara para difuminar la
            costura entre lo viejo y lo nuevo (evita bordes duros). Se clampa a [0, 64]
            (limite del nodo VAEEncodeForInpaint). En modo noise_mask se aplica via
            GrowMask. keyword-only.
        size: lado en px al que se NORMALIZAN imagen Y mascara antes de inpaint. None
            (por defecto) = no escala: el inpaint conserva la resolucion nativa del
            asset y exige que mascara e imagen ya coincidan en tamano (lo recomendado;
            los assets del set ya salen a 512). Un int reescala AMBAS a size x size de
            forma consistente. Solo aplica al modo vae_encode. keyword-only.
        seed: semilla del KSampler. keyword-only.
        lora: LoRA de estilo opcional en models/loras (ej. 'SD15_dark_fantasy.safetensors').
            None = sin LoRA. keyword-only.
        lora_strength: fuerza del LoRA sobre model y clip. Se clampa a [0.0, 2.0].
            keyword-only.
        mode: 'vae_encode' (por defecto, usa VAEEncodeForInpaint con grow_mask nativo)
            o 'noise_mask' (degrada a VAEEncode + SetLatentNoiseMask para servidores sin
            VAEEncodeForInpaint). El caller decide el modo tras consultar /object_info.
            keyword-only.
        negative: prompt negativo. None usa el negativo por defecto pensado para edicion
            local (mezcla limpia, sin costuras, una figura). keyword-only.
        steps, cfg, sampler_name, scheduler, filename_prefix: parametros de generacion.
            keyword-only.
        img_upscale_method: metodo de ImageScale para la imagen cuando size no es None
            ('bilinear' por defecto; 'lanczos' NO esta disponible en este servidor).
            keyword-only.
        mask_upscale_method: metodo de ImageScale para la mascara cuando size no es None
            ('nearest-exact' por defecto, preserva bordes nitidos blanco/negro de la
            mascara). keyword-only.
        crop: modo de recorte de ImageScale ('disabled' o 'center'). keyword-only.

    Returns:
        dict en API format listo para comfyui_submit_workflow: inpaint que repinta SOLO
        la region marcada en blanco por la mascara con '{prompt}, {style}, seamless
        blend...', conservando el resto del asset. Con grow_mask para difuminar la
        costura, escalado consistente opcional y LoRA de estilo opcional.

    Raises:
        ValueError: si input_image, mask_image o prompt estan vacios; si mode no es
            'vae_encode' ni 'noise_mask'; o si la base no tiene los nodos esperados
            (propagado por los helpers).
    """
    from ml.comfyui_build_inpaint_workflow import comfyui_build_inpaint_workflow

    if not input_image or not input_image.strip():
        raise ValueError(
            "comfyui_build_inpaint_asset_workflow: 'input_image' no puede estar vacio"
        )
    if not mask_image or not mask_image.strip():
        raise ValueError(
            "comfyui_build_inpaint_asset_workflow: 'mask_image' no puede estar vacio "
            "(sin mascara no es inpaint; para reescribir el asset entero usa "
            "comfyui_build_asset_variant_workflow)"
        )
    if not prompt or not prompt.strip():
        raise ValueError(
            "comfyui_build_inpaint_asset_workflow: 'prompt' no puede estar vacio"
        )
    if mode not in _VALID_MODES:
        raise ValueError(
            f"comfyui_build_inpaint_asset_workflow: 'mode' debe ser uno de {_VALID_MODES}, "
            f"recibido {mode!r}"
        )

    input_image = input_image.strip()
    mask_image = mask_image.strip()
    prompt = prompt.strip()
    denoise = max(0.0, min(1.0, float(denoise)))
    grow_mask = max(0, min(64, int(grow_mask)))
    lora_strength = max(0.0, min(2.0, float(lora_strength)))
    neg = _INPAINT_ASSET_NEGATIVE if negative is None else negative

    # Prompt de la region: describe SOLO lo que va en la zona enmascarada y empuja a
    # que se funda con el resto (mezcla limpia, iluminacion consistente). A diferencia
    # de asset_variant NO se pide "same pose/silhouette": aqui solo cambia un trozo.
    positive = (
        f"{prompt}, {style}, seamless blend with surroundings, "
        "consistent lighting, matching art style, high detail"
    )

    wf = comfyui_build_inpaint_workflow(
        checkpoint,
        input_image,
        mask_image,
        positive,
        neg,
        denoise=denoise,
        steps=steps,
        cfg=cfg,
        seed=seed,
        sampler_name=sampler_name,
        scheduler=scheduler,
    )

    # Repuntar grow_mask_by (el base lo deja en 6) y el filename_prefix (el base lo
    # deja en 'comfy_inpaint').
    enc_id = _find_class(wf, "VAEEncodeForInpaint")
    if enc_id is not None:
        wf[enc_id]["inputs"]["grow_mask_by"] = grow_mask
    save_id = _find_class(wf, "SaveImage")
    if save_id is not None:
        wf[save_id]["inputs"]["filename_prefix"] = filename_prefix

    if mode == "vae_encode" and size is not None:
        wf = _scale_inpaint_to_size(
            wf,
            size=int(size),
            img_method=img_upscale_method,
            mask_method=mask_upscale_method,
            crop=crop,
        )

    if mode == "noise_mask":
        wf = _to_noise_mask(wf, grow_mask=grow_mask)

    if lora:
        from ml.comfyui_inject_lora import comfyui_inject_lora

        wf = comfyui_inject_lora(
            wf, lora, strength_model=lora_strength, strength_clip=lora_strength
        )

    return wf


if __name__ == "__main__":
    import json

    wf = comfyui_build_inpaint_asset_workflow(
        "enemy_creature_00001_.png",
        "goblin_hand_mask.png",
        "a glowing blue magic orb",
        style="dark fantasy creature, game asset",
        grow_mask=8,
        denoise=1.0,
        seed=7,
    )
    print(
        json.dumps(
            {
                "nodes": list(wf),
                "classes": sorted({n["class_type"] for n in wf.values()}),
                "denoise": wf["3"]["inputs"]["denoise"],
                "grow_mask_by": wf["11"]["inputs"]["grow_mask_by"],
                "positive": wf["6"]["inputs"]["text"],
                "input_image": wf["10"]["inputs"]["image"],
                "mask": wf["12"]["inputs"]["image"],
                "filename_prefix": wf["9"]["inputs"]["filename_prefix"],
            },
            indent=2,
        )
    )
