import os
import moviepy as mp
from pydub import AudioSegment
from pydub.effects import normalize
from scipy.io.wavfile import read, write
import noisereduce as nr
import numpy as np
import gradio as gr

# --- FUNCIONES PARA EL PROCESAMIENTO DE AUDIO ---

def convert_audio_to_wav(input_file_path, wav_file_path):
    """Convierte cualquier archivo de audio a WAV usando pydub."""
    audio = AudioSegment.from_file(input_file_path)  # Cargar el archivo de audio
    audio.export(wav_file_path, format="wav")       # Exportar como WAV

def extract_audio(video_path, output_audio_path):
    """Extrae el audio de un video y lo guarda como un archivo de audio."""
    video = mp.VideoFileClip(video_path)
    video.audio.write_audiofile(output_audio_path)

def denoise_audio(input_audio_path, output_audio_path):
    """Aplica reducción de ruido a un archivo de audio."""
    rate, data = read(input_audio_path)
    if data.ndim > 1:  # Convertir a mono si el audio es estéreo
        data = np.mean(data, axis=1).astype(np.int16)
    reduced_noise = nr.reduce_noise(y=data, sr=rate)
    write(output_audio_path, rate, reduced_noise.astype(np.int16))

def process_audio(input_audio_path, output_audio_path):
    """Aplica mejoras como normalización y ecualización al audio."""
    audio = AudioSegment.from_file(input_audio_path)
    audio = normalize(audio)  # Normaliza el volumen
    # Opcional: Filtrar frecuencias no deseadas (ejemplo: ruido de baja frecuencia)
    audio = audio.low_pass_filter(3000).high_pass_filter(100)
    audio.export(output_audio_path, format="wav")

def replace_audio(video_path, new_audio_path, output_video_path):
    """Reemplaza el audio original de un video con uno procesado."""
    video = mp.VideoFileClip(video_path)
    audio = mp.AudioFileClip(new_audio_path)
    video = video.set_audio(audio)
    video.write_videofile(output_video_path, codec="libx264", audio_codec="aac")


# --- FUNCIONES PRINCIPALES PARA PROCESAR VIDEO O AUDIO ---

def process_video_audio(video_path, output_video_path):
    """Procesa el audio de un video y lo integra nuevamente."""
    temp_audio_path = "temp_audio.wav"
    temp_denoised_audio_path = "temp_denoised_audio.wav"
    temp_processed_audio_path = "temp_processed_audio.wav"
    
    # Extraer el audio del video
    extract_audio(video_path, temp_audio_path)
    
    # Reducir ruido del audio
    denoise_audio(temp_audio_path, temp_denoised_audio_path)
    
    # Procesar y mejorar el audio
    process_audio(temp_denoised_audio_path, temp_processed_audio_path)
    
    # Reemplazar el audio original en el video
    replace_audio(video_path, temp_processed_audio_path, output_video_path)
    
    # Limpiar archivos temporales
    os.remove(temp_audio_path)
    os.remove(temp_denoised_audio_path)
    os.remove(temp_processed_audio_path)

def process_audio_file(input_audio_path, output_audio_path):
    """Procesa únicamente un archivo de audio."""
    temp_denoised_audio_path = "temp_denoised_audio.wav"
    
    # Reducir ruido del audio
    denoise_audio(input_audio_path, temp_denoised_audio_path)
    
    # Procesar y mejorar el audio
    process_audio(temp_denoised_audio_path, output_audio_path)
    
    # Limpiar archivo temporal
    os.remove(temp_denoised_audio_path)


# --- INTERFAZ GRADIO MEJORADA ---

def audio_processing_interface(video_file=None, audio_file=None):
    if video_file is not None:
        output_video_path = "processed_video.mp4"
        process_video_audio(video_file.name, output_video_path)
        return output_video_path
    elif audio_file is not None:
        output_audio_path = "processed_audio.wav"
        process_audio_file(audio_file.name, output_audio_path)
        return output_audio_path


# Interfaz Gradio Mejorada
iface = gr.Interface(
    fn=audio_processing_interface,
    inputs=[
        gr.File(
            label="Sube un video (para procesar su audio)",
            type="filepath",
            file_count="single",
            file_types=[".mp4", ".avi", ".mov", ".mkv"],
        ),
        gr.File(
            label="O sube un archivo de audio",
            type="filepath",
            file_count="single",
            file_types=[".mp3", ".wav", ".ogg", ".flac", ".aac"],
        ),
    ],
    outputs=gr.File(label="Archivo Procesado"),
    live=False,
    title="Procesador de Audio y Video",
    description="Sube un archivo de video o audio para procesarlo. El audio será mejorado (reducción de ruido, normalización) y, si es video, se reemplazará el audio del video.",
)

iface.launch(share=True)