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
import os

# Cargar las variables de entorno desde el entorno de Hugging Face
openai.api_key = os.getenv("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)}"

def generar_noticia(instrucciones, hechos, tamaño, tono, *args):
    """Genera una noticia a partir de instrucciones, hechos y transcripciones."""
    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:
    - Debes seguir los principios de una noticia: es decir, procura siempre responder las 5 W de una noticia en el primer párrafo (Who?, What?, When?, Where?, Why?).
    - Asegúrate de que al menos el 80% de las citas sean directas y estén entrecomilladas.
    - El 20% restante puede ser citas indirectas.
    - No inventes información nueva.
    - Sé riguroso con los hechos proporcionados.
    - 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}
    Escribe una noticia con la siguiente información, incluyendo un título, 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), y el cuerpo del contenido cuyo tamaño es {tamaño} palabras. 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
        )
        noticia = respuesta['choices'][0]['message']['content']
        return noticia, transcripciones_brutas
    except Exception as e:
        return f"Error al generar la noticia: {str(e)}", ""

with gr.Blocks() as demo:
    gr.Markdown("## Generador noticias todo en uno")
    with gr.Row():
        with gr.Column(scale=2):
            instrucciones = gr.Textbox(label="Instrucciones para la noticia", lines=2)
            hechos = gr.Textbox(label="Describe los hechos de la noticia", lines=4)
            tamaño = gr.Number(label="Tamaño del cuerpo de la noticia (en palabras)", value=100)
            tono = gr.Dropdown(label="Tono de la noticia", choices=["serio", "neutral", "divertido"], value="neutral")
        with gr.Column(scale=3):
            inputs_list = [instrucciones, hechos, tamaño, tono]
            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}", type="filepath", file_count="single")
                        inputs_list.append(documento)

    gr.Markdown("---") # Separador visual

    with gr.Row():
        transcripciones_output = gr.Textbox(label="Transcripciones", lines=10)

    gr.Markdown("---") # Separador visual

    with gr.Row():
        generar = gr.Button("Generar noticia")
    with gr.Row():
        noticia_output = gr.Textbox(label="Noticia generada", lines=20)

    generar.click(fn=generar_noticia, inputs=inputs_list, outputs=[noticia_output, transcripciones_output])

demo.launch(share=True)