import gradio as gr
import os
from model import Jamify

# Initialize the Jamify model once
print("Initializing Jamify model...")
jamify_model = Jamify()
print("Jamify model ready.")

def generate_song(reference_audio, lyrics_file, style_prompt, duration):
    # We need to save the uploaded files to temporary paths to pass to the model
    ref_audio_path = reference_audio.name if reference_audio else None
    lyrics_path = lyrics_file.name

    # The model expects paths, so we write the prompt to a temp file if needed
    # (This part of the model could be improved to accept the string directly)

    output_path = jamify_model.predict(
        reference_audio_path=ref_audio_path,
        lyrics_json_path=lyrics_path,
        style_prompt=style_prompt,
        duration_sec=duration
    )
    
    return output_path

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Jamify: Music Generation from Lyrics and Style")
    gr.Markdown("Provide your lyrics, a style reference (either an audio file or a text prompt), and a desired duration to generate a song.")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Inputs")
            lyrics_file = gr.File(label="Lyrics File (.json)", type="filepath")
            duration_slider = gr.Slider(minimum=5, maximum=180, value=30, step=1, label="Duration (seconds)")
            
            with gr.Tab("Style from Audio"):
                reference_audio = gr.File(label="Reference Audio (.mp3, .wav)", type="filepath")
            with gr.Tab("Style from Text"):
                style_prompt = gr.Textbox(label="Style Prompt", lines=3, placeholder="e.g., A high-energy electronic dance track with a strong bassline and euphoric synths.")
            
            generate_button = gr.Button("Generate Song", variant="primary")
            
        with gr.Column():
            gr.Markdown("### Output")
            output_audio = gr.Audio(label="Generated Song")

    generate_button.click(
        fn=generate_song,
        inputs=[reference_audio, lyrics_file, style_prompt, duration_slider],
        outputs=output_audio,
        api_name="generate_song"
    )

    gr.Markdown("### Example Usage")
    gr.Examples(
        examples=[
            [None, "jamify/inputs/Jade Bird - Avalanche.json", "A sad, slow, acoustic country song", 30],
            ["jamify/inputs/Rizzle Kicks, Rachel Chinouriri - Follow Excitement!.mp3", "jamify/inputs/Rizzle Kicks, Rachel Chinouriri - Follow Excitement!.json", "", 45],
        ],
        inputs=[reference_audio, lyrics_file, style_prompt, duration_slider],
        outputs=output_audio,
        fn=generate_song,
        cache_examples=True
    )

demo.queue().launch()