"""In-memory BM25 keyword index (the 'keyword' half of hybrid search)."""
import heapq
import math
import re
import threading
from collections import Counter

STOP = set(
    "a an and are as at be but by for from has have he her his i in is it its of on or she that the their "
    "them they this to was we were what when where which who why will with you your do does did how can "
    "kya hai hain ka ki ke ko se me mai par pe aur tha thi the".split()
)


def tokenize(text: str) -> list[str]:
    return [t for t in re.findall(r"\w+", text.lower()) if len(t) > 1 and t not in STOP]


class BM25Index:
    k1 = 1.5
    b = 0.75

    def __init__(self):
        self._lock = threading.RLock()
        self._dirty = True
        self._clear()

    def _clear(self):
        self.chunks: list[dict] = []
        self.lens: list[int] = []
        self.postings: dict[str, list[tuple[int, int]]] = {}
        self.avgdl = 1.0

    def mark_dirty(self) -> None:
        self._dirty = True

    def build(self, chunks: list[dict]) -> None:
        postings: dict[str, list[tuple[int, int]]] = {}
        lens = []
        for i, c in enumerate(chunks):
            toks = tokenize(c["text"])
            lens.append(len(toks))
            for term, tf in Counter(toks).items():
                postings.setdefault(term, []).append((i, tf))
        with self._lock:
            self.chunks, self.lens, self.postings = chunks, lens, postings
            self.avgdl = (sum(lens) / len(lens)) if lens else 1.0

    def ensure_fresh(self) -> None:
        if not self._dirty:
            return
        with self._lock:
            if not self._dirty:
                return
            self._dirty = False  # cleared first: a concurrent mark_dirty() during the rebuild is not lost
            from . import store

            self.build(store.scroll_all())

    def search(self, query: str, doc_ids=None, pages=None, limit: int = 20) -> list[dict]:
        self.ensure_fresh()
        doc_set = set(doc_ids) if doc_ids else None
        page_set = set(pages) if pages else None
        with self._lock:
            n = len(self.chunks)
            if n == 0:
                return []
            scores: dict[int, float] = {}
            for term in set(tokenize(query)):
                plist = self.postings.get(term)
                if not plist:
                    continue
                df = len(plist)
                idf = math.log(1 + (n - df + 0.5) / (df + 0.5))
                for i, tf in plist:
                    c = self.chunks[i]
                    if doc_set and c["doc_id"] not in doc_set:
                        continue
                    if page_set and c["page"] not in page_set:
                        continue
                    denom = tf + self.k1 * (1 - self.b + self.b * self.lens[i] / self.avgdl)
                    scores[i] = scores.get(i, 0.0) + idf * tf * (self.k1 + 1) / denom
            top = heapq.nlargest(limit, scores.items(), key=lambda kv: kv[1])
            return [{**self.chunks[i], "score": s} for i, s in top]


bm25 = BM25Index()