import os
import PyPDF2
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from transformers import pipeline

# 1. Preprocesamiento de PDFs
def extract_text_from_pdf(pdf_path):
    with open(pdf_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        text = ""
        for page in reader.pages:
            text += page.extract_text()
    return text

def chunk_text(text, chunk_size=1000, overlap=200):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap
    return chunks

# 2. Creación de embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embedding(text):
    return model.encode(text)

# 3. Indexación
index = faiss.IndexFlatL2(384)  # 384 es la dimensión de los embeddings de 'all-MiniLM-L6-v2'

# 4. Recuperación
def search(query, k=3):
    query_vector = create_embedding(query)
    D, I = index.search(np.array([query_vector]), k)
    return I[0]

# 5. Generación de respuestas
generator = pipeline('text-generation', model='gpt2')

def generate_response(query, context):
    prompt = f"Pregunta: {query}\nContexto: {context}\nRespuesta:"
    response = generator(prompt, max_length=150, num_return_sequences=1)[0]['generated_text']
    return response.split("Respuesta:")[-1].strip()

# Procesar PDFs y crear índice
pdf_folder = "/home/www_ssl/assistant/pdf"
all_chunks = []

for filename in os.listdir(pdf_folder):
    if filename.endswith(".pdf"):
        pdf_path = os.path.join(pdf_folder, filename)
        text = extract_text_from_pdf(pdf_path)
        chunks = chunk_text(text)
        all_chunks.extend(chunks)

for chunk in all_chunks:
    embedding = create_embedding(chunk)
    index.add(np.array([embedding]))

# Función principal del asistente
def project_assistant(query):
    relevant_chunk_indices = search(query)
    relevant_chunks = [all_chunks[i] for i in relevant_chunk_indices]
    context = " ".join(relevant_chunks)
    response = generate_response(query, context)
    return response

# Ejemplo de uso
question = "¿Cuáles son las fases principales de un proyecto?"
answer = project_assistant(question)
print(f"Pregunta: {question}")
print(f"Respuesta: {answer}")