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 == 'Basic Split':
        output_dir = os.path.join('basic_splits', file_name_without_ext)
        ensure_directory_exists(output_dir)
        os.chdir('basic_splits')
        from moonarch_basic import BasicSplitter
        splitter = BasicSplitter(file_path)
        splitter.run()
        st.write("Basic split process completed.")
        os.chdir(home_directory)
        ensure_directory_exists(os.path.join('vocal_remover', file_name_without_ext))
        os.chdir('vocal_remover')
        from moonarch_vocal_remover import VocalRemover
        music_sep = VocalRemover(file_path)
        music_sep.run()
        st.write("Vocal remover process completed.")
        os.chdir(home_directory)
        return [
            os.path.join('basic_splits', file_name_without_ext, 'vocals.wav'),
            os.path.join('vocal_remover', file_name_without_ext, 'accompaniment.wav'),
            os.path.join('basic_splits', file_name_without_ext, 'bass.wav'),
            os.path.join('basic_splits', file_name_without_ext, 'drums.wav')
        ]
    
    
# Streamlit app
st.title('Monarch Music Analysis: Basic Splitter')

task_choice = st.selectbox('Choose a task', ['Basic Split'])

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')
            elif 'bass.wav' in file:
                st.write('Bass')
            elif 'drums.wav' in file:
                st.write('Drums')
            elif 'piano.wav' in file:
                st.write('Piano')
            elif 'strings_or_pads.wav' in file:
                st.write('Strings or Pads')
               
            st.audio(file)
            with open(file, 'rb') as f:
                new_file_name = 'music.wav' if file_name == 'accompaniment.wav' else ('strings_or_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)