import io
import base64
from datetime import datetime
from io import BytesIO
import numpy as np
import requests
import streamlit as st
import streamlit_vertical_slider as svs
from pydub import AudioSegment
from streamlit import session_state as st_state

# Add CSS to block download buttons in audio players
st.markdown("""
<style>
/* Hide download button in all audio players - comprehensive targeting */
audio::-webkit-media-controls-panel button[aria-label="Download"],
audio::-webkit-media-controls-enclosure button[aria-label="Download"],
audio::-internal-media-controls-download-button,
audio::-webkit-media-controls-download-button,
audio::-webkit-media-controls-overflow-button,
audio::-webkit-media-controls-mute-button ~ button,
audio::-webkit-media-controls-panel button:last-child,
audio::-webkit-media-controls-panel menu,
audio::-webkit-media-controls-panel menuitem[aria-label*="Download"] {
    display: none !important;
    opacity: 0 !important;
    pointer-events: none !important;
    visibility: hidden !important;
    width: 0 !important;
    height: 0 !important;
    position: absolute !important;
    z-index: -1000 !important;
    clip: rect(0 0 0 0) !important;
    clip-path: inset(50%) !important;
}

/* Make the audio player background opaque to hide any potential UI elements */
audio::-webkit-media-controls-panel {
    background-color: rgba(35, 35, 35, 0.85) !important;
}

/* Prevent download through context menu but allow volume and play control */
audio {
    -webkit-user-select: none !important;
    user-select: none !important;
    -webkit-touch-callout: none !important;
}
</style>
""", unsafe_allow_html=True)

# Add a background image
page_bg_img = '''
<style>
.stApp {
background-image: url("https://songlabai.com/wp-content/uploads/2024/03/4.png");
background-size: cover;
}
</style>
'''
st.markdown(page_bg_img, unsafe_allow_html=True)

# Initialize session state variables for audio processing
if "vocal_audio" not in st_state:
    st_state.vocal_audio = None
if "vocal_sample_rate" not in st_state:
    st_state.vocal_sample_rate = None
if "audio" not in st_state:
    st_state.audio = None
if "audio_pydub" not in st_state:
    st_state.audio_pydub = None
if "audio_sample_rate" not in st_state:
    st_state.audio_sample_rate = None
if "augmented_audio" not in st_state:
    st_state.augmented_audio = None
if "augmented_audio_pydub" not in st_state:
    st_state.augmented_audio_pydub = None
if "augmented_audio_sample_rate" not in st_state:
    st_state.augmented_audio_sample_rate = None
if 'audio_bytes' not in st.session_state:
    st.session_state['audio_bytes'] = None
if 'playing_state' not in st.session_state:
    st.session_state['playing_state'] = 'none'  # 'none', 'generated', 'processed'
if 'audio_metadata' not in st.session_state:
    st.session_state['audio_metadata'] = None

def secure_audio_player(audio_bytes, format="audio/wav"):
    """
    Enhanced secure audio player that completely blocks downloads
    
    Parameters:
    audio_bytes: Audio data as bytes
    format: Audio format mimetype
    """
    if audio_bytes is None:
        return
        
    # Encode to base64
    b64 = base64.b64encode(audio_bytes).decode()
    
    # Custom HTML with additional JavaScript security
    custom_html = f"""
    <div style="pointer-events: auto;" oncontextmenu="return false;">
        <audio controls 
               controlsList="nodownload nofullscreen noremoteplayback" 
               disableRemotePlayback
               oncontextmenu="return false;"
               src="data:{format};base64,{b64}">
        </audio>
    </div>
    <script>
    // Additional JavaScript protection
    document.addEventListener('DOMContentLoaded', function() {{
        const audioElements = document.querySelectorAll('audio');
        audioElements.forEach(audio => {{
            // Prevent context menu
            audio.addEventListener('contextmenu', e => e.preventDefault());
            
            // Monitor and disable download attempts
            audio.addEventListener('enterpictureinpicture', e => e.preventDefault());
            
            // Remove download option from any menus that might appear
            const observer = new MutationObserver(mutations => {{
                const menus = document.querySelectorAll('menu, menuitem');
                menus.forEach(menu => {{
                    if (menu.textContent.toLowerCase().includes('download')) {{
                        menu.remove();
                    }}
                }});
            }});
            
            observer.observe(document.body, {{ 
                childList: true, 
                subtree: true 
            }});
        }});
    </script>
    """
    
    st.markdown(custom_html, unsafe_allow_html=True)

