import streamlit as st
import librosa
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):
    # 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]
    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.")
    
    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('Moonarch Music Analysis: Basic Splitter')
st.header("Experience the Future of Sound")
st.write("""
Everything you need to create and release your music, including samples, plugins, unlimited distribution, and the world's best AI mastering engine.
""")

# Display the Start button
if st.button("Start using Moonarch"):
    st.write("Welcome to Moonarch! Let's start creating amazing music.")

# Add upload option
uploaded_file = st.file_uploader("Upload a song", type=["mp3", "wav"])

if uploaded_file is not None:
    st.write("File uploaded successfully.")
    
    # Placeholder for progress bar
    progress_bar = st.progress(0)
    
    # 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())
    
    # Simulate file processing
    for percent_complete in range(100):
        time.sleep(0.01)
        progress_bar.progress(percent_complete + 1)
    
    st.write("File processing complete.")
    
    if st.button('Basic Split'):
        with st.spinner('Splitting audio...'):
            output_files = process_audio(audio_file_path)
            
            st.success('Audio processed successfully!')

            # Display and download each generated audio
            for file in output_files:
                if os.path.isfile(file):
                    file_name = os.path.basename(file)
                    if 'vocals.wav' in file_name:
                        st.write('Vocal')
                    elif 'bass.wav' in file_name:
                        st.write('Bass')
                    elif 'drums.wav' in file_name:
                        st.write('Drums')
                    elif 'accompaniment.wav' in file_name:
                        st.write('Music')
                    
                    st.audio(file)
                    with open(file, 'rb') as f:
                        st.download_button(
                            label=f'Download {file_name}',
                            data=f,
                            file_name=file_name,
                            mime='audio/wav'
                        )
                else:
                    st.error(f"Error: {file} not found.")