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

# 1. Preprocesamiento de TXT
def extract_text_from_txt(txt_path):
	text = ""
	with open(txt_path, 'r') as file:
		text += file.read()
	return text

# 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]

def search_question(query, k=5):
    query_vector = create_embedding(query)
    Distances, Indices = index.search(np.array([query_vector]), k)
    return Indices[0]

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

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

def generate_answer(query, context):
    response = generator({
		"context": context,
		"question": query}, num_return_sequences=1, truncation=True, max_new_tokens=300)
	
    return response['answer']


# Procesar PDFs y crear índice
pdf_folder = "/home/www_ssl/assistant/data/pdf"
txt_folder = "/home/www_ssl/assistant/data/txt"
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)
        # chunks = text.split(".") # Divide el texto en oraciones
        # all_chunks.extend(chunks)

for filename in os.listdir(txt_folder):
    if filename.endswith(".txt"):
        txt_path = os.path.join(txt_folder, filename)
        text = extract_text_from_txt(txt_path)
        #chunks = chunk_text(text)
        chunks = text.split(".") # Divide el texto en oraciones
        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_question(query)
    relevant_chunks = [all_chunks[i] for i in relevant_chunk_indices]
    context = " ".join(relevant_chunks)
    response = generate_answer(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}")