import gradio as gr
import os
import random
import re
from scipy.io.wavfile import write
from audio_separator.separator import Separator

# Initialize the Separator
separator = Separator()

def separate_audio(audio_file):
    if audio_file is None:
        raise ValueError("Please upload an audio file.")

    # Load the separator model
    separator.load_model()

    # Generate a random file name for the output
    random_id = str(random.randint(10000, 99999))
    output_file = f"outputs/{random_id}.wav"
    
    # Create the output directory if it doesn't exist
    os.makedirs("outputs", exist_ok=True)
    
    # Save the uploaded audio file to the output directory
    write(output_file, audio_file[0], audio_file[1])

    # Perform the separation
    output_files = separator.separate(output_file)
    
    # List the separated files
    files_list = []
    for file in os.listdir("outputs"):
        if re.search(random_id, file):
            files_list.append(os.path.join("outputs", file))
    
    # Ensure there are at least two separated stems
    if len(files_list) < 2:
        raise ValueError("Error: Separation did not produce enough stems.")
        
    stem1_file = files_list[0]
    stem2_file = files_list[1]

    return stem1_file, stem2_file

# Define the Gradio Blocks interface
with gr.Blocks() as demo:
    gr.Markdown("# Audio Separator")
    gr.Markdown("Upload an audio file to separate it into different components.")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(label="Upload Audio", type="numpy")
            separate_button = gr.Button("Separate")
        
        with gr.Column():
            output_1 = gr.Audio(label="Separation Output 1")
            output_2 = gr.Audio(label="Separation Output 2")
    
    
    # Set up button click event
    separate_button.click(fn=separate_audio, inputs=audio_input, outputs=[output_1, output_2])

# Launch the interface
demo.launch()