import streamlit as st
from spleeter.separator import Separator
import os

# Function to perform audio separation
def separate_audio(input_path, output_path, task):
    separator = Separator(task)
    os.makedirs(output_path, exist_ok=True)
    separator.separate_to_file(input_path, output_path)

# Main function to create the Streamlit app
def main():
    st.title("Audio Separation App")

    # File uploader
    uploaded_file = st.file_uploader("Upload an audio file", type=["mp3", "wav"])

    # Task selection
    task_to_do = st.selectbox("Choose your task", ["vocal_remove", "basic", "advance"])

    if task_to_do == 'vocal_remove':
        task = 'spleeter:2stems'
    elif task_to_do == 'basic':
        task = 'spleeter:4stems'
    elif task_to_do == 'advance':
        task = 'spleeter:5stems'
    else:
        st.error("Invalid task")
        return

    # Button to perform separation
    if st.button("Split"):
        if uploaded_file is not None:
            # Save the uploaded file temporarily
            input_path = os.path.join("temp", uploaded_file.name)
            with open(input_path, "wb") as f:
                f.write(uploaded_file.getbuffer())

            # Create output directory path using the base name of the input file
            base_name = os.path.splitext(os.path.basename(input_path))[0]
            output_path = os.path.join("temp", base_name)

            # Perform the separation
            separate_audio(input_path, output_path, task)

            # Display the results
            st.success("Separation completed")

            if task_to_do == 'vocal_remove':
                st.audio(os.path.join(output_path, 'vocals.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'accompaniment.wav'), format='audio/wav', start_time=0)
            elif task_to_do == 'basic':
                st.audio(os.path.join(output_path, 'vocals.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'drums.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'bass.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'other.wav'), format='audio/wav', start_time=0)
            elif task_to_do == 'advance':
                st.audio(os.path.join(output_path, 'vocals.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'piano.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'drums.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'bass.wav'), format='audio/wav', start_time=0)
                st.audio(os.path.join(output_path, 'other.wav'), format='audio/wav', start_time=0)
        else:
            st.error("Please upload an audio file")

if __name__ == "__main__":
    main()