"""Visual reranker: Qwen3-VL-Reranker-2B (Apache-2.0) scores (query, page
image) pairs and reorders the fused top pages. Query-time only (no rebuild).
Fail-open: on ANY error the input order is returned unchanged."""
from pathlib import Path

import config

_PROMPT = "Retrieve images or text relevant to the user's query."
_re = None


def get_visual_reranker():
    """Lazy singleton; returns None when loading fails (fail-open upstream)."""
    global _re
    if _re is False:
        return None
    if _re is None:
        try:
            from sentence_transformers import CrossEncoder
            _re = CrossEncoder(config.VISUAL_RERANK_MODEL, device=config.DEVICE)
        except Exception:
            _re = False
    return _re or None


def _png(sha, page):
    import ingest_lib
    return ingest_lib.PAGES_DIR / str(sha) / f"p{int(page):04d}.png"


def rerank_keys(query, keys):
    """keys: [(sha, page)] -> reordered list. Entries without an existing
    page PNG keep their relative order at the end."""
    if not keys:
        return keys
    re = get_visual_reranker()
    if re is None:
        return keys
    try:
        pairs = []
        scored_idx = []
        rest = []
        for i, (sha, page) in enumerate(keys):
            png = None
            if sha and page is not None:
                try:
                    png = _png(sha, page)
                except Exception:
                    png = None
            if png is not None and Path(png).exists():
                pairs.append((query, {"image": str(png)}))
                scored_idx.append(i)
            else:
                rest.append((i, (sha, page)))
        if not pairs:
            return keys
        try:
            scores = re.predict(pairs, prompt=_PROMPT)
        except TypeError:
            scores = re.predict(pairs)
        order = sorted(zip(scored_idx, [float(s) for s in scores]),
                       key=lambda t: -t[1])
        out = [keys[i] for i, _ in order]
        out.extend(k for _, k in rest)
        return out
    except Exception:
        return keys