import json
import faiss
import numpy as np
from pathlib import Path
from sentence_transformers import SentenceTransformer


class RAGRetriever:
    def __init__(self, index_dir: str,
                 embed_model: str = "pritamdeka/S-PubMedBert-MS-MARCO"):
        self.index_dir = Path(index_dir)
        self.index = faiss.read_index(str(self.index_dir / "faiss.index"))
        with open(self.index_dir / "chunks.jsonl") as f:
            self.chunks = [json.loads(line) for line in f]
        self.encoder = SentenceTransformer(embed_model)

    def search(self, query: str, k: int = 4, min_score: float = 0.3) -> list[dict]:
        q = self.encoder.encode([query], normalize_embeddings=True)
        scores, idxs = self.index.search(np.asarray(q, dtype="float32"), k)
        return [
            {**self.chunks[i], "score": float(scores[0][j])}
            for j, i in enumerate(idxs[0])
            if i >= 0 and scores[0][j] >= min_score
        ]