import fitz  # PyMuPDF
import faiss
from sentence_transformers import SentenceTransformer
import numpy as np

# model = SentenceTransformer('all-MiniLM-L6-v2')
from huggingface_hub import snapshot_download
from sentence_transformers import SentenceTransformer

# Download model to a local folder in your Space
model_path = snapshot_download('sentence-transformers/all-MiniLM-L6-v2')

# Load it from the local path
model = SentenceTransformer(model_path)


def extract_text_from_pdf(pdf_file):
    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
    full_text = ""
    for page in doc:
        full_text += page.get_text()
    return full_text

def chunk_text(text, chunk_size=300):
    words = text.split()
    return [' '.join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]

def embed_chunks(chunks):
    embeddings = model.encode(chunks)
    return np.array(embeddings)

def create_faiss_index(embeddings):
    dim = embeddings.shape[1]
    index = faiss.IndexFlatL2(dim)
    index.add(embeddings)
    return index

def get_top_chunks(query, index, chunks, k=3):
    query_embedding = model.encode([query])
    distances, indices = index.search(query_embedding, k)
    return [chunks[i] for i in indices[0]]