def convert_audio_segment_to_float_array(audio_pydub):
    """
    Convert a pydub AudioSegment to a NumPy array of type float32.

    Args:
    audio_pydub (AudioSegment): The AudioSegment object to be converted.

    Returns:
    np.ndarray: A NumPy array containing the audio data as float32.
    """
    # Get the raw audio data as a sequence of samples
    samples = audio_pydub.get_array_of_samples()

    # Convert the samples to a NumPy array and normalize to float32
    audio_array = np.array(samples).astype(np.float32)

    # Normalize the audio array to range between -1.0 and 1.0
    max_val = 2**15  # Assuming 16-bit audio, modify this if using different bit depths
    audio_array /= max_val
    return audio_array

def get_wordpress_beats():
    """
    Get beats from WordPress using the custom plugin endpoint
    
    Returns:
    list: List of beat dictionaries with id, title, url, etc.
    """
    # Use the custom plugin's REST API endpoint
    api_url = "https://songlabai.com/wp-json/songlab/v1/beats"
    
    try:
        response = requests.get(api_url)
        print(f"WordPress API response status: {response.status_code}")
        
        if response.status_code == 200:
            beats = response.json()
            print(f"Found {len(beats)} beats in WordPress")
            return beats
        else:
            print(f"Error fetching beats: {response.status_code}")
            return []
            
    except Exception as e:
        print(f"Exception in get_wordpress_beats: {str(e)}")
        return []

def download_wordpress_beat(beat_url):
    """
    Download a beat from WordPress
    
    Parameters:
    beat_url: URL of the beat file
    
    Returns:
    BytesIO: Beat data as BytesIO object or None if failed
    """
    try:
        response = requests.get(beat_url, stream=True)
        if response.status_code == 200:
            return BytesIO(response.content)
        else:
            print(f"Error downloading beat: {response.status_code}")
            return None
    except Exception as e:
        print(f"Exception in download_wordpress_beat: {str(e)}")
        return None

def load_wordpress_beat_for_processing(beat_url, beat_name):
    """
    Download and load a WordPress beat for processing
    
    Parameters:
    beat_url: URL of the beat file
    beat_name: Name of the beat
    
    Returns:
    tuple: (success, message)
    """
    try:
        print(f"Downloading beat from: {beat_url}")
        
        # Download the beat
        beat_buffer = download_wordpress_beat(beat_url)
        if not beat_buffer:
            return False, "Failed to download beat"
        
        # Get file extension from beat URL
        file_extension = beat_url.split('.')[-1].lower() if '.' in beat_url else "mp3"
        
        # Load into AudioSegment
        beat_segment = AudioSegment.from_file(beat_buffer, format=file_extension)
        
        print(f"Loaded beat: {beat_name}, duration: {len(beat_segment)/1000:.2f}s")
        
        # Convert to WAV if needed
        if file_extension != 'wav':
            wav_buffer = BytesIO()
            beat_segment.export(wav_buffer, format='wav')
            wav_buffer.seek(0)
            beat_segment = AudioSegment.from_wav(wav_buffer)
        
        # Store in session state
        st_state.audio_pydub = beat_segment
        st_state.audio_sample_rate = beat_segment.frame_rate
        st_state.augmented_audio_pydub = beat_segment
        st_state.augmented_audio_sample_rate = beat_segment.frame_rate
        
        # Convert to bytes for player
        wav_bytes = BytesIO()
        beat_segment.export(wav_bytes, format='wav')
        wav_bytes.seek(0)
        st.session_state.audio_bytes = wav_bytes.getvalue()
        
        # Set playing state
        st.session_state.playing_state = 'generated'
        
        # Set metadata
        st.session_state.audio_metadata = {
            'genre': 'Beat',
            'energy_level': 'Unknown',
            'tempo': 'Unknown',
            'description': f'Selected beat: {beat_name}',
            'duration': len(beat_segment)/1000  # Convert ms to seconds
        }
        
        return True, "Beat loaded successfully"
    except Exception as e:
        import traceback
        print(f"Error loading beat: {str(e)}")
        print(traceback.format_exc())
        return False, f"Error loading beat: {str(e)}"

