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 check if the file exists
def file_exists(file_path):
    return os.path.isfile(file_path)

# 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.")
    os.chdir(home_directory)
    
    return [
        os.path.join('basic_splits', file_name_without_ext, 'vocals.wav'),
        os.path.join('basic_splits', file_name_without_ext, 'bass.wav'),
        os.path.join('basic_splits', file_name_without_ext, 'drums.wav'),
        os.path.join('basic_splits', file_name_without_ext, 'accompaniment.wav')
    ]

# Streamlit app
st.title('Moonarch Music Analysis')
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 file_exists(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.")