"""Construye el workflow ComfyUI que EXTIENDE EL LIENZO de un asset (outpaint).

Es el quinto vertice del eje `transform` del catalogo gamedev-2d. Los otros cuatro
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.
  - inpaint (inpaint_asset): repintan SOLO una region marcada por una mascara EXTERNA
    del mismo tamano que el asset, conservando el resto.

Este builder cubre el dolor que ninguno resuelve: tienes un asset terminado y necesitas
MAS LIENZO del que tiene -- extender un fondo/parallax a otra resolucion o aspect ratio,
ampliar un card_art o un splash mas alla de sus bordes, completar la escena hacia un lado
sin regenerarla entera. Eso es outpaint: se agranda el canvas por uno o varios lados y se
genera contenido NUEVO coherente con el original en esa zona ampliada.

Diferencia clave con inpaint (su hermano mas cercano):

  - inpaint_asset RECIBE una mascara externa que el caller tiene que pintar, del MISMO
    tamano que el asset, y edita una region INTERIOR; las dimensiones del asset no cambian.
  - outpaint_asset NO recibe mascara: el lienzo se agranda y la mascara de "lo nuevo"
    (la franja anadida) la GENERA el grafo con el nodo nativo `ImagePadForOutpaint`, que
    a la vez extiende la imagen (relleno) y emite la MASK feathered de la zona ampliada.
    Las dimensiones del asset CRECEN por los lados extendidos.

Mecanismo:

    CheckpointLoaderSimple -> [LoraLoader opcional de estilo] -> KSampler
    LoadImage(asset) -> ImagePadForOutpaint(left/top/right/bottom/feathering) -> (IMAGE, MASK)
        ImagePadForOutpaint.IMAGE -> VAEEncodeForInpaint.pixels
        ImagePadForOutpaint.MASK  -> VAEEncodeForInpaint.mask
    CLIPTextEncode(prompt de lo nuevo + estilo + "seamless extension") -> KSampler.positive
    KSampler(denoise alto) -> VAEDecode -> SaveImage

`ImagePadForOutpaint` (nodo nativo de ComfyUI) hace lo que en inpaint hacian LoadImage +
LoadImageMask juntos: amplia el lienzo `left/top/right/bottom` pixeles, rellena la franja
nueva, y emite una MASK donde la zona ampliada es blanca (a regenerar) y la original
negra (a conservar), con un borde difuminado de `feathering` pixeles para que la costura
entre lo viejo y lo nuevo quede suave. `VAEEncodeForInpaint` codifica ese latente
respetando la MASK; el KSampler (denoise alto, 1.0) genera lo nuevo solo en la franja.

`grow_mask` dilata adicionalmente el borde de la mascara hacia DENTRO de la imagen
original (sobre el feathering del pad) para reforzar el blend si aparece costura; por
defecto 0 porque el feathering del propio pad ya suele bastar.

Compone:
  - comfyui_build_inpaint_workflow -> base inpaint (Checkpoint/LoadImage/VAEEncodeForInpaint/
    KSampler/VAEDecode/SaveImage). El LoadImageMask de la base NO se usa: outpaint genera
    su propia mascara con ImagePadForOutpaint, asi que ese nodo se elimina y se reconecta
    VAEEncodeForInpaint a las dos salidas del pad.
  - comfyui_inject_lora -> LoRA de estilo opcional (coherencia con el set).

El unico codigo propio es: el prompt gamedev (extension de escena + estilo + mezcla
limpia), la insercion del ImagePadForOutpaint entre LoadImage y VAEEncodeForInpaint, el
redondeo de las extensiones al multiplo de 8 que exige el nodo (y SD), y el repunte de
grow_mask_by / filename_prefix.

class_types/inputs verificados contra /object_info del servidor (8GB lowvram):
CheckpointLoaderSimple, LoadImage, ImagePadForOutpaint (image, left/top/right/bottom INT
step 8, feathering INT default 40; outputs IMAGE+MASK), VAEEncodeForInpaint, CLIPTextEncode,
KSampler, VAEDecode, SaveImage, LoraLoader.

Funcion pura: sin red, sin I/O. No muta dicts de entrada (construye desde cero via la base).
NO valida que input_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 os
import sys

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

# Negativo por defecto pensado para EXTENSION de lienzo: la franja nueva debe continuar la
# escena sin costura, sin repetir/espejar el contenido original ni meter un segundo sujeto,
# texto o marcas. NO restringe material/color (el prompt manda en la zona nueva).
_OUTPAINT_ASSET_NEGATIVE = (
    "visible seam, hard border, abrupt edge, halo, duplicate, repeated pattern, "
    "mirrored, tiling artifact, second subject, blurry, lowres, deformed, "
    "text, watermark, signature, logo, jpeg artifacts"
)


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 _round8(px: int) -> int:
    """Redondea al multiplo de 8 mas cercano (>= 0).

    `ImagePadForOutpaint` declara `step: 8` para left/top/right/bottom y SD trabaja sobre
    latentes de 8 px, asi que extender un numero no multiplo de 8 produce dimensiones que
    el VAE tendria que recortar. Se normaliza aqui para que las dims finales sean limpias.
    """
    px = max(0, int(px))
    return ((px + 4) // 8) * 8


def comfyui_build_outpaint_asset_workflow(
    input_image: str,
    prompt: str,
    *,
    left: int = 0,
    right: int = 0,
    top: int = 0,
    bottom: int = 0,
    feather: int = 40,
    checkpoint: str = "IMG_dreamshaper_8.safetensors",
    denoise: float = 1.0,
    style: str = "game background",
    grow_mask: int = 0,
    seed: int = 0,
    lora: str | None = None,
    lora_strength: float = 1.0,
    negative: str | None = None,
    steps: int = 28,
    cfg: float = 7.0,
    sampler_name: str = "dpmpp_2m",
    scheduler: str = "karras",
    filename_prefix: str = "outpaint_asset",
) -> dict:
    """Construye el dict (API format) de un outpaint que EXTIENDE el lienzo de un asset.

    Agranda el canvas del asset `left/top/right/bottom` pixeles por lado y genera el
    contenido nuevo de esa franja con `prompt`, conservando el asset original en su sitio.
    La mascara de "lo nuevo" la genera el grafo (`ImagePadForOutpaint`), no el caller: por
    eso, a diferencia de `comfyui_build_inpaint_asset_workflow`, NO se pasa `mask_image`.

    Args:
        input_image: nombre del archivo del asset a extender dentro de la carpeta input/
            del servidor ComfyUI (un fondo, parallax, card_art, splash...). Lo carga
            LoadImage. Subelo antes con POST /upload/image o copialo a ~/ComfyUI/input/.
            No puede estar vacio.
        prompt: que generar en la franja ampliada ("more forest and trees", "open sky and
            clouds", "continuing stone wall"). Describe la CONTINUACION de la escena, no un
            objeto centrado. No puede estar vacio.
        left: pixeles a extender por la izquierda. Se redondea al multiplo de 8 mas cercano
            y se clampa a [0, 16384]. keyword-only.
        right: pixeles a extender por la derecha (mismo tratamiento). keyword-only.
        top: pixeles a extender por arriba (mismo tratamiento). keyword-only.
        bottom: pixeles a extender por abajo (mismo tratamiento). keyword-only.
            Al menos uno de left/right/top/bottom debe ser > 0 tras redondear: sin
            extension no hay outpaint (para editar una region INTERIOR usa inpaint_asset).
        feather: pixeles de difuminado del borde entre el asset original y la franja nueva
            (`feathering` del nodo). Default 40 (el del nodo). Mas alto = transicion mas
            gradual pero invade mas la imagen original; no debe superar mucho la extension
            del lado mas pequeno. Se clampa a [0, 16384]. keyword-only.
        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 (la franja nueva).
            En outpaint suele ser 1.0 (por defecto): la zona ampliada parte de relleno y se
            genera por completo. Se clampa a [0.0, 1.0]. keyword-only.
        style: descriptor de estilo que mantiene la franja coherente con el asset y el set
            ("game background", "pixel art landscape", "dark fantasy scenery"). Pasa el
            MISMO style + checkpoint + (lora) que el resto del set. keyword-only.
        grow_mask: pixeles que se dilata ADEMAS el borde de la mascara hacia dentro de la
            imagen original (sobre el feathering del pad) para reforzar el blend si aparece
            costura. Default 0 (el feathering del pad suele bastar). Se clampa a [0, 64]
            (limite de VAEEncodeForInpaint). 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.
        negative: prompt negativo. None usa el negativo por defecto pensado para extension
            de escena (continuacion sin costura, sin repetir/espejar el original, un solo
            sujeto). keyword-only.
        steps, cfg, sampler_name, scheduler, filename_prefix: parametros de generacion.
            keyword-only.

    Returns:
        dict en API format listo para comfyui_submit_workflow: outpaint que extiende el
        lienzo por los lados pedidos y genera lo nuevo con '{prompt}, {style}, seamless
        extension...', conservando el asset original. Nodos: CheckpointLoaderSimple '4',
        LoadImage '10', ImagePadForOutpaint (id nuevo), VAEEncodeForInpaint '11',
        CLIPTextEncode '6'/'7', KSampler '3', VAEDecode '8', SaveImage '9' (+ LoraLoader si
        lora). El LoadImageMask de la base inpaint se elimina: la mascara la genera el pad.

    Raises:
        ValueError: si input_image o prompt estan vacios; si las cuatro extensiones quedan
            en 0 tras redondear (no hay nada que extender); o si la base no tiene los nodos
            esperados (propagado).
    """
    from ml.comfyui_build_inpaint_workflow import comfyui_build_inpaint_workflow

    if not input_image or not input_image.strip():
        raise ValueError(
            "comfyui_build_outpaint_asset_workflow: 'input_image' no puede estar vacio"
        )
    if not prompt or not prompt.strip():
        raise ValueError(
            "comfyui_build_outpaint_asset_workflow: 'prompt' no puede estar vacio"
        )

    input_image = input_image.strip()
    prompt = prompt.strip()

    left = min(16384, _round8(left))
    right = min(16384, _round8(right))
    top = min(16384, _round8(top))
    bottom = min(16384, _round8(bottom))
    if (left + right + top + bottom) == 0:
        raise ValueError(
            "comfyui_build_outpaint_asset_workflow: al menos uno de left/right/top/bottom "
            "debe ser > 0 (multiplo de 8) para extender el lienzo; sin extension no es "
            "outpaint. Para editar una region INTERIOR usa "
            "comfyui_build_inpaint_asset_workflow."
        )

    feather = max(0, min(16384, int(feather)))
    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 = _OUTPAINT_ASSET_NEGATIVE if negative is None else negative

    # Prompt de la franja nueva: pide CONTINUAR la escena (no un objeto centrado) y fundirla
    # con el original. A diferencia de inpaint_asset no se habla de "esta region" sino de
    # extender naturalmente hacia los lados ampliados.
    positive = (
        f"{prompt}, {style}, seamless extension of the scene, "
        "natural continuation, consistent lighting, matching art style, high detail"
    )

    # La base inpaint exige un nombre de mascara; usamos un placeholder porque el
    # LoadImageMask se elimina justo despues (outpaint genera su mascara con el pad).
    wf = comfyui_build_inpaint_workflow(
        checkpoint,
        input_image,
        "__outpaint_pad_placeholder__",
        positive,
        neg,
        denoise=denoise,
        steps=steps,
        cfg=cfg,
        seed=seed,
        sampler_name=sampler_name,
        scheduler=scheduler,
    )

    load_id = _find_class(wf, "LoadImage")
    enc_id = _find_class(wf, "VAEEncodeForInpaint")
    mask_id = _find_class(wf, "LoadImageMask")
    if load_id is None or enc_id is None:
        raise ValueError(
            "comfyui_build_outpaint_asset_workflow: la base inpaint no expone "
            "LoadImage/VAEEncodeForInpaint; no se puede armar el outpaint"
        )

    # Eliminar el LoadImageMask de la base: la mascara la genera ImagePadForOutpaint.
    if mask_id is not None:
        del wf[mask_id]

    # Insertar ImagePadForOutpaint entre LoadImage y VAEEncodeForInpaint. El pad extiende
    # el lienzo y emite (IMAGE extendida [salida 0], MASK de la franja nueva [salida 1]).
    pad_id = _new_id(wf)
    wf[pad_id] = {
        "class_type": "ImagePadForOutpaint",
        "inputs": {
            "image": [load_id, 0],
            "left": left,
            "top": top,
            "right": right,
            "bottom": bottom,
            "feathering": feather,
        },
    }

    # Reconectar VAEEncodeForInpaint a las dos salidas del pad: pixels <- IMAGE, mask <- MASK.
    enc = wf[enc_id]["inputs"]
    enc["pixels"] = [pad_id, 0]
    enc["mask"] = [pad_id, 1]
    enc["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 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_outpaint_asset_workflow(
        "seamless_00004_.png",
        "more forest and foliage extending to the right",
        right=256,
        feather=40,
        style="game background, top-down 2d",
        denoise=1.0,
        seed=7,
    )
    pad_id = next(
        nid for nid, n in wf.items() if n["class_type"] == "ImagePadForOutpaint"
    )
    enc_id = next(
        nid for nid, n in wf.items() if n["class_type"] == "VAEEncodeForInpaint"
    )
    print(
        json.dumps(
            {
                "nodes": list(wf),
                "classes": sorted({n["class_type"] for n in wf.values()}),
                "has_LoadImageMask": any(
                    n["class_type"] == "LoadImageMask" for n in wf.values()
                ),
                "pad_inputs": wf[pad_id]["inputs"],
                "enc_pixels": wf[enc_id]["inputs"]["pixels"],
                "enc_mask": wf[enc_id]["inputs"]["mask"],
                "enc_grow_mask_by": wf[enc_id]["inputs"]["grow_mask_by"],
                "denoise": wf["3"]["inputs"]["denoise"],
                "positive": wf["6"]["inputs"]["text"],
                "input_image": wf["10"]["inputs"]["image"],
                "filename_prefix": wf["9"]["inputs"]["filename_prefix"],
            },
            indent=2,
        )
    )
