import os
import openai
import whisper
import tempfile
import gradio as gr
from pydub import AudioSegment
import fitz  # PyMuPDF para manejar PDFs
import docx  # Para manejar archivos .docx
import pandas as pd  # Para manejar archivos .xlsx y .csv
#from google.colab import userdata  # Importa userdata de google.colab

# Load environment variables from the Hugging Face environment
openai.api_key = os.getenv("OPENAI_API_KEY")

# Configura tu clave API de OpenAI usando Google Colab userdata
#openai.api_key = userdata.get('OPENAI_API_KEY')

# Cargar el modelo Whisper de mayor calidad una vez
model = whisper.load_model("large")

def preprocess_audio(audio_file):
    """Preprocesa el archivo de audio para mejorar la calidad."""
    try:
        audio = AudioSegment.from_file(audio_file)
        audio = audio.apply_gain(-audio.dBFS + (-20))
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
            audio.export(temp_file.name, format="mp3")
            return temp_file.name
    except Exception as e:
        return f"Error al preprocesar el archivo de audio: {str(e)}"

def transcribir_audio(audio_file):
    """Transcribe un archivo de audio."""
    try:
        archivo_path = preprocess_audio(audio_file) if isinstance(audio_file, str) else preprocess_audio(tempfile.NamedTemporaryFile(delete=False, suffix=".mp3", mode='w+b').name)
        resultado = model.transcribe(archivo_path)
        return resultado.get("text", "Error en la transcripción")
    except Exception as e:
        return f"Error al procesar el archivo de audio: {str(e)}"

def leer_documento(documento_path):
    """Lee el contenido de un documento PDF, DOCX, XLSX o CSV."""
    try:
        if documento_path.endswith(".pdf"):
            doc = fitz.open(documento_path)
            return "\n".join([pagina.get_text() for pagina in doc])
        elif documento_path.endswith(".docx"):
            doc = docx.Document(documento_path)
            return "\n".join([parrafo.text for parrafo in doc.paragraphs])
        elif documento_path.endswith(".xlsx"):
            return pd.read_excel(documento_path).to_string()
        elif documento_path.endswith(".csv"):
            return pd.read_csv(documento_path).to_string()
        else:
            return "Tipo de archivo no soportado. Por favor suba un documento PDF, DOCX, XLSX o CSV."
    except Exception as e:
        return f"Error al leer el documento: {str(e)}"

# Diccionario con las plantillas de Historias AMP
plantillas_amp = {
    "Plantilla 1": "<div class=\"AmpStorie-container1\"><div class=\"AmpStorie-text1\">{contenido}</div></div>",
    "Plantilla 2": "<div class=\"AmpStorie-text2\">{contenido}</div>",
    "Plantilla 3": "<div class=\"AmpStorie-text3\">{contenido}</div>",
    "Plantilla 4": "<div class=\"AmpStorie-container AmpStorie-container_full\"><div class=\"AmpStorie-text4\">{contenido}</div></div>",
    "Plantilla 5": "<div class=\"AmpStorie-container AmpStorie-container_full\"><div class=\"AmpStorie-elementContainer\"><div class=\"AmpStorie-line\"></div><div class=\"AmpStorie-text5\">{contenido}</div></div></div>",
}

