"""Cross-encoder reranking of retrieved source units.

A cross-encoder reads the question and one candidate text together and scores
their relevance; it orders candidates that BM25 and E5 already retrieved. It
never adds a candidate, removes evidence of a named scope or writes text.

Model choice (measured on CPU, 60 question-text pairs): the multilingual
cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 needs about 1.2 s and ranked the
Slovak probes as well as BAAI/bge-reranker-v2-m3, which needs about 12 s.
The model is prepared into the image (runtime assets); RETRIEVAL_RERANKER=off
disables it for the documented A/B comparison.
"""
from __future__ import annotations

import os
import threading
from functools import lru_cache
from pathlib import Path

DEFAULT_MODEL = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"
MAX_PAIRS = 160
_lock = threading.Lock()


def model_name() -> str:
    return os.getenv("RETRIEVAL_RERANKER_MODEL", DEFAULT_MODEL)


def enabled() -> bool:
    return os.getenv("RETRIEVAL_RERANKER", "on").lower() != "off"


def model_path() -> str:
    """The prepared local copy in the image, else the Hugging Face cache."""
    root = os.getenv("ZP_RUNTIME_ASSETS")
    if root and (Path(root) / "reranker" / "config.json").is_file():
        return str(Path(root) / "reranker")
    return model_name()


@lru_cache(maxsize=1)
def _model():
    from sentence_transformers import CrossEncoder
    local = os.getenv("HF_HUB_OFFLINE") == "1" or model_path() != model_name()
    return CrossEncoder(model_path(), max_length=256, device="cpu", local_files_only=local)


def available() -> bool:
    if not enabled():
        return False
    try:
        _model()
        return True
    except Exception:
        return False


def scores(question: str, texts: list[str]) -> list[float] | None:
    """Relevance score per text, or None when reranking is off or unavailable."""
    if not texts or not available():
        return None
    texts = texts[:MAX_PAIRS]
    with _lock:
        values = _model().predict([(question, t[:1000]) for t in texts], batch_size=32, show_progress_bar=False)
    return [float(v) for v in values]


def order(question: str, rows: list[dict], *, key=lambda r: r["text"]) -> list[dict]:
    """Rows sorted by relevance (stable for ties); unchanged when unavailable."""
    values = scores(question, [key(r) for r in rows])
    if values is None:
        return rows
    ranked = sorted(range(len(values)), key=lambda i: (-values[i], i))
    return [rows[i] for i in ranked] + rows[len(values):]