"""Construye un workflow ComfyUI para Flux (schnell o dev) en "API format".

API format: cada clave es un node_id (string); cada nodo tiene class_type +
inputs. Las conexiones entre nodos son listas [node_id, output_index]. Este es
el formato que acepta POST /prompt, distinto del formato de la UI (graph con
links explicitos).

Flux NO se carga como un checkpoint clasico (no CheckpointLoaderSimple). El
modelo de difusion se carga con UNETLoader; los dos text encoders (clip_l + t5xxl)
con DualCLIPLoader (type="flux"); el VAE con VAELoader. El muestreo usa el camino
"custom advanced" (RandomNoise -> KSamplerSelect + BasicScheduler -> BasicGuider
-> SamplerCustomAdvanced), que es el patron canonico de los ejemplos oficiales de
Flux y el que produce resultados estables con los modelos fp8 distilados.

Diferencias schnell vs dev:
- schnell: modelo distilado, ~4 pasos, sin FluxGuidance (la guia va fija dentro
  del modelo). Rapido. El conditioning del prompt va directo a BasicGuider.
- dev: ~20 pasos, el conditioning pasa antes por FluxGuidance (guidance ~3.5),
  que sube la adherencia al prompt a costa de tiempo. Mejor calidad.

Flux ignora el prompt negativo, por eso solo se codifica el positivo.

Funcion pura: sin red, sin I/O. Determinista para los mismos argumentos. La
validacion de existencia de modelos en disco se hace pasando `available` (mapa
de modelos que el servidor expone via /object_info); recibir ese mapa como
argumento no rompe la pureza (el caller hace la unica peticion de red).
"""

# Modelos por defecto para cada variante (nombres tal como los expone el
# servidor ComfyUI en /object_info; verificados contra UNETLoader.unet_name,
# DualCLIPLoader.clip_name1/2 y VAELoader.vae_name).
_DEFAULT_UNET = {
    "schnell": "IMG_flux1-schnell-fp8-e4m3fn.safetensors",
    "dev": "IMG_flux1-dev-fp8-e4m3fn.safetensors",
}
_DEFAULT_STEPS = {"schnell": 4, "dev": 20}

# Carpeta destino por rol de modelo, para mensajes de error utiles. ComfyUI
# acepta tanto la carpeta "diffusion_models" (moderna) como "unet" (legacy) para
# el UNET; los text encoders en "text_encoders" o "clip"; el VAE en "vae".
_MODEL_DIRS = {
    "unet": "models/diffusion_models/ (o models/unet/)",
    "clip": "models/text_encoders/ (o models/clip/)",
    "vae": "models/vae/",
}


