import noisereduce as nr
import scipy.io.wavfile as wavfile
import numpy as np
import gradio as gr
import os
import tempfile
import shutil

def denoise_audio_file(input_path, output_path):
    rate, data = wavfile.read(input_path)
    
    if len(data.shape) > 1:
        reduced_noise = np.zeros_like(data, dtype=np.float32)
        for channel in range(data.shape[1]):
            reduced_noise[:, channel] = nr.reduce_noise(y=data[:, channel], sr=rate)
    else:
        reduced_noise = nr.reduce_noise(y=data, sr=rate)
    
    wavfile.write(output_path, rate, reduced_noise.astype(data.dtype))
    return output_path

def process_single_file(file):
    if not file.name.endswith('.wav'):
        raise gr.Error("Please upload a WAV file")
    
    # Use the original filename for the denoised file, but in a temp dir
    name, ext = os.path.splitext(os.path.basename(file.name))
    base_filename = f"{name}_denoised{ext}"
    temp_dir = tempfile.mkdtemp()
    output_path = os.path.join(temp_dir, base_filename)
    denoise_audio_file(file.name, output_path)
    return output_path

def process_batch_files(files):
    output_files = []
    temp_dir = tempfile.mkdtemp()
    for file in files:
        if file.name.endswith('.wav'):
            name, ext = os.path.splitext(os.path.basename(file.name))
            base_filename = f"{name}_denoised{ext}"
            output_path = os.path.join(temp_dir, base_filename)
            denoise_audio_file(file.name, output_path)
            output_files.append(output_path)
    return output_files

with gr.Blocks(title="Audio Noise Reducer") as demo:
    gr.Markdown("# 🎧 Audio Noise Reduction")
    gr.Markdown("Upload WAV files to remove background noise using AI-powered processing.")
    
    with gr.Tab("Single File Processing"):
        with gr.Row():
            with gr.Column():
                single_file = gr.File(label="Upload WAV File", file_types=[".wav"])
                single_btn = gr.Button("Process File")
            with gr.Column():
                single_output = gr.File(label="Download Denoised File")
                single_status = gr.Textbox(label="Processing Status", interactive=False)
                
        single_btn.click(
            fn=process_single_file,
            inputs=single_file,
            outputs=single_output,
            api_name="process_single"
        )

    with gr.Tab("Batch Processing"):
        with gr.Row():
            with gr.Column():
                batch_files = gr.File(label="Upload WAV Files", file_count="multiple", file_types=[".wav"])
                batch_btn = gr.Button("Process Files")
            with gr.Column():
                batch_output = gr.Files(label="Download Denoised Files")
                batch_status = gr.Textbox(label="Processing Status", interactive=False)
                
        batch_btn.click(
            fn=process_batch_files,
            inputs=batch_files,
            outputs=batch_output,
            api_name="process_batch"
        )

    demo.queue()
    
if __name__ == "__main__":
    demo.launch()