import gradio as gr
from moviepy.editor import VideoFileClip, AudioFileClip
import librosa
import noisereduce as nr
import soundfile as sf
import numpy as np
from scipy.signal import butter, filtfilt
import os

def remove_wind_noise(video_path):
    try:
        # === TEMP FILES ===
        temp_audio = "temp_audio.wav"
        clean_audio = "clean_audio.wav"
        output_video = "cleaned_video.mp4"

        # === STEP 1: Extract audio ===
        video = VideoFileClip(video_path)
        video.audio.write_audiofile(temp_audio, logger=None)

        # === STEP 2: Load audio ===
        audio_data, sr = librosa.load(temp_audio, sr=None)

        # === STEP 3: High-pass filter ===
        def highpass_filter(data, sr, cutoff=100, order=3):
            nyquist = 0.5 * sr
            normal_cutoff = cutoff / nyquist
            b, a = butter(order, normal_cutoff, btype="high", analog=False)
            return filtfilt(b, a, data)

        filtered_audio = highpass_filter(audio_data, sr)

        # === STEP 4: Noise reduction ===
        noise_clip = filtered_audio[:int(sr * 0.5)]
        reduced_noise = nr.reduce_noise(
            y=filtered_audio,
            y_noise=noise_clip,
            sr=sr,
            prop_decrease=0.5,
            stationary=False
        )

        # === STEP 5: Save cleaned audio ===
        sf.write(clean_audio, reduced_noise, sr)

        # === STEP 6: Merge cleaned audio with video ===
        clean_audio_clip = AudioFileClip(clean_audio)
        final_video = video.set_audio(clean_audio_clip)
        final_video.write_videofile(output_video, codec="libx264", audio_codec="aac", logger=None)

        # Clean up temporary files
        for f in [temp_audio, clean_audio]:
            if os.path.exists(f):
                os.remove(f)

        return output_video

    except Exception as e:
        return f"❌ Error: {str(e)}"


# === GRADIO INTERFACE ===
description = """
🎧 **Wind Noise Remover**
Upload a video (.MOV, .MP4) and remove background wind noise using Python + MoviePy + Librosa.
"""

app = gr.Interface(
    fn=remove_wind_noise,
    inputs=gr.Video(label="Upload your video"),
    outputs=gr.Video(label="Cleaned video"),
    title="Wind Noise Remover",
    description=description,
)

if __name__ == "__main__":
    app.launch()