import gradio as gr
import replicate

def generate_music(genre, instruments, tempo, mood, duration, key_scale):
    prompt = f"Generate a {genre} track featuring {instruments}, in a {mood} mood, with a consistent melody throughout, with a tempo of {tempo} BPM, in the key of {key_scale}."
    try:
        # Assuming replicate.run() is a placeholder for actual music generation API call
        output = replicate.run(
            "bamartin1618/weather-channel:77dae70c23138abd20014ba686f62da780c5aed738e00085ce698eb5e344c65a",
            input={
                "top_k": 250,
                "top_p": 1,
                "prompt": prompt,
                "duration": duration,
                "temperature": 1,
                "continuation": False,
                "output_format": "wav",
                "continuation_start": 0,
                "multi_band_diffusion": False,
                "normalization_strategy": "loudness",
                "classifier_free_guidance": 6
            }
        )
    except Exception as e:  # Consider specifying more specific exception types
        return f'An error occurred during music generation: {str(e)}'

    if output:
        return f'<a href="{output}" target="_blank">Click here to listen to the generated music</a>'
    else:
        return 'Music generation failed or no URL returned.'

def get_audio_path(selection):
    audio_files = {
        "Smooth Jazz Evening": "smooth_jazz_evening.wav",
        "Uplifting Soft Rock": "uplifting_soft_rock.wav",
        "Mellow Jazz Morning": "mellow_jazz_morning.wav",
        "Soft Rock Serenade": "soft_rock_serenade.wav",
    }
    # Check if the selection exists in the dictionary to avoid KeyError
    if selection in audio_files:
        return audio_files[selection]
    else:
        # Handle unexpected selection
        return 'Error: Selected audio file does not exist.'

with gr.Blocks() as app:
    with gr.Tab("Generate Music"):
        genre = gr.Dropdown(["Smooth Jazz", "Soft Rock"], label="Genre")
        instruments = gr.CheckboxGroup(["Piano", "Saxophone", "Guitar", "Strings"], label="Instruments")
        tempo = gr.Slider(minimum=60, maximum=180, step=5, value=120, label="Tempo (BPM)")
        mood = gr.Dropdown(["Relaxing", "Cheerful"], label="Mood")
        duration = gr.Slider(minimum=1, maximum=20, step=1, value=15, label="Duration (Seconds)")
        key_scale = gr.Dropdown(["C Major", "A Minor", "G Major", "E Minor"], label="Key & Scale")
        generate_btn = gr.Button("Generate")
        generate_btn.click(
            generate_music, 
            inputs=[genre, instruments, tempo, mood, duration, key_scale], 
            outputs=gr.HTML()
        )

    with gr.Tab("Samples"):
        audio_selection = gr.Radio(
            ["Smooth Jazz Evening", "Uplifting Soft Rock", "Mellow Jazz Morning", "Soft Rock Serenade"],
            label="Select Music Sample"
        )
        play_audio_btn = gr.Button("Play Audio")
        preloaded_audio_output = gr.Audio(label="Preloaded Audio")
        play_audio_btn.click(
            get_audio_path, 
            inputs=audio_selection, 
            outputs=preloaded_audio_output
        )

app.launch()