def display_audio_player(is_final=False):
    """
    Display audio player for either original or processed audio
    """
    if is_final and st_state.augmented_audio_pydub is not None:
        # Display processed audio
        audio_bytes = st_state.augmented_audio_pydub.export(format="wav").read()
        st.markdown("### 🎧 Final Processed Audio")
        secure_audio_player(audio_bytes)
        
        # Display download button
        st.download_button(
            label="⬇️ Download Processed Audio",
            data=audio_bytes,
            file_name=f"processed_audio_{datetime.now().strftime('%Y%m%d_%H%M%S')}.wav",
            mime="audio/wav"
        )
        
    elif not is_final and st.session_state.get('audio_bytes') is not None:
        st.markdown("### 🎵 Loaded Audio")
        secure_audio_player(st.session_state.audio_bytes)
        
        # Display download button for original audio
        st.download_button(
            label="⬇️ Download Original Audio",
            data=st.session_state.audio_bytes,
            file_name=f"original_audio_{datetime.now().strftime('%Y%m%d_%H%M%S')}.wav",
            mime="audio/wav"
        )

# Main application code
st.title("Music Mixer & Processor")
st.markdown("Load audio, add vocals, and apply professional-quality effects to your tracks")

# Add tabs to separate uploaded music and beat selection
tab1, tab2 = st.tabs(["Upload Your Music", "Choose from Available Beats"])

with tab1:
    # File uploader UI
    uploaded_music = st.file_uploader(
        "Upload music file", 
        type=["wav", "mp3", "ogg", "flac"], 
        help="Upload any music file to process",
        key="user_music_uploader"
    )
    
    if uploaded_music:
        # Show a preview player for the uploaded music
        st.audio(uploaded_music, format=f"audio/{uploaded_music.name.split('.')[-1]}")
        
        # Process button
        if st.button("Load for Processing", key="load_uploaded_music"):
            try:
                # Read the uploaded music into memory
                with st.spinner("Loading audio..."):
                    # Convert the uploaded file to an AudioSegment
                    audio_data = uploaded_music.read()
                    
                    # Create a temporary file-like object
                    audio_buffer = BytesIO(audio_data)
                    audio_buffer.seek(0)
                    
                    # Load into AudioSegment based on file extension
                    file_extension = uploaded_music.name.split('.')[-1].lower()
                    audio_segment = AudioSegment.from_file(audio_buffer, format=file_extension)
                    
                    # Convert to the same format used in your application
                    if file_extension != 'wav':
                        wav_buffer = BytesIO()
                        audio_segment.export(wav_buffer, format='wav')
                        wav_buffer.seek(0)
                        audio_segment = AudioSegment.from_wav(wav_buffer)
                    
                    # Store in session state
                    st_state.audio_pydub = audio_segment
                    st_state.audio_sample_rate = audio_segment.frame_rate
                    st_state.augmented_audio_pydub = audio_segment
                    st_state.augmented_audio_sample_rate = audio_segment.frame_rate
                    
                    # Convert to bytes for player
                    wav_bytes = BytesIO()
                    audio_segment.export(wav_bytes, format='wav')
                    wav_bytes.seek(0)
                    st.session_state.audio_bytes = wav_bytes.getvalue()
                    
                    # Set playing state
                    st.session_state.playing_state = 'generated'
                    
                    # Set metadata
                    st.session_state.audio_metadata = {
                        'genre': 'Unknown (Uploaded)',
                        'energy_level': 'Unknown',
                        'tempo': 'Unknown',
                        'description': f'Uploaded file: {uploaded_music.name}',
                        'duration': len(audio_segment)/1000  # Convert ms to seconds
                    }
                    
                    # Success message
                    st.success("🎉 Music loaded successfully! Scroll down to use processing tools.")
                    
                    # Force refresh to update the UI with the loaded audio
                    st.rerun()
                    
            except Exception as e:
                st.error(f"❌ Error loading audio file: {str(e)}")
                st.info("Please make sure the file is a valid audio file in one of the supported formats.")

