import sounddevice as sd
import numpy as np
import torchaudio
import gradio as gr
from scipy.io.wavfile import write
import noisereduce as nr
import torch
from torchaudio.pipelines import WAV2VEC2_ASR_BASE_960H
from datetime import timedelta

# Configuración
duration = 5  # Duración de la grabación en segundos
sample_rate = 44100  # Frecuencia de muestreo

# Cargar el modelo de transcripción de voz
bundle = WAV2VEC2_ASR_BASE_960H
asr_model = bundle.get_model()
labels = bundle.get_labels()

def record_audio():
    # Grabar audio
    print("Grabando...")
    audio = sd.rec(int(duration * sample_rate), samplerate=sample_rate, channels=1, dtype='float64')
    sd.wait()  # Esperar a que termine la grabación
    print("Grabación terminada.")
    
    # Guardar como archivo WAV
    write('recorded_audio.wav', sample_rate, audio)
    
    return 'recorded_audio.wav'

def transcribe_audio(file):
    # Cargar el audio grabado
    waveform, sample_rate = torchaudio.load(file)

    # Reducir el ruido
    reduced_noise = nr.reduce_noise(y=waveform.numpy()[0], sr=sample_rate)
    
    # Convertir de nuevo a tensor
    reduced_waveform = torch.tensor(reduced_noise).unsqueeze(0)

    # Asegurarse de que el audio está en la frecuencia de muestreo del modelo
    if sample_rate != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        reduced_waveform = resampler(reduced_waveform)

    # Realizar la transcripción usando el modelo de ASR
    with torch.no_grad():
        logits = asr_model(reduced_waveform)

    # Obtener las predicciones de las etiquetas
    predicted_ids = torch.argmax(logits, dim=-1)

    # Convertir IDs a texto
    transcription = ''.join([labels[i] for i in predicted_ids[0].tolist() if i < len(labels)])
    
    return transcription.strip()

def format_as_srt(transcription):
    # Dividir la transcripción en partes de ejemplo para los subtítulos
    # Esto puede ajustarse según sea necesario para definir la duración de los subtítulos
    words = transcription.split()
    srt_output = []
    start_time = 0
    end_time = 0

    for i, word in enumerate(words):
        start_time = end_time
        end_time = start_time + 1  # Duración fija de 1 segundo por palabra (ajustar según necesidad)
        
        # Formato SRT
        srt_output.append(f"{i + 1}")
        srt_output.append(f"{str(timedelta(seconds=start_time)).split('.')[0].replace(',', '.')},000 --> {str(timedelta(seconds=end_time)).split('.')[0].replace(',', '.')},000")
        srt_output.append(f"{word}\n")

    return ''.join(srt_output)

def process_audio():
    # Grabar audio y luego transcribir
    audio_file = record_audio()
    transcription = transcribe_audio(audio_file)
    srt_content = format_as_srt(transcription)

    # Guardar el contenido SRT en un archivo
    with open('subtitles.srt', 'w') as f:
        f.write(srt_content)

    return transcription, 'subtitles.srt'

# Crear ejemplos para la interfaz
examples = [
    ["Graba un saludo", "Hola, ¿cómo estás?"],
    ["Graba una presentación", "Soy un apasionado de la programación."],
    ["Graba una explicación", "El reconocimiento de voz es fascinante."],
]

# Crear la interfaz Gradio
iface = gr.Interface(
    fn=process_audio,
    inputs=None,
    outputs=["text", "file"],
    title="🎤 Grabador de Audio y Transcriptor a Subtítulos SRT 📜",
    description="👋 Bienvenido a nuestra aplicación de grabación y transcripción de audio a subtítulos. Graba tu voz y obtén la transcripción junto con un archivo SRT de subtítulos.",
    examples=examples,
    theme="default",
    layout="vertical",
    css="""
    .title { color: #4A90E2; font-weight: bold; }
    .description { font-size: 16px; color: #555; }
    .footer { text-align: center; font-size: 12px; color: #777; }
    """,
)

iface.launch()