def generar_historia_amp(instrucciones, hechos, tono, plantilla_seleccionada, *args):
    """Genera una historia AMP a partir de instrucciones, hechos y transcripciones."""
    tamaño = 40  # Tamaño fijo de 40 palabras
    base_de_conocimiento = {"instrucciones": instrucciones, "hechos": hechos, "contenido_documentos": [], "audio_data": []}
    num_audios = 5 * 3  # 5 audios * 3 campos (audio, nombre, cargo)
    audios = args[:num_audios]
    documentos = args[num_audios:]

    for documento in documentos:
        if documento is not None:
            base_de_conocimiento["contenido_documentos"].append(leer_documento(documento.name))

    for i in range(0, len(audios), 3):
        audio_file, nombre, cargo = audios[i:i+3]
        if audio_file is not None:
            base_de_conocimiento["audio_data"].append({"audio": audio_file, "nombre": nombre, "cargo": cargo})

    transcripciones_texto, transcripciones_brutas, total_citas_directas = "", "", 0

    for idx, data in enumerate(base_de_conocimiento["audio_data"]):
        if data["audio"] is not None:
            transcripcion = transcribir_audio(data["audio"])
            transcripcion_texto = f'"{transcripcion}" - {data["nombre"]}, {data["cargo"]}'
            transcripcion_bruta = f'[Audio {idx + 1}]: "{transcripcion}" - {data["nombre"]}, {data["cargo"]}'
            if total_citas_directas < len(base_de_conocimiento["audio_data"]) * 0.8:
                transcripciones_texto += transcripcion_texto + "\n"
                total_citas_directas += 1
            else:
                transcripciones_texto += f'{data["nombre"]} mencionó que {transcripcion}' + "\n"
            transcripciones_brutas += transcripcion_bruta + "\n\n"

    contenido_documentos = "\n\n".join(base_de_conocimiento["contenido_documentos"])

    prompt_interno = """
    Instrucciones para el modelo:
    - Genera 6 bullets o puntos clave relacionados con las instrucciones y hechos proporcionados.
    - Cada bullet debe tener aproximadamente {tamaño} palabras.
    - recuerda: siempre debes generar seis bullets, de {tamaño} palabras cada uno
    - No inventes información nueva.
    - Sé riguroso con los hechos proporcionados.
    - Genera una narrativa conectada y coherente entre los bullets, formando una historia completa.
    - Al procesar los documentos cargados, extrae y resalta citas importantes y testimonios textuales de las fuentes.
    - Al procesar los documentos cargados, extrae y resalta cifras clave.
    """

    prompt = f"""
    {prompt_interno}
    Genera 6 bullets o puntos clave relacionados con la siguiente información, incluyendo un título y un gancho de 15 palabras (el gancho es lo que se conoce en inglés como hook, información adicional que complementa el título). El tono debe ser {tono}.
    Instrucciones: {base_de_conocimiento["instrucciones"]}
    Hechos: {base_de_conocimiento["hechos"]}
    Contenido adicional de los documentos: {contenido_documentos}
    Utiliza las siguientes transcripciones como citas directas e indirectas (sin cambiar ni inventar contenido):
    {transcripciones_texto}
    """

    try:
        respuesta = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1
        )
        texto = respuesta['choices'][0]['message']['content']
        lines = texto.split("\n")

        # Encontrar título y gancho
        titulo = lines[0]
        gancho = lines[1]

        # Generar los bullets
        bullets = []
        bullet = ""
        word_count = 0

        for line in lines[2:]:
            words = line.split()
            if word_count + len(words) > tamaño:
                bullets.append(bullet.strip())
                bullet = line
                word_count = len(words)
            else:
                bullet += " " + line
                word_count += len(words)

        if bullet:
            bullets.append(bullet.strip())

        # Encapsular los bullets en la plantilla seleccionada
        contenido_bullets = ""
        for bullet in bullets:
            if bullet:
                contenido_bullets += plantillas_amp[plantilla_seleccionada].format(contenido=bullet) + "\n\n"

        historia_amp = f"{titulo}\n\n{gancho}\n\n{contenido_bullets}"
        return historia_amp, transcripciones_brutas
    except Exception as e:
        return f"Error al generar la historia AMP: {str(e)}", ""


demo = gr.Blocks() # Crea el contexto de Gradio.Blocks()
with demo:
    with gr.Column(scale=2):
        instrucciones = gr.Textbox(label="Instrucciones para la historia AMP", lines=2)
        hechos = gr.Textbox(label="Describe los hechos de la historia AMP", lines=4)
        tono = gr.Dropdown(label="Tono de la historia AMP", choices=["serio", "neutral", "divertido"], value="neutral")
        plantilla = gr.Dropdown(label="Plantilla AMP", choices=list(plantillas_amp.keys()), value="Plantilla 1")

    with gr.Column(scale=3):
        inputs_list = [instrucciones, hechos, tono, plantilla]
        with gr.Tabs():
            for i in range(1, 6):
                with gr.TabItem(f"Audio {i}"):
                    audio = gr.Audio(type="filepath", label=f"Audio {i}")
                    nombre = gr.Textbox(label="Nombre", scale=1)
                    cargo = gr.Textbox(label="Cargo", scale=1)
                    inputs_list.extend([audio, nombre, cargo])
            for i in range(1, 6):
                with gr.TabItem(f"Documento {i}"):
                    documento = gr.File(label=f"Documento {i}")
                    inputs_list.append(documento)

    btn_generar = gr.Button(value="Generar historia AMP")
    output_historia = gr.Textbox(label="Historia AMP generada", lines=20)
    output_transcripciones = gr.Textbox(label="Transcripciones brutas de los audios", lines=10)

    btn_generar.click(generar_historia_amp, inputs_list, [output_historia, output_transcripciones])

demo.launch(debug=True)