from sentence_transformers import CrossEncoder

reranker = CrossEncoder(
    "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1",
    max_length=512,
    device="cpu",
)


def rerank(query: str, results: list[dict], k: int = 5) -> list[dict]:
    if not results:
        return []

    pairs = [(query, r["payload"]["full_text"]) for r in results]
    scores = reranker.predict(pairs)

    scored = []
    for i, score in enumerate(scores):
        scored.append({
            "ticket_id": results[i]["payload"]["ticket_id"],
            "score": float(score),
            "full_text": results[i]["payload"]["full_text"],
            "category": results[i]["payload"].get("category", ""),
        })

    scored.sort(key=lambda x: x["score"], reverse=True)
    return scored[:k]