def comfyui_build_flux_workflow(
    prompt: str,
    *,
    variant: str = "schnell",
    width: int = 1024,
    height: int = 1024,
    steps: int | None = None,
    guidance: float = 3.5,
    seed: int = 0,
    unet_name: str | None = None,
    clip_l_name: str = "clip_l.safetensors",
    t5xxl_name: str = "t5xxl_fp8_e4m3fn_scaled.safetensors",
    vae_name: str = "ae.safetensors",
    weight_dtype: str = "default",
    sampler_name: str = "euler",
    scheduler: str = "simple",
    filename_prefix: str = "flux",
    available: dict | None = None,
) -> dict:
    """Construye el dict del workflow Flux (schnell o dev) en API format.

    Cadena de nodos: UNETLoader + DualCLIPLoader + VAELoader -> CLIPTextEncode
    [-> FluxGuidance solo en dev] -> BasicGuider; RandomNoise + KSamplerSelect +
    BasicScheduler + EmptyLatentImage -> SamplerCustomAdvanced -> VAEDecode ->
    SaveImage.

    Args:
        prompt: prompt positivo (lo que se quiere ver). Flux ignora el negativo.
        variant: "schnell" (rapido, ~4 pasos, sin FluxGuidance) o "dev"
            (~20 pasos, con FluxGuidance). keyword-only.
        width: ancho del latente/imagen en px (multiplo de 8). keyword-only.
        height: alto del latente/imagen en px (multiplo de 8). keyword-only.
        steps: pasos de sampling. Si None, default por variante (schnell=4,
            dev=20). keyword-only.
        guidance: valor de FluxGuidance. Solo se aplica en variant="dev"; en
            schnell se ignora (el modelo distilado lleva la guia fija).
            keyword-only.
        seed: semilla de RandomNoise (cambia para variar la imagen). keyword-only.
        unet_name: nombre del modelo de difusion en UNETLoader. Si None, default
            por variante. keyword-only.
        clip_l_name: nombre del encoder CLIP-L en DualCLIPLoader. keyword-only.
        t5xxl_name: nombre del encoder T5-XXL en DualCLIPLoader. keyword-only.
        vae_name: nombre del VAE en VAELoader. keyword-only.
        weight_dtype: dtype de los pesos del UNET ("default", "fp8_e4m3fn",
            "fp8_e4m3fn_fast", "fp8_e5m2"). keyword-only.
        sampler_name: sampler para KSamplerSelect (ej. "euler"). keyword-only.
        scheduler: scheduler para BasicScheduler (ej. "simple"). keyword-only.
        filename_prefix: prefijo del PNG generado por SaveImage en output/.
            keyword-only.
        available: mapa opcional para validar que los modelos existen en el
            servidor, con claves opcionales "unet", "clip", "vae", cada una una
            lista de nombres disponibles (tal como /object_info los expone). Si
            se pasa y algun modelo elegido no esta en su lista, se lanza
            FileNotFoundError indicando que falta y en que carpeta colocarlo.
            Si es None (default), no se valida disco. keyword-only.

    Returns:
        dict en API format listo para comfyui_submit_workflow. Las claves son
        node_ids string y cada valor tiene class_type + inputs.

    Raises:
        ValueError: si variant no es "schnell" ni "dev".
        FileNotFoundError: si `available` se pasa y algun modelo (unet/clip/vae)
            no esta disponible en el servidor; el mensaje lista los que faltan y
            la carpeta donde colocarlos. La funcion NO crashea de forma opaca:
            falla con un error claro y accionable antes de enviar nada a la GPU.
    """
    if variant not in ("schnell", "dev"):
        raise ValueError(
            f"comfyui_build_flux_workflow: variant '{variant}' invalido; "
            f"usa 'schnell' o 'dev'"
        )

    unet = unet_name or _DEFAULT_UNET[variant]
    n_steps = steps if steps is not None else _DEFAULT_STEPS[variant]

    # Error path: validar contra los modelos que expone el servidor, si el caller
    # nos pasa el mapa. Pura (no toca disco; recibe las listas ya obtenidas).
    if available is not None:
        missing = []
        checks = (
            ("unet", unet, available.get("unet")),
            ("clip", clip_l_name, available.get("clip")),
            ("clip", t5xxl_name, available.get("clip")),
            ("vae", vae_name, available.get("vae")),
        )
        for role, name, names in checks:
            if names is not None and name not in names:
                missing.append(
                    f"  - '{name}' (rol {role}) no esta en el servidor; "
                    f"colocalo en {_MODEL_DIRS[role]}"
                )
        if missing:
            raise FileNotFoundError(
                "comfyui_build_flux_workflow: faltan modelos Flux en el "
                "servidor:\n" + "\n".join(missing)
            )

    # Loaders (Flux no usa CheckpointLoaderSimple).
    workflow: dict = {
        "12": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": unet, "weight_dtype": weight_dtype},
        },
        "11": {
            "class_type": "DualCLIPLoader",
            "inputs": {
                "clip_name1": t5xxl_name,
                "clip_name2": clip_l_name,
                "type": "flux",
            },
        },
        "10": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": vae_name},
        },
        "5": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": width, "height": height, "batch_size": 1},
        },
        "6": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": prompt, "clip": ["11", 0]},
        },
    }

    # Conditioning hacia BasicGuider. En dev pasa por FluxGuidance; en schnell va
    # directo (el modelo distilado no usa guidance externo).
    if variant == "dev":
        workflow["21"] = {
            "class_type": "FluxGuidance",
            "inputs": {"conditioning": ["6", 0], "guidance": guidance},
        }
        guider_cond = ["21", 0]
    else:
        guider_cond = ["6", 0]

    workflow.update(
        {
            "25": {
                "class_type": "RandomNoise",
                "inputs": {"noise_seed": seed},
            },
            "16": {
                "class_type": "KSamplerSelect",
                "inputs": {"sampler_name": sampler_name},
            },
            "17": {
                "class_type": "BasicScheduler",
                "inputs": {
                    "model": ["12", 0],
                    "scheduler": scheduler,
                    "steps": n_steps,
                    "denoise": 1.0,
                },
            },
            "22": {
                "class_type": "BasicGuider",
                "inputs": {"model": ["12", 0], "conditioning": guider_cond},
            },
            "13": {
                "class_type": "SamplerCustomAdvanced",
                "inputs": {
                    "noise": ["25", 0],
                    "guider": ["22", 0],
                    "sampler": ["16", 0],
                    "sigmas": ["17", 0],
                    "latent_image": ["5", 0],
                },
            },
            "8": {
                "class_type": "VAEDecode",
                "inputs": {"samples": ["13", 0], "vae": ["10", 0]},
            },
            "9": {
                "class_type": "SaveImage",
                "inputs": {"filename_prefix": filename_prefix, "images": ["8", 0]},
            },
        }
    )
    return workflow


if __name__ == "__main__":
    import json

    wf = comfyui_build_flux_workflow(
        "a red apple on a wooden table, sharp focus, studio light",
        variant="schnell",
        width=1024,
        height=1024,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
