"""Cross-encoder reranker that fixes the order of candidates already recalled by the server.

Embedding recall scores each memory on its own, so when a distractor shares the query's wording the
right memory can land second: the README's weak case, "Who is presenting at the quarterly review?",
ranks "The team agreed to ship the billing fix before the quarterly review." just above the memory
that actually answers it. A cross-encoder reads the (query, memory) pair together and scores
relevance directly, which separates those two. It runs on CPU over the handful of candidates the
server already returned, so it adds one small forward pass per recall and never touches the store.
"""
from __future__ import annotations

from collections.abc import Callable, Sequence
from functools import lru_cache
from typing import Any

MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"  # ~80 MB, CPU, trained for query-passage relevance


@lru_cache(maxsize=1)
def _model():
    # Imported lazily so the module is cheap to import and the model loads once, on first recall.
    from sentence_transformers import CrossEncoder

    return CrossEncoder(MODEL)


def rerank(query: str, candidates: Sequence[Any], *, text: Callable[[Any], str] = lambda c: c.text) -> list[Any]:
    """Return ``candidates`` ordered by cross-encoder relevance to ``query``, highest first.

    The sort is stable, so candidates the reranker cannot separate keep the server's order, and a
    list shorter than two is returned unchanged. ``text`` reads the string to score from each
    candidate (``m.text`` for a recalled memory).
    """
    items = list(candidates)
    if len(items) < 2:
        return items
    scores = _model().predict([(query, text(c)) for c in items])
    order = sorted(range(len(items)), key=lambda i: scores[i], reverse=True)
    return [items[i] for i in order]