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

# 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=5*60*1000):
    audio = AudioSegment.from_file(audio_path)
    audio_length = len(audio)
    segments = []

    # Divide el audio en fragmentos de la duración especificada
    for i in range(0, audio_length, segment_duration_ms):
        segment = audio[i:i + segment_duration_ms]  # Cada fragmento de hasta la duración especificada
        segment_path = f"segment_{i // (60 * 1000)}.wav"  # Nombre del archivo con el índice del minuto
        segment.export(segment_path, format="wav")  # Exporta el fragmento como un archivo WAV
        
        # Verifica el tamaño del archivo y asegura que no supere el límite de 10MB
        if os.path.getsize(segment_path) > 10 * 1024 * 1024:  # 10 MB
            print(f"Warning: Segment {segment_path} exceeds 10MB, reducing segment duration.")
            return dividir_audio(audio_path, segment_duration_ms // 2)  # Recursivamente reduce el tamaño del segmento
        
        segments.append(segment_path)
    
    return segments

# Función para crear el PDF
def crear_pdf(transcripcion):
    pdf = FPDF()
    pdf.add_page()

    pdf.set_font("Arial", size=12)
    
    # Agregar la transcripción al PDF línea por línea
    for line in transcripcion.split("\n"):
        pdf.multi_cell(0, 10, line)
    
    # Guardar el archivo PDF en la carpeta actual
    pdf_path = "/mnt/data/transcripcion_audio.pdf"
    pdf.output(pdf_path)
    
    return pdf_path

# Función para transcribir los segmentos de audio y generar el PDF
def transcribe_audio_y_pdf(audio_file, progress=gr.Progress()):
    # 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 5 minutos, dividirlo en segmentos de 5 minutos
    if audio_duration_minutes > 5:
        segments = dividir_audio(audio_file, segment_duration_ms=5 * 60 * 1000)  # 5 minutos en milisegundos
    else:
        segments = [audio_file]  # Si es menor de 5 minutos, no dividir

    # Almacenar todas las transcripciones
    all_transcriptions = []

    # Procesar cada segmento individualmente y mostrar progresivamente
    for index, segment_path in enumerate(segments):
        with open(segment_path, "rb") as audio:
            try:
                progress(index / len(segments))  # Actualizar el progreso
                output = replicate.run(
                    "vaibhavs10/incredibly-fast-whisper:3ab86df6c8f54c11309d4d1f930ac292bad43ace52d10c80d87eb258b3c9f79c",
                    input={
                        "task": "transcribe",
                        "audio": audio,  # El archivo de audio cargado
                        "language": "None",  # Detecta automáticamente el idioma
                        "timestamp": "chunk",  # Incluye marcas de tiempo
                        "batch_size": 64,
                        "diarise_audio": False
                    }
                )
                transcription = output['text']
                all_transcriptions.append(f"Segment {index + 1}:\n{transcription}")

            except Exception as e:
                return f"Error transcribiendo el segmento {index + 1}: {e}"

    # Combina todas las transcripciones en una sola cadena
    full_transcription = "\n".join(all_transcriptions)

    # Crear PDF y devolver la ruta del archivo
    pdf_path = crear_pdf(full_transcription)
    return full_transcription, pdf_path

# Crea la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Transcriptor de Audio a Texto (Genera PDF)")
    audio_input = gr.Audio(type="filepath", label="Sube tu archivo de audio")
    output_text = gr.Textbox(label="Transcripción")
    output_pdf = gr.File(label="Descarga el PDF")

    transcribe_button = gr.Button("Transcribir y Crear PDF")
    transcribe_button.click(fn=transcribe_audio_y_pdf, inputs=audio_input, outputs=[output_text, output_pdf])

# Ejecuta la aplicación
demo.launch()