import gradio as gr
from pydub import AudioSegment
import librosa
import noisereduce as nr
import numpy as np
from io import BytesIO
import tempfile
import os

# Function to load audio
def load_audio(audio_file_path):
    # Open the file and read its bytes
    with open(audio_file_path, "rb") as f:
        audio_bytes = f.read()
    
    # Load the audio using pydub from the file bytes
    audio_segment = AudioSegment.from_file(BytesIO(audio_bytes))
    audio_array = np.array(audio_segment.get_array_of_samples(), dtype=np.float32)
    sample_rate = audio_segment.frame_rate
    return audio_array, sample_rate

# Function for noise reduction
def reduce_noise(audio, sample_rate):
    reduced_noise_audio = nr.reduce_noise(y=audio, sr=sample_rate)
    return reduced_noise_audio

# Function for pitch shifting
def pitch_shift(audio, sample_rate, n_steps):
    shifted_audio = librosa.effects.pitch_shift(audio, sr=sample_rate, n_steps=n_steps)
    return shifted_audio

# Function for time-stretching
def time_stretch(audio, rate):
    stretched_audio = librosa.effects.time_stretch(audio, rate)
    return stretched_audio

# Function to save audio to a temporary file and return the path
def save_audio(audio, sample_rate):
    # Create a temporary file to save the audio
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    temp_file_name = temp_file.name
    
    # Use pydub to save the audio to the temp file
    audio_segment = AudioSegment(
        np.int16(audio).tobytes(), frame_rate=sample_rate, sample_width=2, channels=1
    )
    audio_segment.export(temp_file_name, format="wav")
    
    temp_file.close()
    return temp_file_name

# Main function to apply audio effects
def apply_effects(audio_file, noise_reduction, pitch_shift_steps, time_stretch_rate):
    audio, sr = load_audio(audio_file)
    
    if noise_reduction:
        audio = reduce_noise(audio, sr)
    
    if pitch_shift_steps != 0:
        audio = pitch_shift(audio, sr, pitch_shift_steps)
    
    if time_stretch_rate != 1.0:
        audio = time_stretch(audio, time_stretch_rate)
    
    return save_audio(audio, sr)

# Gradio UI
def build_ui():
    with gr.Blocks() as demo:
        # Center-aligned Title
        gr.Markdown("<h1 style='text-align: center;'>Shyam's AI Audio Studio</h1>")
        
        # Center-aligned Description
        gr.Markdown(
            """
            <div style='text-align: center;'>
            <p>Welcome to <strong>Shyam's AI Audio Studio</strong>!</p>
            <p>This tool allows you to upload audio files and apply various effects like:</p>
            <ul style="display: inline-block; text-align: left;">
            <li>Noise Reduction</li>
            <li>Pitch Shifting (up or down by semitones)</li>
            <li>Time Stretching (speed up or slow down)</li>
            </ul>
            <p>Experiment with the sliders to fine-tune the effects and get your desired sound!</p>
            </div>
            """
        )
        
        # Input components
        audio_input = gr.Audio(type="filepath", label="Upload Audio File")
        noise_reduction = gr.Checkbox(label="Apply Noise Reduction", value=True)
        pitch_shift_steps = gr.Slider(label="Pitch Shift (in semitones)", minimum=-12, maximum=12, value=0, step=1)
        time_stretch_rate = gr.Slider(label="Time Stretch Rate", minimum=0.5, maximum=2.0, value=1.0, step=0.1)
        
        # Output component
        audio_output = gr.File(label="Download Edited Audio")
        
        # Button to trigger the process
        edit_button = gr.Button("Apply Effects")
        
        # Link the button to the effect function
        edit_button.click(
            apply_effects, 
            inputs=[audio_input, noise_reduction, pitch_shift_steps, time_stretch_rate],
            outputs=audio_output
        )
    
    return demo

# Launch the Gradio app
if __name__ == "__main__":
    ui = build_ui()
    ui.launch()