"""Construye un workflow ComfyUI de texto->audio (ACE-Step) 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).

El grafo usa los nodos de audio NATIVOS de ComfyUI 0.26.0 para el modelo
ACE-Step (abierto, Apache 2.0): CheckpointLoaderSimple ->
TextEncodeAceStepAudio (tags + lyrics) -> EmptyAceStepLatentAudio ->
ModelSamplingSD3 -> KSampler -> VAEDecodeAudio -> SaveAudio. El negative se
construye con ConditioningZeroOut sobre el positive (patron oficial de ACE-Step).

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


def comfyui_build_audio_workflow(
    ckpt_name: str,
    prompt: str,
    *,
    lyrics: str = "",
    seconds: float = 10.0,
    seed: int = 0,
    steps: int = 50,
    cfg: float = 5.0,
    sampler_name: str = "euler",
    scheduler: str = "simple",
    shift: float = 5.0,
    lyrics_strength: float = 1.0,
    filename_prefix: str = "audio/comfy_audio",
) -> dict:
    """Construye el dict del workflow texto->audio para ACE-Step.

    Cadena de nodos: CheckpointLoaderSimple -> TextEncodeAceStepAudio (positivo)
    + ConditioningZeroOut (negativo) + EmptyAceStepLatentAudio -> ModelSamplingSD3
    -> KSampler -> VAEDecodeAudio -> SaveAudio. SaveAudio escribe un .flac en la
    carpeta output/<filename_prefix> del servidor ComfyUI.

    Args:
        ckpt_name: nombre del checkpoint ACE-Step tal como lo ve el servidor
            (ej. "AUDIO_ace_step_v1_3.5b.safetensors"). Debe estar entre los que
            devuelve comfyui_object_info en CheckpointLoaderSimple.
        prompt: descripcion del sonido o estilo musical (va al campo "tags" de
            TextEncodeAceStepAudio). Ej. "8-bit coin pickup sound, retro game".
        lyrics: letra cantada para musica con voz. Vacio "" para SFX o musica
            instrumental.
        seconds: duracion del audio en segundos (min 1.0). Controla el tamano
            del latente via EmptyAceStepLatentAudio.
        seed: semilla del KSampler (cambia para variar el resultado).
        steps: pasos de sampling del KSampler (50 recomendado para ACE-Step).
        cfg: classifier-free guidance scale (5.0 recomendado para ACE-Step).
        sampler_name: nombre del sampler (ej. "euler").
        scheduler: scheduler del sampler (ej. "simple").
        shift: shift del ModelSamplingSD3 aplicado al MODEL antes del sampling
            (5.0 recomendado para ACE-Step). Mejora la coherencia temporal.
        lyrics_strength: fuerza del condicionamiento de la letra (1.0 por
            defecto; sin efecto practico cuando lyrics esta vacio).
        filename_prefix: prefijo del .flac generado por SaveAudio en output/.

    Returns:
        dict en API format listo para comfyui_submit_workflow. Las claves son
        node_ids ("3".."11") y cada valor tiene class_type + inputs.
    """
    return {
        "4": {
            "class_type": "CheckpointLoaderSimple",
            "inputs": {"ckpt_name": ckpt_name},
        },
        "6": {
            "class_type": "TextEncodeAceStepAudio",
            "inputs": {
                "clip": ["4", 1],
                "tags": prompt,
                "lyrics": lyrics,
                "lyrics_strength": lyrics_strength,
            },
        },
        "10": {
            "class_type": "ConditioningZeroOut",
            "inputs": {"conditioning": ["6", 0]},
        },
        "5": {
            "class_type": "EmptyAceStepLatentAudio",
            "inputs": {"seconds": seconds, "batch_size": 1},
        },
        "11": {
            "class_type": "ModelSamplingSD3",
            "inputs": {"model": ["4", 0], "shift": shift},
        },
        "3": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed,
                "steps": steps,
                "cfg": cfg,
                "sampler_name": sampler_name,
                "scheduler": scheduler,
                "denoise": 1.0,
                "model": ["11", 0],
                "positive": ["6", 0],
                "negative": ["10", 0],
                "latent_image": ["5", 0],
            },
        },
        "8": {
            "class_type": "VAEDecodeAudio",
            "inputs": {"samples": ["3", 0], "vae": ["4", 2]},
        },
        "9": {
            "class_type": "SaveAudio",
            "inputs": {"filename_prefix": filename_prefix, "audio": ["8", 0]},
        },
    }


if __name__ == "__main__":
    import json

    wf = comfyui_build_audio_workflow(
        ckpt_name="AUDIO_ace_step_v1_3.5b.safetensors",
        prompt="8-bit coin pickup sound, retro game, short",
        seconds=4.0,
        seed=42,
    )
    print(json.dumps(wf, indent=2))
