import os
import numpy as np
import zipfile
from pydub import AudioSegment
import streamlit as st
from io import BytesIO

# Noise generation functions
def generate_gaussian_noise(duration, sample_rate=44100, std_dev=0.1):
    """Generate Gaussian noise."""
    samples = np.random.normal(0, std_dev, size=(duration * sample_rate,))
    audio_segment = AudioSegment(
        samples.tobytes(),
        frame_rate=sample_rate,
        sample_width=2,  # 16-bit audio
        channels=1
    )
    return audio_segment

def generate_impulse_noise(duration, sample_rate=44100, impulse_rate=0.05):
    """Generate impulse noise."""
    samples = np.zeros(int(duration * sample_rate))
    num_impulses = int(sample_rate * duration * impulse_rate)
    impulse_indices = np.random.randint(0, len(samples), num_impulses)
    samples[impulse_indices] = np.random.uniform(-1, 1, size=num_impulses)
    
    audio_segment = AudioSegment(
        samples.tobytes(),
        frame_rate=sample_rate,
        sample_width=2,  # 16-bit audio
        channels=1
    )
    return audio_segment

def add_noise(audio_file, noise_type='gaussian', noise_level=0.1, start_time=None, end_time=None):
    """Add noise to a clean audio file."""
    audio = AudioSegment.from_file(audio_file)
    duration = len(audio)

    # Check for specific duration
    if start_time is not None and end_time is not None:
        start_time = int(start_time * 1000)
        end_time = int(end_time * 1000)
        if start_time < 0 or end_time > duration or start_time >= end_time:
            st.error("Invalid start or end time.")
            return None
        noise_duration = (end_time - start_time) // 1000
    else:
        start_time = 0
        end_time = duration
        noise_duration = duration // 1000

    # Generate noise
    if noise_type == 'gaussian':
        noise = generate_gaussian_noise(noise_duration)
    elif noise_type == 'impulse':
        noise = generate_impulse_noise(noise_duration)
    else:
        st.error("Invalid noise type.")
        return None

    # Normalize noise
    noise_array = np.array(noise.get_array_of_samples())
    noise_array = noise_array * noise_level / np.max(np.abs(noise_array))
    noise_array = np.clip(noise_array, -1.0, 1.0)

    noise = AudioSegment(
        noise_array.tobytes(),
        frame_rate=noise.frame_rate,
        sample_width=noise.sample_width,
        channels=noise.channels
    )

    # Overlay noise for the specified segment
    silence = AudioSegment.silent(duration=end_time - start_time)
    noisy_segment = silence.overlay(noise)
    
    noisy_audio = audio[:start_time] + noisy_segment + audio[end_time:]
    return noisy_audio

# Streamlit App
def main():
    st.title("🎵 Audio Noise Generator")
    st.sidebar.title("Instructions")
    st.sidebar.write("""
    - Upload one or more audio files (.wav or .mp3).
    - Choose noise type (Gaussian or Impulse).
    - Add noise to the full audio or a specific time duration.
    - Save the processed audio files.
    """)

    # Upload files
    uploaded_files = st.file_uploader("Upload Audio Files", type=["wav", "mp3"], accept_multiple_files=True)

    # Noise type selection
    noise_type = st.selectbox("Select Noise Type", ["gaussian", "impulse"])

    # Noise level
    noise_level = st.slider("Select Noise Level", 0.01, 1.0, 0.1, 0.01)

    # Duration selection
    add_to_specific_duration = st.checkbox("Add noise to specific duration")
    if add_to_specific_duration:
        start_time = st.number_input("Start Time (in seconds)", min_value=0.0, value=0.0, step=0.1)
        end_time = st.number_input("End Time (in seconds)", min_value=0.1, value=1.0, step=0.1)
    else:
        start_time, end_time = None, None

    # Process and save files
    if st.button("Process Audio"):
        if uploaded_files:
            zip_buffer = BytesIO()
            with zipfile.ZipFile(zip_buffer, "w") as zf:
                for uploaded_file in uploaded_files:
                    st.write(f"Processing {uploaded_file.name}...")
                    audio_data = BytesIO(uploaded_file.read())
                    
                    noisy_audio = add_noise(audio_data, noise_type, noise_level, start_time, end_time)
                    if noisy_audio:
                        # Rename file based on noise type
                        base_name, ext = os.path.splitext(uploaded_file.name)
                        output_name = f"{noise_type}_{base_name}.wav"

                        # Save to in-memory buffer
                        buffer = BytesIO()
                        noisy_audio.export(buffer, format="wav")
                        buffer.seek(0)

                        # Add to zip archive
                        zf.writestr(output_name, buffer.read())
                        st.success(f"Added {noise_type} noise to {uploaded_file.name}.")
            
            zip_buffer.seek(0)
            st.write("### Download All Processed Files")
            st.download_button(
                label="Download All as ZIP",
                data=zip_buffer,
                file_name="processed_audio_files.zip",
                mime="application/zip"
            )
        else:
            st.warning("Please upload at least one audio file.")

if __name__ == "__main__":
    main()