import numpy as np
import torchaudio
import gradio as gr
from pydub import AudioSegment
import noisereduce as nr
import torch
from torchaudio.pipelines import WAV2VEC2_ASR_BASE_960H
from datetime import timedelta
import os

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

def transcribe_audio(file):
    # Cargar el audio grabado usando pydub
    audio_segment = AudioSegment.from_file(file)
    audio_segment = audio_segment.set_frame_rate(16000)  # Cambiar la frecuencia de muestreo a 16 kHz
    audio_segment.export("temp.wav", format="wav")

    # Cargar el archivo WAV
    waveform, sample_rate = torchaudio.load("temp.wav")

    # 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
    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(file):
    # Asegurarse de que el archivo de audio está siendo tratado correctamente
    # El archivo puede venir como una tupla (ruta, tipo) o como un solo archivo
    # Verifica si el archivo es una tupla
    if isinstance(file, tuple):
        audio_file_path = file[0]
    else:
        audio_file_path = file
    
    # Comprobar si audio_file_path es realmente una cadena
    if not isinstance(audio_file_path, str):
        raise ValueError("El archivo de audio no es una cadena válida.")
    
    # Verificar si el archivo realmente existe
    import os
    if not os.path.exists(audio_file_path):
        raise ValueError(f"El archivo {audio_file_path} no existe o no es accesible.")
    
    # Transcribir el audio y generar subtítulos
    transcription = transcribe_audio(audio_file_path)
    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 la interfaz Gradio
iface = gr.Interface(
    fn=process_audio,
    inputs="audio",
    outputs=["text", "file"],
    title="🎤 Grabador de Audio y Transcriptor a Subtítulos SRT 📜",
    description="Graba tu voz y obtén la transcripción junto con un archivo SRT de subtítulos.",
)

iface.launch(share=True)