"""
Game Audio Generator - nutzt AudioCraft direkt
Aufruf: python generate_audio.py
"""
import torch
import os
import time

OUTPUT_DIR = "/app/ComfyUI/output/audio"
os.makedirs(OUTPUT_DIR, exist_ok=True)

TRACKS = [
    # Menü Musik
    ("music", "dark fantasy main menu music, epic orchestral, dramatic, mysterious, dark medieval atmosphere, slow build, no vocals", 40, "menu_music"),

    # Krieger Attacken
    ("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 generate_music(prompt, duration, filename):
    from audiocraft.models import MusicGen
    from audiocraft.data.audio import audio_write

    print(f"  [MusicGen] {filename}: {prompt[:50]}...")
    model = MusicGen.get_pretrained("facebook/musicgen-small")
    model.set_generation_params(duration=duration)
    wav = model.generate([prompt])
    out = os.path.join(OUTPUT_DIR, filename)
    audio_write(out, wav[0].cpu(), model.sample_rate, strategy="loudness")
    print(f"  -> gespeichert: {out}.wav")

def generate_sfx(prompt, duration, filename):
    from audiocraft.models import AudioGen
    from audiocraft.data.audio import audio_write

    print(f"  [AudioGen] {filename}: {prompt[:50]}...")
    model = AudioGen.get_pretrained("facebook/audiogen-medium")
    model.set_generation_params(duration=duration)
    wav = model.generate([prompt])
    out = os.path.join(OUTPUT_DIR, filename)
    audio_write(out, wav[0].cpu(), model.sample_rate, strategy="loudness")
    print(f"  -> gespeichert: {out}.wav")

if __name__ == "__main__":
    print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
    print(f"Output: {OUTPUT_DIR}\n")

    music_model = None
    sfx_model = None

    for typ, prompt, duration, filename in TRACKS:
        print(f"\n[{typ.upper()}] {filename}")
        if typ == "music":
            generate_music(prompt, duration, filename)
        else:
            generate_sfx(prompt, duration, filename)

    print("\nFertig! Alle Audio-Dateien wurden generiert.")