import gradio as gr
from pydub import AudioSegment
from scipy.signal import convolve
import numpy as np
import tempfile

# Reverb function using convolution with a simple impulse response
def add_reverb(audio_segment, decay=0.5, delay=200):
    # Convert AudioSegment to numpy array
    samples = np.array(audio_segment.get_array_of_samples())
    
    # Create a simple reverb impulse response
    impulse_response = np.zeros(delay * 10)
    impulse_response[::delay] = decay ** np.arange(len(impulse_response) // delay)
    
    # Apply convolution for reverb
    processed_samples = convolve(samples, impulse_response, mode='full')[:len(samples)]
    processed_samples = np.int16(processed_samples / np.max(np.abs(processed_samples)) * 32767)
    
    # Convert numpy array back to AudioSegment
    processed_audio = AudioSegment(
        processed_samples.tobytes(),
        frame_rate=audio_segment.frame_rate,
        sample_width=audio_segment.sample_width,
        channels=audio_segment.channels
    )
    return processed_audio

# Wrapper function to apply selected effect
def apply_effect(audio_file, effect, reverb_decay, reverb_delay):
    audio = AudioSegment.from_file(audio_file)
    
    if effect == "Reverb":
        audio = add_reverb(audio, decay=reverb_decay, delay=int(reverb_delay))
    
    # Save the modified audio to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
        audio.export(temp_file.name, format="wav")
    return temp_file.name

# Gradio Interface
def build_ui():
    with gr.Blocks() as demo:
        gr.Markdown("### Audio Effects with Gradio")
        
        with gr.Row():
            audio_input = gr.Audio(source="upload", type="filepath", label="Upload Audio")
            audio_output = gr.Audio(label="Processed Audio")
        
        with gr.Row():
            effect = gr.Radio(["Reverb"], label="Choose Effect")
            reverb_decay = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Reverb Decay")
            reverb_delay = gr.Slider(50, 500, value=200, step=10, label="Reverb Delay (ms)")
        
        apply_button = gr.Button("Apply Effect")
        
        apply_button.click(
            fn=apply_effect,
            inputs=[audio_input, effect, reverb_decay, reverb_delay],
            outputs=[audio_output]
        )
    
    return demo

# Launch Gradio App
demo = build_ui()
demo.launch()