import ffmpeg
import numpy as np
import librosa
import os
import time
from transformers import pipeline
import gradio as gr

# Modelos de Hugging Face para español
#transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-small", language="es")
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-small")
summarizer = pipeline("summarization", model="mrm8488/bert2bert_shared-spanish-finetuned-summarization")

# Variable global para estados y transcripciones
state = {"status": "Esperando transmisión...", "transcriptions": [], "summary": ""}

# Función para esperar inicio de transmisión RTMP
def wait_for_stream(rtmp_url):
    state["status"] = "Esperando transmisión..."
    print(state["status"])
    while True:
        try:
            probe = ffmpeg.probe(rtmp_url, format='flv')
            if probe:
                state["status"] = "¡Transmisión detectada!"
                print(state["status"])
                break
        except ffmpeg.Error:
            time.sleep(5)

# Procesar transmisión RTMP en tiempo real
def process_rtmp(rtmp_url):
    audio_output = "stream_audio.wav"
    transcription = []

    state["status"] = "Transcribiendo en tiempo real..."
    print(state["status"])

    # Iniciar FFmpeg para extraer audio en tiempo real
    process = (
        ffmpeg
        .input(rtmp_url, format='flv')
        .output(audio_output, format='wav', acodec='pcm_s16le', ac=1, ar=16000)
        .overwrite_output()
        .run_async(pipe_stdout=True, pipe_stderr=True)
    )

    try:
        while True:
            if os.path.exists(audio_output):
                audio_data, _ = librosa.load(audio_output, sr=16000)
                if len(audio_data) > 0:
                    text = transcriber(np.array(audio_data))["text"]
                    transcription.append(text)
                    state["transcriptions"].append(text)
                    print(f"Transcripción: {text}")
            time.sleep(2)  # Procesar cada 2 segundos
    except KeyboardInterrupt:
        process.terminate()

    state["status"] = "Transmisión finalizada"
    print(state["status"])
    return " ".join(transcription)

# Generar resumen
def finalize_summary(transcription):
    state["status"] = "Generando resumen..."
    print(state["status"])
    summary = summarizer(transcription, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
    state["summary"] = summary
    state["status"] = "Resumen listo"
    print(state["status"])
    return summary

# Flujo principal
def process_and_finalize():
    rtmp_url = "rtmp://37.27.213.138/live/stream"

    # Esperar inicio de transmisión
    wait_for_stream(rtmp_url)

    # Procesar transmisión y transcribir en tiempo real
    transcription = process_rtmp(rtmp_url)

    # Generar resumen
    summary = finalize_summary(transcription)

    return summary

# Interfaz Gradio
def display_status():
    # Mostrar estados y transcripciones en tiempo real
    return f"Estado: {state['status']}\n\nTranscripciones:\n" + "\n".join(state["transcriptions"]) + f"\n\nResumen final:\n{state['summary']}"

demo = gr.Interface(
    fn=display_status,
    inputs=None,
    outputs="text",
    title="Estado de Transmisión y Resumen",
    description="Muestra el estado de la transmisión, transcripciones en tiempo real y el resumen generado."
)

if __name__ == "__main__":
    demo.launch()

# import gradio as gr
# import numpy as np
# import librosa
# from transformers import pipeline

# # Cargar el modelo de transcripción Whisper
# transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-small")

# # Función para procesar y transcribir el audio
# def transcribe(audio):
#     if audio is None:
#         raise ValueError("No se ha proporcionado un archivo de audio.")

#     # Cargar el archivo de audio como un array NumPy
#     audio_data, _ = librosa.load(audio, sr=16000)  # Resample a 16 kHz
#     result = transcriber(np.array(audio_data))
#     return result["text"]

# # Crear la interfaz Gradio
# demo = gr.Interface(
#     fn=transcribe,
#     inputs=gr.Audio(type="filepath"),  # Subida de archivos de audio
#     outputs="text",
#     title="Transcripción de Audio en Vivo",
#     description="Sube un archivo de audio para transcribir su contenido automáticamente."
# )

# # Lanzar la aplicación
# if __name__ == "__main__":
#     demo.launch()