with tab2:
    # Show beat selection
    st.markdown("### 🎵 Select a Beat")
    
    # Fetch beats from WordPress
    with st.spinner("Loading beats from WordPress..."):
        available_beats = get_wordpress_beats()
    
    if not available_beats:
        st.info("No beats are currently available. Please upload your own music or try again later.")
    else:
        # Create a dictionary for easy lookup
        beats_dict = {beat['title']: beat for beat in available_beats}
        
        selected_beat_title = st.selectbox(
            "Choose a beat to master and add vocals to",
            options=["-- Select a beat --"] + list(beats_dict.keys()),
            key="user_beat_selection"
        )
        
        if selected_beat_title and selected_beat_title != "-- Select a beat --":
            # Get the selected beat info
            selected_beat = beats_dict[selected_beat_title]
            
            # Display information about the beat
            st.markdown(f"#### Beat Preview")
            st.markdown(f"**Beat:** {selected_beat_title}")
            
            # Show additional metadata if available
            if selected_beat.get('genre') or selected_beat.get('bpm') or selected_beat.get('key'):
                meta_items = []
                if selected_beat.get('genre'):
                    meta_items.append(f"Genre: {selected_beat['genre']}")
                if selected_beat.get('bpm'):
                    meta_items.append(f"BPM: {selected_beat['bpm']}")
                if selected_beat.get('key'):
                    meta_items.append(f"Key: {selected_beat['key']}")
                if selected_beat.get('duration'):
                    meta_items.append(f"Duration: {selected_beat['duration']}s")
                
                if meta_items:
                    st.caption(" | ".join(meta_items))
            
            # Get file extension from URL
            file_extension = selected_beat['url'].split('.')[-1].lower() if '.' in selected_beat['url'] else "mp3"
            
            # Download and preview the beat
            try:
                with st.spinner("Loading beat preview..."):
                    beat_buffer = download_wordpress_beat(selected_beat['url'])
                    if beat_buffer:
                        # Display audio player for preview
                        st.audio(beat_buffer, format=f"audio/{file_extension}")
                        
                        # Add button to load the beat for processing
                        if st.button("Load this Beat for Processing", key="load_selected_beat"):
                            with st.spinner("Loading beat for processing..."):
                                success, message = load_wordpress_beat_for_processing(
                                    selected_beat['url'], 
                                    selected_beat_title
                                )
                                
                                if success:
                                    st.success("🎉 Beat loaded successfully! Scroll down to use processing tools.")
                                    st.rerun()
                                else:
                                    st.error(f"❌ {message}")
                    else:
                        st.error("Failed to load beat preview")
            except Exception as e:
                st.error(f"Error previewing beat: {str(e)}")
                import traceback
                print(traceback.format_exc())

# Add this helper info section
with st.expander("ℹ️ About uploading music and using beats"):
    st.markdown("""
    ### Ways to use music in this app
    
    You have two options for working with existing music:
    
    1. **Upload Your Own Music**
       - Upload your own audio files to process with our tools
       - Support for WAV, MP3, OGG, and FLAC formats
    
    2. **Use Available Beats**
       - Select from pre-made beats provided by Songlab AI
       - Perfect for adding vocals or creating remixes
    
    ### What you can do with the music
    
    - **Add vocals** to instrumental tracks
    - **Apply different effects** to experiment with various sounds
    - **Master** the audio with professional-quality tools
    - **Adjust pitch and volume** to perfect your sound
    """)

# This divider separates the upload section from post-processing options
st.divider()

