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)
    segments = []

    for i in range(0, len(audio), segment_duration_ms):
        segment = audio[i:i + segment_duration_ms]
        segment_path = f"segment_{i // (60 * 1000)}.wav"
        segment.export(segment_path, format="wav")
        segments.append(segment_path)
    
    return segments

# Función para crear el PDF con barra de progreso
def crear_pdf(transcripcion, progress=gr.Progress()):
    pdf = FPDF()
    pdf.add_page()

    pdf.set_font("Arial", size=12)
    lines = transcripcion.split("\n")

    # Agregar la transcripción al PDF línea por línea con progreso
    for i, line in enumerate(lines):
        pdf.multi_cell(0, 10, line)
        progress(i / len(lines))  # Actualiza la barra de progreso

    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()):
    audio = AudioSegment.from_file(audio_file)
    if len(audio) > 5 * 60 * 1000:
        segments = dividir_audio(audio_file)
    else:
        segments = [audio_file]

    all_transcriptions = []
    for i, segment in enumerate(segments):
        with open(segment, "rb") as audio:
            output = replicate.run(
                "model/replicate",
                input={"task": "transcribe", "audio": audio}
            )
            transcription = output['text']
            all_transcriptions.append(f"Segment {i + 1}:\n{transcription}")
            yield "\n".join(all_transcriptions), None

    full_transcription = "\n".join(all_transcriptions)
    pdf_path = crear_pdf(full_transcription, progress)
    yield full_transcription, pdf_path

# Definir el script JavaScript para auto-scroll
scroll_jscode = """
    function autoScroll() {
        var textbox = document.getElementById('transcription');
        textbox.scrollTop = textbox.scrollHeight;
    }
"""

# Crea la interfaz de Gradio con scroll automático en el campo de texto
with gr.Blocks() as demo:
    gr.Markdown("# Transcriptor con Progreso y Scroll")
    audio_input = gr.Audio(type="filepath", label="Sube tu archivo de audio")
    output_text = gr.Textbox(lines=20, label="Transcripción", elem_id="transcription")
    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])
    
    # JavaScript para auto-scroll al final de la caja de texto
    demo.load(js=scroll_jscode)

demo.launch()