from sentence_transformers import CrossEncoder

RERANK_MODEL = "BAAI/bge-reranker-v2-m3"
RERANK_THRESHOLD = 0.005

_model = None


def _get_model():
    global _model
    if _model is None:
        print(f"Loading reranker: {RERANK_MODEL}")
        _model = CrossEncoder(RERANK_MODEL, max_length=512)
    return _model


def rerank(query: str, hits: list[dict], top_k: int = 5) -> list[dict]:
    """Rerank hits bằng cross-encoder, loại bỏ kết quả dưới threshold."""
    if not hits:
        return []

    model = _get_model()
    pairs = [(query, h["text"]) for h in hits]
    scores = model.predict(pairs)

    for hit, score in zip(hits, scores):
        hit["rerank_score"] = float(score)

    ranked = sorted(hits, key=lambda x: x["rerank_score"], reverse=True)
    filtered = [h for h in ranked if h["rerank_score"] >= RERANK_THRESHOLD]

    return filtered[:top_k]