import gradio as gr
import os
import time
from basic_pitch.inference import predict

# Configuration
OUTPUT_FOLDER = "midi_outputs"
CLEANUP_THRESHOLD_SECONDS = 60  # 1 minute
 
# Ensure the output directory exists
if not os.path.exists(OUTPUT_FOLDER):
    os.makedirs(OUTPUT_FOLDER) 
 
def cleanup_old_files():
    """Removes MIDI files older than 1 minute from the output folder."""
    now = time.time()
    for filename in os.listdir(OUTPUT_FOLDER):
        file_path = os.path.join(OUTPUT_FOLDER, filename)
        # Check if it's a file and if it's older than our threshold
        if os.path.isfile(file_path) and filename.endswith(".mid"):
            file_age = now - os.path.getmtime(file_path)
            if file_age > CLEANUP_THRESHOLD_SECONDS:
                try:
                    os.remove(file_path)
                    print(f"Cleaned up old file: {filename}")
                except Exception as e:
                    print(f"Failed to delete {filename}: {e}")

def convert_audio_to_midi(audio_filepath): 
    if audio_filepath is None:
        return None
    
    cleanup_old_files()
        
    try:
        # predict returns (model_output, midi_data, note_events)
        _, midi_data, _ = predict(audio_filepath)
        
        timestamp = int(time.time())
        filename = f"output_{timestamp}.mid"
        output_midi_path = os.path.join(OUTPUT_FOLDER, filename)
        
        midi_data.write(output_midi_path)
        
        return output_midi_path # Gradio handles making this a downloadable URL
    except Exception as e:
        print(f"Error: {e}")
        return None

# Build the Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎵 MIDI Converter
        Upload audio to convert to MIDI. Files are saved to the `/midi_outputs` folder.
        **Note:** Outputs are automatically deleted after 1 minute to save space.
        """
    )
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(type="filepath", label="Upload Audio File")
            convert_btn = gr.Button("Convert to MIDI", variant="primary")
            
        with gr.Column():
            midi_output = gr.File(label="Download MIDI (Expires in 60s)")

    convert_btn.click(
        fn=convert_audio_to_midi,
        inputs=audio_input,
        outputs=midi_output
    )

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