import os
import faiss
import torch
from transformers import AutoTokenizer, AutoModel
from sentence_transformers import SentenceTransformer
from PyPDF2 import PdfReader

class RAGRetriever:
    def __init__(self):
        self.encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
        self.index = faiss.IndexFlatL2(384)
        self.contexts = []
        self.ids = []

    def add_document(self, text):
        sentences = text.split("\n")
        clean_sentences = [s.strip() for s in sentences if s.strip()]
        embeddings = self.encoder.encode(clean_sentences)
        self.index.add(embeddings)
        self.contexts.extend(clean_sentences)

    def retrieve(self, query, top_k=3):
        q_vec = self.encoder.encode([query])
        D, I = self.index.search(q_vec, top_k)
        return [self.contexts[i] for i in I[0]]

def extract_text_from_file(file_path):
    ext = os.path.splitext(file_path)[-1].lower()
    if ext == ".txt":
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()
    elif ext == ".pdf":
        reader = PdfReader(file_path)
        return "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
    else:
        return ""