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 sumario de 20 palabras, 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)