#!/usr/bin/env python3
"""rag_reranker.py — Re-ranking con cross-encoder"""
from sentence_transformers import CrossEncoder

class Reranker:
    def __init__(self, model_name="BAAI/bge-reranker-v2-m3", use_fp16=True):
        self.model = CrossEncoder(model_name, max_length=512, device="cpu")
    
    def rerank(self, query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:
        pairs = [(query, c["contenido"]) for c in candidates]
        scores = self.model.predict(pairs)
        for i, score in enumerate(scores):
            candidates[i]["rerank_score"] = float(score)
        candidates.sort(key=lambda x: x["rerank_score"], reverse=True)
        return candidates[:top_k]

def search_with_rerank(query: str, alpha: float = 0.4,
                       top_k_hybrid: int = 20, top_k_final: int = 5):
    from rag_hybrid import hybrid_search
    candidates = hybrid_search(query, alpha=alpha, top_k=top_k_hybrid)
    candidates_dict = [{"contenido": c[1], "score": c[0]} for c in candidates]
    reranker = Reranker()
    return reranker.rerank(query, candidates_dict, top_k=top_k_final)