"""
Game Audio Generator via ComfyUI API
Aufruf: python generate_via_api.py
Voraussetzung: ComfyUI Container läuft auf localhost:8188
"""
import json
import requests
import time
import urllib.request
import os

COMFY_URL = "http://localhost:8188"
OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "output", "audio")
os.makedirs(OUTPUT_DIR, exist_ok=True)

TRACKS = [
    # (typ, prompt, dauer, dateiname)
    ("music", "dark fantasy main menu music, epic orchestral, dramatic, mysterious, dark medieval atmosphere, slow build, no vocals", 40, "menu_music"),
    ("sfx",   "fast sword swing whoosh attack sound effect", 1, "attack_autoattack"),
    ("sfx",   "heavy powerful sword strike crushing impact metal clang", 2, "attack_heavy_strike"),
    ("sfx",   "massive ground slam earthquake shockwave stone cracking impact", 2, "attack_tektonischer_schlag"),
    ("sfx",   "warrior battle cry rage grunt determination short", 2, "attack_durchbeissen"),
    ("sfx",   "blade slash wet blood hit dark melee attack", 2, "attack_blutrausch"),
    ("sfx",   "fast spinning whirlwind sword multiple swings whoosh", 2, "attack_wirbelwind"),
    ("sfx",   "chains rattling metal binding restraint whoosh throw", 2, "attack_zornfesseln"),
    ("sfx",   "berserker rage roar battle fury intense warrior scream", 3, "attack_berserker"),
]

def build_prompt(typ, prompt_text, duration, filename):
    node_type = "MusicGen" if typ == "music" else "AudioGen"
    widgets = [prompt_text, duration, "facebook/musicgen-small", 0] if typ == "music" else [prompt_text, duration, 0]
    return {
        "prompt": {
            "1": {
                "class_type": node_type,
                "inputs": {
                    "prompt": prompt_text,
                    "duration": duration,
                    **({"model": "facebook/musicgen-small", "seed": 0} if typ == "music" else {"seed": 0})
                }
            }
        },
        "extra_data": {"extra_pnginfo": {"filename_hint": filename}}
    }

def queue_prompt(payload):
    r = requests.post(f"{COMFY_URL}/api/prompt", json=payload)
    r.raise_for_status()
    return r.json()["prompt_id"]

def wait_for_completion(prompt_id, timeout=300):
    start = time.time()
    while time.time() - start < timeout:
        r = requests.get(f"{COMFY_URL}/api/history/{prompt_id}")
        data = r.json()
        if prompt_id in data:
            return data[prompt_id]
        time.sleep(2)
    raise TimeoutError(f"Timeout nach {timeout}s")

def get_output_files():
    r = requests.get(f"{COMFY_URL}/api/history")
    return r.json()

if __name__ == "__main__":
    print(f"Verbinde mit ComfyUI: {COMFY_URL}")
    try:
        requests.get(f"{COMFY_URL}/api/system_stats").raise_for_status()
        print("ComfyUI erreichbar\n")
    except Exception as e:
        print(f"ComfyUI nicht erreichbar: {e}")
        exit(1)

    for typ, prompt, duration, filename in TRACKS:
        print(f"[{typ.upper()}] {filename}: {prompt[:55]}...")
        payload = build_prompt(typ, prompt, duration, filename)
        try:
            prompt_id = queue_prompt(payload)
            print(f"  queued: {prompt_id}")
            result = wait_for_completion(prompt_id)
            print(f"  fertig!")
        except Exception as e:
            print(f"  Fehler: {e}")

    print(f"\nAlle Jobs abgeschickt. Output unter: {OUTPUT_DIR}")
    print(f"Oder direkt in ComfyUI unter output/audio/")