# Post-processing options - only show if we have audio to process
if st_state.audio_pydub is not None:
    st.header("Audio Processing Tools")
    
    # Display the original audio
    display_audio_player(is_final=False)
    
    # Vocal upload and mixing
    st.subheader("Vocal Processing")
    
    vocal_file = st.file_uploader(
        "Upload Vocal File", type=["mp3", "wav", "ogg", "flac", "aac"],
        help="Upload vocals to mix with your instrumental track"
    )
    
    if vocal_file:
        # Preview the vocals
        st.audio(vocal_file)
        
        try:
            # Load the vocal file
            vocal_data = vocal_file.read()
            file_extension = vocal_file.name.split('.')[-1].lower()
            
            # Create buffer and load into AudioSegment
            vocal_buffer = BytesIO(vocal_data)
            vocal_audio = AudioSegment.from_file(vocal_buffer, format=file_extension)
            
            # Store in session state
            st_state.vocal_audio = vocal_audio
            st.success("✅ Vocals loaded successfully")
        except Exception as e:
            st.error(f"Error loading vocal file: {str(e)}")
    
    # Mixing options
    if st_state.vocal_audio is not None:
        st.subheader("Vocal Mixing")
        
        # Vocal volume
        vocal_volume = st.slider(
            "Vocal Volume", 
            min_value=-10, 
            max_value=10, 
            value=0,
            help="Adjust the volume of the vocals relative to the instrumental"
        )
        
        # Vocal position
        vocal_position = st.slider(
            "Vocal Position (ms)", 
            min_value=0, 
            max_value=5000, 
            value=0,
            help="Adjust when the vocals start (in milliseconds)"
        )
        
        # Apply vocal mix
        if st.button("Mix Vocals with Instrumental"):
            try:
                # Adjust vocal volume if needed
                adjusted_vocals = st_state.vocal_audio + vocal_volume
                
                # Mix vocals at the specified position
                st_state.augmented_audio_pydub = st_state.audio_pydub.overlay(
                    adjusted_vocals, position=vocal_position
                )
                
                # Update derived variables
                st_state.augmented_audio = convert_audio_segment_to_float_array(
                    st_state.augmented_audio_pydub
                )
                st_state.augmented_audio_sample_rate = st_state.augmented_audio_pydub.frame_rate
                st.session_state.playing_state = 'processed'
                
                st.success("✅ Vocals mixed successfully!")
                # We'll display the processed audio later
            except Exception as e:
                st.error(f"Error mixing vocals: {str(e)}")
    
    # Audio Effects
    st.subheader("Audio Effects & Mastering")
    
    # Volume Balance, Pitch Shift, and Effect Buttons
    vol_col, pitch_col, buttons_col = st.columns([2, 2, 2.5])
    
    with buttons_col:
        with st.container(height=371, border=True):
            st.markdown("### Effect Controls")
            apply_stereo = st.button("Apply Stereo Effect", 
                                    help="Creates a stereo widening effect")
            
            reverse = st.button("Apply Audio Reverse",
                               help="Reverses the audio track")
            
            # Low-pass filter
            low_pass = st.button("Apply Low-Pass Filter",
                                help="Reduces high frequencies, creating a warmer sound")
            
            # Add compression button
            compress = st.button("Apply Compression",
                                help="Reduces dynamic range, making quieter parts louder")
            
            reset_post_processing = st.button("Reset All Effects",
                                             help="Remove all applied effects")

    with vol_col:
        with st.container(border=True):
            volume_balance = svs.vertical_slider(
                "Volume Balance",
                min_value=-10.0,
                max_value=10.0,
                default_value=0.0,
                step=0.1,
                slider_color="green",
                track_color="lightgray",
                thumb_color="red",
                thumb_shape="pill",
            )
            vol_button = st.button("Apply Volume Change")

    # Pitch shifting
    with pitch_col:
        with st.container(border=True):
            pitch_semitones = svs.vertical_slider(
                label="Pitch (semitones)",
                min_value=-12,
                max_value=12,
                default_value=0,
                step=1,
                slider_color="red",
                track_color="lightgray",
                thumb_color="red",
                thumb_shape="pill",
            )
            pitch_shift_button = st.button("Apply Pitch Shift")

    effect_applied = False
    
    if st_state.augmented_audio_pydub is not None:
        if vol_button:
            st_state.augmented_audio_pydub = st_state.augmented_audio_pydub + volume_balance
            effect_applied = True
            st.session_state.playing_state = 'processed'

        if apply_stereo:
            st_state.augmented_audio_pydub = st_state.augmented_audio_pydub.pan(
                -0.5
            ).overlay(st_state.augmented_audio_pydub.pan(0.5))
            effect_applied = True
            st.session_state.playing_state = 'processed'
            
        if reverse:
            st_state.augmented_audio_pydub = st_state.augmented_audio_pydub.reverse()
            effect_applied = True
            st.session_state.playing_state = 'processed'
            
        if pitch_shift_button:
            st_state.augmented_audio_pydub = st_state.augmented_audio_pydub._spawn(
                st_state.augmented_audio_pydub.raw_data,
                overrides={
                    "frame_rate": int(
                        st_state.augmented_audio_pydub.frame_rate
                        * (2 ** (pitch_semitones / 12.0))
                    )
                },
            )
            effect_applied = True
            st.session_state.playing_state = 'processed'
            
        if low_pass:
            # Simple low-pass filter approximation using pydub
            # For a better filter, you would use scipy or another DSP library
            from pydub.effects import low_pass_filter
            st_state.augmented_audio_pydub = low_pass_filter(st_state.augmented_audio_pydub, 1500)
            effect_applied = True
            st.session_state.playing_state = 'processed'
            
        if compress:
            # Apply compression using pydub
            # This is a simple approximation of compression
            from pydub.effects import normalize
            st_state.augmented_audio_pydub = normalize(st_state.augmented_audio_pydub)
            effect_applied = True
            st.session_state.playing_state = 'processed'

        if reset_post_processing and st_state.audio_pydub is not None:
            st_state.augmented_audio_pydub = st_state.audio_pydub
            st.session_state.playing_state = 'generated'
            effect_applied = True
    
        # Display the final audio if processed
        if effect_applied or st.session_state.playing_state == 'processed':
            display_audio_player(is_final=True)
            
        # Force rerun if any effect was applied to show the updated audio
        if effect_applied:
            st.rerun()
else:
    # Show a message prompting the user to load audio first
    st.info("👆 Please upload a music file or select a beat from the options above to start processing.")