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

BASE = Path(__file__).resolve().parent.parent
INDEX = BASE / "rag" / "vector.index"
CHUNKS = BASE / "rag" / "chunks.jsonl"

embedder = SentenceTransformer("paraphrase-multilingual-mpnet-base-v2")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", max_length=512)

# Load chunks
texts = [json.loads(line)["text"] for line in open(CHUNKS, "r", encoding="utf-8")]
index = faiss.read_index(str(INDEX))


def search(query, top_k=5, retrieve_k=20):
    # Step 1 — embed query
    q = embedder.encode([query], convert_to_numpy=True).astype("float32")
    faiss.normalize_L2(q)

    # Step 2 — retrieve K candidate chunks
    D, I = index.search(q, retrieve_k)

    candidates = [(i, texts[i]) for i in I[0]]

    # Step 3 — rerank using cross encoder (much more accurate)
    pairs = [[query, c[1]] for c in candidates]
    scores = reranker.predict(pairs)

    # sort by reranker score (descending)
    ranked = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)

    # return top-k final chunks
    final_chunks = [c[1][1] for c in ranked[:top_k]]

    return "\n\n".join(final_chunks)
