import streamlit as st
import librosa
import soundfile as sf
import os
import time
import numpy as np

# Set the home directory path for moonarch
home_directory = os.getcwd()

# Function to ensure directory exists
def ensure_directory_exists(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

# Function to process the audio
def process_audio(file_path, choice):
    # Ensure to navigate to the home directory first
    os.chdir(home_directory)
    
    file_name_without_ext = os.path.splitext(os.path.basename(file_path))[0]
    
    if choice == 'Vocal Remove':
        output_dir = os.path.join('vocal_remover', file_name_without_ext)
        ensure_directory_exists(output_dir)
        os.chdir('vocal_remover')
        from moonarch_vocal_remover import VocalRemover
        splitter = VocalRemover(file_path)
        splitter.run()
        st.write("Vocal removal process completed.")
        os.chdir(home_directory)
        return [os.path.join(output_dir, 'vocals.wav'), os.path.join(output_dir, 'accompaniment.wav')]
    

    

    
# Streamlit app
st.title('Monarch Music Analysis: Vocal Remover')

task_choice = st.selectbox('Choose a task', ['Vocal Remove'])

uploaded_file = st.file_uploader('Upload an audio file', type=['mp3', 'wav'])

if uploaded_file is not None:
    # Save uploaded file to a temporary location
    audio_file_path = os.path.join(home_directory, f'{uploaded_file.name}')
    with open(audio_file_path, 'wb') as f:
        f.write(uploaded_file.getbuffer())
    
    y, sr = librosa.load(audio_file_path, sr=None)  # sr=None to use the native sample rate

    if st.button('Start'):
        progress_bar = st.progress(0)
        
        # Simulate progress bar for processing
        for percent_complete in range(100):
            time.sleep(0.01)
            progress_bar.progress(percent_complete + 1)

        output_files = process_audio(audio_file_path, task_choice)
        
        st.success('Audio processed successfully!')

        for file in output_files:
            file_name = os.path.basename(file)
            if 'vocals.wav' in file:
                st.write('Vocal')
            elif 'accompaniment.wav' in file:
                st.write('Music')
            
                
            st.audio(file)
            with open(file, 'rb') as f:
                new_file_name = 'music.wav' if file_name == 'accompaniment.wav' else ('strings_pads.wav' if file_name == 'other.wav' else file_name)
                st.download_button(label=f'Download {new_file_name}', data=f, file_name=new_file_name)