from transformers import pipeline
from pydub import AudioSegment
from pydub.utils import make_chunks
import os
import gradio as gr

# Inicializar el modelo de Hugging Face
modelo = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-large-xlsr-53-spanish")

# Función para dividir audios largos en fragmentos
def dividir_audio(input_path, output_dir, chunk_length_ms=30000):
    """
    Divide un archivo de audio en fragmentos más pequeños.
    Args:
        input_path (str): Ruta al archivo de audio original.
        output_dir (str): Carpeta donde se guardarán los fragmentos.
        chunk_length_ms (int): Duración de cada fragmento en milisegundos.
    """
    audio = AudioSegment.from_file(input_path)
    chunks = make_chunks(audio, chunk_length_ms)
    os.makedirs(output_dir, exist_ok=True)

    for i, chunk in enumerate(chunks):
        chunk_name = os.path.join(output_dir, f"chunk_{i}.wav")
        chunk.export(chunk_name, format="wav")
        print(f"Guardado: {chunk_name}")

# Función para transcribir audios
def transcribir_audio(audio_path):
    """
    Transcribe un archivo de audio usando Hugging Face.
    Args:
        audio_path (str): Ruta al archivo de audio.
    Returns:
        str: Texto transcrito.
    """
    print(f"Procesando: {audio_path}")
    result = modelo(audio_path)
    return result["text"]

# Función principal para manejar la transcripción desde la interfaz
def procesar_audio(input_audio):
    output_dir = "chunks"
    chunk_length_ms = 30000  # Dividir en fragmentos de 30 segundos

    # Guardar el archivo cargado temporalmente
    input_audio_path = "input_audio.wav"
    input_audio.export(input_audio_path, format="wav")

    print("Dividiendo audio...")
    dividir_audio(input_audio_path, output_dir, chunk_length_ms)

    print("Transcribiendo fragmentos...")
    transcripcion_final = []
    for filename in sorted(os.listdir(output_dir)):
        if filename.endswith(".wav"):
            filepath = os.path.join(output_dir, filename)
            texto = transcribir_audio(filepath)
            transcripcion_final.append(texto)

    # Concatenar la transcripción final
    transcripcion_completa = " ".join(transcripcion_final)

    # Devolver la transcripción completa
    return transcripcion_completa

# Crear la interfaz de usuario con Gradio
interface = gr.Interface(
    fn=procesar_audio,
    inputs=gr.Audio(source="upload", type="file", label="Sube tu archivo de audio"),
    outputs=gr.Textbox(label="Transcripción"),
    title="Transcriptor de Audio",
    description="Sube un archivo de audio de una reunión y obtén la transcripción completa."
)

# Ejecutar la aplicación
if __name__ == "__main__":
    interface.launch()