from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import textwrap
import pickle
import os
import logging

logger = logging.getLogger(__name__)

class RAG:
    def __init__(self):
        self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
        self.index = faiss.IndexFlatL2(384)
        self.metadata = {}
        self.index_file = "faiss_index/index.faiss"
        self.metadata_file = "faiss_index/metadata.pkl"
        os.makedirs("faiss_index", exist_ok=True)

        if os.path.exists(self.index_file):
            self.index = faiss.read_index(self.index_file)
            with open(self.metadata_file, "rb") as f:
                self.metadata = pickle.load(f)

    async def embed_document(self, file_id: str, text: str):
        chunks = textwrap.wrap(text, 500)
        if not chunks:
            logger.warning(f"Empty document provided for file_id {file_id}")
            raise ValueError("Document contains no text to embed.")

        embeddings = self.embedder.encode(chunks)
        ids = np.array([i for i in range(len(self.metadata), len(self.metadata) + len(chunks))])
        self.index.add(embeddings.astype("float32"))

        for i, chunk in enumerate(chunks):
            self.metadata[ids[i]] = {"file_id": file_id, "text": chunk}

        faiss.write_index(self.index, self.index_file)
        with open(self.metadata_file, "wb") as f:
            pickle.dump(self.metadata, f)

    async def query_document(self, question: str, file_id: str) -> list:
        query_embedding = self.embedder.encode([question])[0].astype("float32")
        distances, indices = self.index.search(np.array([query_embedding]), k=3)

        return [
            self.metadata[idx]["text"]
            for idx in indices[0]
            if idx in self.metadata and self.metadata[idx]["file_id"] == file_id
        ]