import gradio as gr
import asyncio
import wave
import os
import numpy as np
import nest_asyncio
from google import genai
from google.genai import types

# Enable nested asyncio (required for Gradio + asyncio)
nest_asyncio.apply()

# Configure Gemini API
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
client = genai.Client(http_options={'api_version': 'v1alpha'}, api_key=GOOGLE_API_KEY)

# Save PCM audio to WAV file
def save_wave_file(filename, pcm, channels=1, rate=24000, sample_width=2):
    with wave.open(filename, "wb") as wf:
        wf.setnchannels(channels)
        wf.setsampwidth(sample_width)
        wf.setframerate(rate)
        wf.writeframes(pcm)

# Async music generation function
async def generate_music(prompt, bpm, temperature):
    audio_chunks = []

    async def receive_audio(session):
        async for message in session.receive():
            chunk = message.server_content.audio_chunks[0].data
            audio_chunks.append(chunk)

    try:
        async with (
            client.aio.live.music.connect(model='models/lyria-realtime-exp') as session,
            asyncio.TaskGroup() as tg,
        ):
            tg.create_task(receive_audio(session))

            await session.set_weighted_prompts([
                types.WeightedPrompt(text=prompt, weight=1.0)
            ])
            await session.set_music_generation_config(
                types.LiveMusicGenerationConfig(bpm=int(bpm), temperature=float(temperature))
            )

            await session.play()
            await asyncio.sleep(5)  # Streaming duration
            await session.pause()

        all_pcm = b"".join(audio_chunks)

        # Save WAV file
        output_path = "generated_music.wav"
        save_wave_file(output_path, all_pcm)

        # Convert PCM to numpy array for audio playback
        audio_np = np.frombuffer(all_pcm, dtype=np.int16)
        return (24000, audio_np), output_path, "Music generated successfully!"

    except Exception as e:
        return None, None, f"❌ Error: {str(e)}"

# Wrapper for Gradio
def generate_music_gradio(prompt, bpm, temperature):
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(generate_music(prompt, bpm, temperature))

# Gradio UI
with gr.Blocks(title="Gemini Lyria Music Generator") as demo:
    gr.Markdown("## 🎶 Gemini Lyria Music Generator")

    with gr.Group():
        gr.Markdown("### 🎛 Input")
        with gr.Row():
            prompt_input = gr.Textbox(
                label="Music Style / Prompt",
                placeholder="e.g., ambient synth, minimal techno, classical piano"
            )
        with gr.Row():
            bpm_input = gr.Slider(label="BPM", minimum=60, maximum=180, value=90)
            temp_input = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=1.0)
        generate_btn = gr.Button("🎧 Generate Music")

    with gr.Group():
        gr.Markdown("### 🎵 Output")
        with gr.Row():
            output_audio = gr.Audio(label="Generated Audio", type="numpy")
            download_file = gr.File(label="Download WAV")
        status_output = gr.Textbox(label="Status", interactive=False)

    with gr.Group():
        gr.Markdown("### 🔁 Examples")
        examples = gr.Examples(
            examples=[
                ["minimal techno", 125, 1.0],
                ["classical piano in a rainy mood", 70, 0.9],
                ["ambient space drone", 90, 1.5],
                ["lo-fi chill beats", 80, 1.0],
                ["orchestral epic music", 110, 1.2],
            ],
            inputs=[prompt_input, bpm_input, temp_input]
        )

    generate_btn.click(
        fn=generate_music_gradio,
        inputs=[prompt_input, bpm_input, temp_input],
        outputs=[output_audio, download_file, status_output]
    )

demo.launch()