import gradio as gr
import torch
from transformers import AutoProcessor, MusicgenForConditionalGeneration
import numpy as np

# Wir nutzen das melody-Modell, da es Audio-Inputs verarbeiten kann
MODEL_ID = "facebook/musicgen-melody"

print("Lade MusicGen-Modell... Bitte warten.")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = MusicgenForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=torch.float16)
model.to("cpu") # Erzwinge CPU, da die meisten Free-Spaces keine GPU haben
print("MusicGen ist startklar!")

def generate_game_music(prompt: str, audio_input):
    if not prompt.strip():
        raise gr.Error("Bitte gib zuerst einen Musik-Stil ein.")
    
    try:
        # Inputs vorbereiten
        input_kwargs = {"text": [prompt], "padding": True, "return_tensors": "pt"}
        
        # Falls Audio hochgeladen wurde, als Melodie-Referenz nutzen
        if audio_input is not None:
            sr, audio = audio_input
            # Audio in das richtige Format bringen
            audio = torch.tensor(audio).float().t().unsqueeze(0)
            input_kwargs["audio_values"] = audio
            input_kwargs["sampling_rate"] = sr

        inputs = processor(**input_kwargs)

        with torch.no_grad():
            audio_values = model.generate(**inputs, max_new_tokens=256)
        
        # Formatierung für Gradio: (Samples, Channels)
        audio_data = audio_values[0].cpu().numpy().transpose(1, 0)
        sampling_rate = model.config.audio_encoder.sampling_rate
        audio_data = np.clip(audio_data, -1.0, 1.0)
        
        return (sampling_rate, audio_data)

    except Exception as e:
        raise gr.Error(f"Fehler bei der Generierung: {str(e)}")

# Interface mit gr.Blocks
with gr.Blocks(title="🕹️ GamerJam AI Music Studio") as demo:
    gr.Markdown("# 🕹️ GamerJam.de AI Music Studio")
    gr.Markdown("Generiere Soundtracks. Du kannst optional ein Audio-Sample als Stil-Vorlage hochladen.")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Musikstil beschreiben", placeholder="e.g. 16-bit retro arcade loop")
            audio_input = gr.Audio(label="Optional: Referenz-Audio hochladen", type="numpy")
            submit_btn = gr.Button("Soundtrack generieren")
        
        with gr.Column():
            output_audio = gr.Audio(label="Generierter Game-Soundtrack")

    submit_btn.click(
        fn=generate_game_music,
        inputs=[prompt_input, audio_input],
        outputs=output_audio
    )

if __name__ == "__main__":
    demo.launch()