import os
import replicate
import gradio as gr
from pydub import AudioSegment

# Asegúrate de que REPLICATE_API_TOKEN esté configurado en las variables de entorno
replicate_token = os.getenv("REPLICATE_API_TOKEN")

if not replicate_token:
    raise ValueError("No se ha encontrado el token de API de Replicate.")

# Función para dividir el archivo de audio en segmentos de duración definida (en milisegundos)
def dividir_audio(audio_path, segment_duration_ms=60000):
    audio = AudioSegment.from_file(audio_path)
    audio_length = len(audio)
    segments = []

    # Divide el audio en fragmentos
    for i in range(0, audio_length, segment_duration_ms):
        segment = audio[i:i+segment_duration_ms]
        segment_path = f"segment_{i // 1000}.wav"
        segment.export(segment_path, format="wav")  # Exporta el fragmento como un archivo WAV
        segments.append(segment_path)
    
    return segments

# Función para transcribir el audio
def transcribe_audio(audio_file):
    # Cargar el archivo de audio completo
    audio = AudioSegment.from_file(audio_file)
    audio_duration_minutes = len(audio) / (1000 * 60)  # Duración en minutos

    # Si el audio dura más de 10 minutos, dividirlo en segmentos de 1 minuto
    if audio_duration_minutes > 10:
        segments = dividir_audio(audio_file, segment_duration_ms=60000)
    else:
        segments = [audio_file]  # Si es menor de 10 minutos, no dividir

    # Almacenar todas las transcripciones
    all_transcriptions = []

    # Procesar cada segmento individualmente
    for segment_path in segments:
        with open(segment_path, "rb") as audio:
            output = replicate.run(
                "vaibhavs10/incredibly-fast-whisper:3ab86df6c8f54c11309d4d1f930ac292bad43ace52d10c80d87eb258b3c9f79c",
                input={
                    "task": "transcribe",
                    "audio": audio,  # El archivo de audio cargado en Gradio
                    "language": "None",  # Detecta automáticamente el idioma
                    "timestamp": "chunk",  # Incluye marcas de tiempo
                    "batch_size": 64,
                    "diarise_audio": False
                }
            )
            # Almacenar la transcripción del segmento
            all_transcriptions.append(output['text'])
    
    # Combina todas las transcripciones en una sola cadena
    full_transcription = "\n".join(all_transcriptions)
    return full_transcription  # Devuelve la transcripción completa

# Crea la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Transcripción de Audio usando Whisper")
    audio_input = gr.Audio(type="filepath", label="Sube tu archivo de audio")
    output_text = gr.Textbox(label="Transcripción")
    
    transcribe_button = gr.Button("Transcribir")
    transcribe_button.click(fn=transcribe_audio, inputs=audio_input, outputs=output_text)

# Ejecuta la aplicación
demo.launch()