import gradio as gr
from transformers import pipeline
import os
from TTS.api import TTS
from pydub import AudioSegment

# Inicializar los modelos de Coqui TTS para dos voces
tts_juan = TTS(model_name="tts_models/es/mai/tacotron2-DDC", progress_bar=False, gpu=False)
tts_maria = TTS(model_name="tts_models/es/css10/vits", progress_bar=False, gpu=False)

# Modelo de resumen compatible con español
resumidor = pipeline("summarization", model="mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization")

# Función para convertir texto a audio con voces diferentes
def convertir_texto_a_audio(texto, nombre_archivo, voz):
    if voz == "Juan":
        tts_juan.tts_to_file(text=texto, file_path=nombre_archivo)
    elif voz == "María":
        tts_maria.tts_to_file(text=texto, file_path=nombre_archivo)
    else:
        # Voz predeterminada
        tts_juan.tts_to_file(text=texto, file_path=nombre_archivo)

def procesar_conversacion(conversacion):
    lineas = conversacion.strip().split('\n')
    dialogos = []
    for linea in lineas:
        if linea.startswith("Juan:"):
            dialogos.append(("Juan", linea.replace("Juan:", "").strip()))
        elif linea.startswith("María:"):
            dialogos.append(("María", linea.replace("María:", "").strip()))
    return dialogos

def generar_audio_conversacion(dialogos):
    audios = []
    for idx, (persona, texto) in enumerate(dialogos):
        ruta_audio = os.path.join("/tmp", f"dialogo_{idx}.wav")
        convertir_texto_a_audio(texto, ruta_audio, persona)
        audios.append(ruta_audio)
    return audios

def combinar_audios(rutas_audios):
    conversacion_audio = AudioSegment.empty()
    for ruta in rutas_audios:
        audio = AudioSegment.from_file(ruta)
        conversacion_audio += audio
    ruta_final = os.path.join("/tmp", "conversacion_final.wav")
    conversacion_audio.export(ruta_final, format="wav")
    return ruta_final

# Creamos el modelo de generación de diálogo
dialogador = pipeline("text-generation", model="microsoft/DialoGPT-small")

def transformar_texto(archivo):
    if archivo is None:
        raise gr.Error("Por favor, sube un archivo de texto.")

    try:
        with open(archivo.name, 'r', encoding='utf-8') as f:
            texto = f.read()
    except Exception as e:
        raise gr.Error(f"Error al leer el archivo: {e}")

    max_palabras = 2000  # Nuevo límite de palabras
    num_palabras = len(texto.split())
    if num_palabras > max_palabras:
        raise gr.Error(f"El texto excede el límite de {max_palabras} palabras.")

    # Resumimos el texto para que el prompt no sea demasiado largo
    try:
        resumen = resumidor(
            texto,
            max_length=130,  # Ajusta según las capacidades del modelo
            min_length=30,
            do_sample=False,
            clean_up_tokenization_spaces=True
        )[0]['summary_text']
    except Exception as e:
        raise gr.Error(f"Error al resumir el texto: {e}")

    try:
        # Generamos la conversación utilizando el resumen
        prompt = f"""
Crear una conversación educativa entre Juan y María sobre el siguiente tema:

{resumen}

La conversación debe tener el formato:

Juan: [pregunta]
María: [respuesta]
...

Por favor, asegúrate de que la conversación cubra los puntos principales del texto.
"""
        output = dialogador(
            prompt,
            max_length=1024,  # Aumentamos el límite de tokens generados
            num_return_sequences=1,
            clean_up_tokenization_spaces=True
        )
        if not output:
            raise gr.Error("El modelo no generó ninguna salida. Por favor, intenta con un texto más corto o diferente.")
        conversacion = output[0]['generated_text']
    except Exception as e:
        raise gr.Error(f"Error al generar la conversación: {e}")

    # Procesamos la conversación y generamos los audios
    dialogos = procesar_conversacion(conversacion)
    if not dialogos:
        raise gr.Error("No se pudo procesar la conversación generada. Asegúrate de que el modelo está generando el formato correcto.")
    rutas_audios = generar_audio_conversacion(dialogos)
    ruta_audio_final = combinar_audios(rutas_audios)

    return ruta_audio_final

# Crear la interfaz de usuario
interfaz = gr.Interface(
    fn=transformar_texto,
    inputs=gr.File(label="Sube tu documento de texto (.txt)"),
    outputs=gr.Audio(type="filepath", label="Escucha la conversación generada"),
    title="📝➡️🎙️ Transformador de Texto a Conversación",
    description="Sube un documento de texto y conviértelo en una conversación educativa entre dos personas.",
    theme="default",
    allow_flagging="never"
)

interfaz.launch()