# app.py
import torch
from transformers import (
    DPRContextEncoder, DPRContextEncoderTokenizerFast,
    DPRQuestionEncoder, DPRQuestionEncoderTokenizerFast,
    BartForConditionalGeneration, BartTokenizer
)
from datasets import Dataset
import faiss
import numpy as np
import gradio as gr

# Importar funciones de extracción
from extract_text import extract_text_from_pdf, extract_text_from_docx, extract_text_from_image

# Inicializar modelos y variables globales
ctx_encoder = DPRContextEncoder.from_pretrained('facebook/dpr-ctx_encoder-single-nq-base')
ctx_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained('facebook/dpr-ctx_encoder-single-nq-base')

q_encoder = DPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base')
q_tokenizer = DPRQuestionEncoderTokenizerFast.from_pretrained('facebook/dpr-question_encoder-single-nq-base')

generator = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
gen_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')

# Inicializar dataset y índice
dataset = Dataset.from_dict({'text': []})
embeddings = np.empty((0, ctx_encoder.config.hidden_size), dtype='float32')
index = faiss.IndexFlatIP(ctx_encoder.config.hidden_size)

# Función para actualizar el índice con nuevo texto
def actualizar_indice(nuevo_texto):
    global dataset, embeddings, index

    # Añadir nuevo documento al dataset
    dataset = dataset.add_item({'text': nuevo_texto})

    # Codificar el nuevo documento
    inputs = ctx_tokenizer(nuevo_texto, truncation=True, padding='longest', return_tensors='pt')
    embedding = ctx_encoder(**inputs).pooler_output.detach().numpy()

    # Actualizar embeddings y índice
    embeddings = np.vstack([embeddings, embedding])
    index.add(embedding)

# Función para recuperar documentos relevantes
def retrieve_docs(question, k=5):
    inputs = q_tokenizer(question, return_tensors='pt')
    question_embedding = q_encoder(**inputs).pooler_output.detach().numpy()

    distances, indices = index.search(question_embedding, k)
    retrieved_texts = [dataset[i]['text'] for i in indices[0]]
    return retrieved_texts

# Función para generar respuesta
def generate_answer(question):
    retrieved_docs = retrieve_docs(question)
    context = ' '.join(retrieved_docs)

    input_text = f"Pregunta: {question} Contexto: {context}"
    inputs = gen_tokenizer([input_text], max_length=1024, return_tensors='pt', truncation=True)
    summary_ids = generator.generate(inputs['input_ids'], num_beams=4, max_length=100, early_stopping=True)
    answer = gen_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return answer

# Función principal de la aplicación
def responder(archivo, pregunta):
    texto_extraido = ''
    if archivo is not None:
        file_path = archivo.name
        if file_path.endswith('.pdf'):
            texto_extraido = extract_text_from_pdf(file_path)
        elif file_path.endswith('.docx'):
            texto_extraido = extract_text_from_docx(file_path)
        elif file_path.lower().endswith(('.png', '.jpg', '.jpeg')):
            texto_extraido = extract_text_from_image(file_path)
        else:
            return "Formato de archivo no soportado."

        # Actualizar el índice con el nuevo texto
        actualizar_indice(texto_extraido)

        # Generar respuesta
        respuesta = generate_answer(pregunta)
        return respuesta
    else:
        return "Por favor, sube un archivo."

# Configurar la interfaz de Gradio
interfaz = gr.Interface(
    fn=responder,
    inputs=[
        gr.inputs.File(label="Sube un archivo (PDF, DOCX, Imagen)"),
        gr.inputs.Textbox(lines=2, placeholder="Escribe tu pregunta aquí...")
    ],
    outputs="text",
    title="Aplicación RAG con Extracción de Texto",
    description="Sube un archivo y haz una pregunta sobre su contenido."
)

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