fn_registry/python/functions/ml/comfyui_append_styles.py

"""comfyui_append_styles — merge+dedup de estilos nuevos sobre el styles.json de WAS.

El selector de estilos de ComfyUI (nodos WAS `Prompt Styles Selector` /
`Prompt Multiple Styles Selector`) lee de
`~/ComfyUI/custom_nodes/was-node-suite-comfyui/styles.json`, un dict cuyo formato exacto es:

    { "NombreEstilo": {"prompt": "modificadores de estilo", "negative_prompt": "..."}, ... }

El selector múltiple CONCATENA los `prompt` de los estilos elegidos, por lo que cada `prompt`
debe contener MODIFICADORES de estilo (no la descripción del sujeto) y NO el placeholder
`{prompt}`.

Esta función fusiona un dict de estilos nuevos sobre el archivo existente de forma SEGURA y NO
destructiva:

  - Hace un backup con timestamp del styles.json antes de tocarlo (nunca sobrescribe sin copia).
  - Preserva TODOS los estilos existentes (dedup por nombre: los existentes ganan salvo
    `overwrite=True`).
  - Valida cada entrada nueva: debe ser un dict con `prompt` no vacío. Si falta `negative_prompt`
    se rellena con un negativo por defecto razonable; las entradas inválidas se descartan
    (reportadas, no abortan el merge).
  - Escribe el resultado de forma atómica (a un .tmp y `os.replace`).

Devuelve un resumen con conteos (antes/después, añadidos, duplicados saltados, inválidos) para
que el caller verifique el efecto sin volver a leer el archivo.

Impura: lee y escribe disco. No usa red. No mata procesos. No borra el original (sólo backup +
reemplazo atómico).
"""
from __future__ import annotations

import json
import os
import shutil
import time

# Negativo por defecto cuando un estilo nuevo no trae `negative_prompt`. Sobrio y SFW,
# alineado con el estilo de los negativos que ya viven en el styles.json de WAS.
DEFAULT_NEGATIVE = (
    "ugly, deformed, noisy, blurry, low quality, distorted, disfigured, "
    "bad anatomy, watermark, signature, text, NSFW"
)

DEFAULT_STYLES_PATH = os.path.join(
    os.path.expanduser("~"),
    "ComfyUI",
    "custom_nodes",
    "was-node-suite-comfyui",
    "styles.json",
)


def _validate_entry(value: object) -> dict | None:
    """Normaliza una entrada de estilo. Devuelve el dict válido o None si es inválida.

    Una entrada válida es un dict con `prompt` (str no vacío). `negative_prompt` se rellena
    con `DEFAULT_NEGATIVE` si falta o está vacío. Campos extra se descartan (el formato WAS
    sólo usa `prompt` y `negative_prompt`).
    """
    if not isinstance(value, dict):
        return None
    prompt = value.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        return None
    neg = value.get("negative_prompt")
    if not isinstance(neg, str) or not neg.strip():
        neg = DEFAULT_NEGATIVE
    return {"prompt": prompt.strip(), "negative_prompt": neg.strip()}


def comfyui_append_styles(
    new_styles: dict,
    styles_path: str = DEFAULT_STYLES_PATH,
    overwrite: bool = False,
    backup: bool = True,
    dry_run: bool = False,
) -> dict:
    """Fusiona `new_styles` sobre el styles.json de WAS preservando los existentes.

    Args:
        new_styles: dict {nombre: {"prompt": str, "negative_prompt": str}} de estilos a añadir.
            Las entradas inválidas (sin `prompt`) se descartan; las que no traen
            `negative_prompt` reciben uno por defecto.
        styles_path: ruta del styles.json. Default: el de la instalación WAS del usuario.
        overwrite: si False (default), un nombre que ya existe en el archivo NO se pisa (se
            cuenta como duplicado saltado). Si True, los nuevos pisan a los existentes.
        backup: si True (default), copia el archivo a `<path>.bak.<epoch>` antes de escribir.
        dry_run: si True, calcula el merge y los conteos pero NO escribe nada (ni backup).

    Returns:
        dict resumen: {
            "styles_path", "backup_path", "total_before", "total_after",
            "added": [nombres añadidos], "overwritten": [nombres pisados],
            "skipped_existing": [nombres saltados por existir y overwrite=False],
            "invalid": [nombres descartados por inválidos], "dry_run": bool
        }

    Raises:
        FileNotFoundError: si `styles_path` no existe (no se crea de cero para no enmascarar
            una instalación rota; el caller debe asegurar que el archivo está).
        ValueError: si `new_styles` no es un dict o el archivo existente no contiene un dict.
    """
    if not isinstance(new_styles, dict):
        raise ValueError("comfyui_append_styles: new_styles debe ser un dict")
    if not os.path.isfile(styles_path):
        raise FileNotFoundError(f"comfyui_append_styles: no existe styles.json en {styles_path!r}")

    with open(styles_path, "r", encoding="utf-8") as fh:
        existing = json.load(fh)
    if not isinstance(existing, dict):
        raise ValueError(
            f"comfyui_append_styles: el styles.json en {styles_path!r} no es un dict de estilos"
        )

    total_before = len(existing)
    merged = dict(existing)  # copia: no mutar el cargado hasta validar todo

    added: list[str] = []
    overwritten: list[str] = []
    skipped_existing: list[str] = []
    invalid: list[str] = []

    for name, value in new_styles.items():
        norm = _validate_entry(value)
        if norm is None:
            invalid.append(str(name))
            continue
        if name in existing:
            if overwrite:
                merged[name] = norm
                overwritten.append(name)
            else:
                skipped_existing.append(name)
            continue
        merged[name] = norm
        added.append(name)

    backup_path = ""
    if not dry_run:
        if backup:
            backup_path = f"{styles_path}.bak.{int(time.time())}"
            shutil.copy2(styles_path, backup_path)
        # Escritura atómica: escribir a .tmp en el mismo dir y reemplazar.
        tmp_path = f"{styles_path}.tmp.{os.getpid()}"
        with open(tmp_path, "w", encoding="utf-8") as fh:
            json.dump(merged, fh, ensure_ascii=False, indent=4)
        os.replace(tmp_path, styles_path)

    return {
        "styles_path": styles_path,
        "backup_path": backup_path,
        "total_before": total_before,
        "total_after": len(merged),
        "added": added,
        "overwritten": overwritten,
        "skipped_existing": skipped_existing,
        "invalid": invalid,
        "dry_run": dry_run,
    }


if __name__ == "__main__":
    import sys

    # CLI de conveniencia: lee un dict de estilos JSON de stdin (o de un archivo dado como
    # primer arg) y lo fusiona. Con --dry-run no escribe. Imprime el resumen como JSON.
    args = sys.argv[1:]
    dry = "--dry-run" in args
    over = "--overwrite" in args
    path_args = [a for a in args if not a.startswith("--")]
    if path_args:
        with open(path_args[0], "r", encoding="utf-8") as fh:
            payload = json.load(fh)
    else:
        payload = json.load(sys.stdin)
    res = comfyui_append_styles(payload, overwrite=over, dry_run=dry)
    print(json.dumps(res, ensure_ascii=False, indent=2))