import os
import gradio as gr
from audio_separator.separator import Separator

# Function for separating the audio
def separate_audio(audio):
    if audio is None:
        return None, None

    # Step 1: Get the audio file path
    input_file = audio

    # Step 2: Set up output directory and models for separation
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)

    separator = Separator(output_dir=output_dir)

    # Define output paths
    vocals_path = os.path.join(output_dir, 'Vocals.wav')
    instrumental_path = os.path.join(output_dir, 'Instrumental.wav')

    with gr.Progress() as progress:
        # Load model
        progress(0, "Loading model...")
        separator.load_model(model_filename='model_bs_roformer_ep_317_sdr_12.9755.ckpt')
        progress(20, "Model loaded. Starting separation...")

        # Step 3: Splitting track into Vocal and Instrumental
        voc_inst = separator.separate(input_file)

        # Check if separation was successful
        if len(voc_inst) != 2:
            return None, None
        
        # Save the separated files
        os.rename(voc_inst[0], instrumental_path)
        os.rename(voc_inst[1], vocals_path)

        progress(100, "Separation complete!")
    
    # Return paths to the processed files
    return instrumental_path, vocals_path

# Define the Gradio Interface
with gr.Blocks(theme="NoCrypt/miku@1.2.2") as demo:
    gr.Markdown("# Audio Separator Gradio Demo")

    with gr.Row():
        with gr.Column():
            link_input = gr.Audio(label="Upload Audio File", type="filepath")
            separate_button = gr.Button("Separate Audio")
        
        with gr.Column():
            instrumental_output = gr.Audio(label="Instrumental Output", type="filepath")
            vocals_output = gr.Audio(label="Vocals Output", type="filepath")
            
    # Define button functionality
    separate_button.click(
        separate_audio,
        inputs=[link_input],
        outputs=[
            instrumental_output,
            vocals_output,
        ]
    )

# Launch the Gradio app
demo.launch(debug=True, share=True)