import streamlit as st
import os
import tempfile
import subprocess

# Set Streamlit app title
st.title("Audio Separation App")

# Function to process the audio file
def separate_audio(audio_path):

    print(f"{audio_path=}")
    head, tail = os.path.split(audio_path)

    gradio_temp_path = head
    audio_filename = tail.split('.')[0]
    print(f"{gradio_temp_path=}")
    print(f"{audio_filename=}")
    
    command = f"spleeter separate -p spleeter:2stems {audio_path}"
    command = command.split()
    print(f"{command=}")

    result = subprocess.run(command)
    print(result)
    
    print("--------")
    accompaniment_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/accompaniment.wav"
    vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
    print(f"{accompaniment_path=}")
    print(os.path.exists(accompaniment_path))
    print(f"{vocals_path=}")
    print(os.path.exists(vocals_path))

    return vocals_path, accompaniment_path


def separate_audio_by_stem(audio_path, stem_count):

    print(f"{audio_path=}")
    head, tail = os.path.split(audio_path)

    gradio_temp_path = head
    audio_filename = tail.split('.')[0]
    print(f"{gradio_temp_path=}")
    print(f"{audio_filename=}")
    print(f"{stem_count=}")
    
    command = f"spleeter separate -p spleeter:{stem_count}stems {audio_path}"
    command = command.split()
    print(f"{command=}")

    result = subprocess.run(command)
    print(result)

    if stem_count == 2:
        accompaniment_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/accompaniment.wav"
        vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"

        print(f"{accompaniment_path=} \t exists: {os.path.exists(accompaniment_path)}")
        print(f"{vocals_path=} \t exists: {os.path.exists(vocals_path)}")
        
        return [
            {'description': 'Accompaniment', 'path':accompaniment_path},
            {'description': 'Vocals', 'path':vocals_path},
        ]
    
    elif stem_count == 4:
        
        vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
        drums_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"
        bass_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"
        other_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"
        
        print(f"{vocals_path=} \t exists: {os.path.exists(vocals_path)}")
        print(f"{drums_path=} \t exists: {os.path.exists(drums_path)}")
        print(f"{bass_path=} \t exists: {os.path.exists(bass_path)}")
        print(f"{other_path=} \t exists: {os.path.exists(other_path)}")

        return [
            {'description': 'Vocals', 'path':vocals_path},
            {'description': 'Drums', 'path':drums_path},
            {'description': 'Bass', 'path':bass_path},
            {'description': 'Other', 'path':other_path},
        ]
        
    elif stem_count == 5:

        piano_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/piano.wav"
        vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
        drums_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"
        bass_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"
        other_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"
        
        print(f"{piano_path=} \t exists: {os.path.exists(vocals_path)}")
        print(f"{vocals_path=} \t exists: {os.path.exists(vocals_path)}")
        print(f"{drums_path=} \t exists: {os.path.exists(drums_path)}")
        print(f"{bass_path=} \t exists: {os.path.exists(bass_path)}")
        print(f"{other_path=} \t exists: {os.path.exists(other_path)}")

        return [
            {'description': 'Vocals', 'path':vocals_path},
            {'description': 'Piano', 'path':piano_path},
            {'description': 'Drums', 'path':drums_path},
            {'description': 'Bass', 'path':bass_path},
            {'description': 'Other', 'path':other_path},
        ]
    
# Streamlit app content
st.write("Upload an audio file")

uploaded_file = st.file_uploader("Choose a file", type=["wav","mp3"])
selected_stem_count = st.radio("Select stem count", (2,4,5))

if uploaded_file is not None:

    if st.button("Submit"):

        # Save the uploaded file to a temporary location
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:        
            temp_file.write(uploaded_file.read())
            temp_file_path = temp_file.name
        
        # Process the uploaded audio file
        separate_audios = separate_audio_by_stem(temp_file_path, selected_stem_count)
    
        # Display the output files for download
        st.write("Output Files:")
        for audio in separate_audios:
            st.write(audio['description'])
            st.audio(audio['path'], format="audio/m4a", start_time=0)