import gradio as gr
import numpy as np
from mega import Mega
import os
import glob

# Load files
spectrogram_path = "ui/temp.npy"
generated_song_path = "ui/temp.wav"

def rate_song(user_id, rating, model_name, song_name, similarity):
    # Log in to Mega
    mega = Mega()
    mega_user_name = os.environ.get('MEGA_USERNAME')
    mega_password = os.environ.get('MEGA_PASSWORD')
    m = mega.login(mega_user_name, mega_password)
    
    # Construct file names and paths for uploading
    dynamic_song_name = f"{user_id}_{model_name}_{song_name}_{similarity}_{rating}.wav"
    dynamic_spec_name = f"{user_id}_{model_name}_{song_name}_{similarity}_{rating}.npy"
    folder = m.find('orpheus_data')

    # Upload files
    m.upload(generated_song_path, folder[0], dest_filename=dynamic_song_name)
    m.upload(spectrogram_path, folder[0], dest_filename=dynamic_spec_name)
    
    return "Files uploaded successfully!"

with gr.Blocks() as rating_demo:
    song_name = gr.Markdown("# Original Song")
    gr.Audio(generated_song_path, label="Generated Song", format="wav")
    rating_slider = gr.Slider(minimum=0, maximum=10, value=3, label="Rating")
    submit_rating_button = gr.Button("Submit Rating")

    # Outputs
    upload_status = gr.Textbox(label="Upload Status")

    # Collect session state and submit rating
    submit_rating_button.click(fn=rate_song, inputs=["user_id", rating_slider, "model_name", "song_name", "similarity"], outputs=upload_status)

rating_demo.launch()