from __future__ import annotations

from typing import List
import pandas as pd

from schemas.candidates import RankedItem, RankedList, CandidateSet


def rerank_candidates(plan, candidate_set: CandidateSet, reranker, catalog_df: pd.DataFrame, k: int = 10) -> RankedList:
    """Rerank candidate union using cross-encoder reranker."""
    catalog = catalog_df.set_index("assessment_id")
    scored: List[RankedItem] = []
    for cand in candidate_set.candidates:
        if cand.assessment_id not in catalog.index:
            continue
        doc_text = catalog.loc[cand.assessment_id].get("doc_text", "")
        score = reranker.score(plan.rerank_query, doc_text)
        scored.append(RankedItem(assessment_id=cand.assessment_id, score=float(score)))
    scored.sort(key=lambda x: x.score, reverse=True)
    return RankedList(items=scored[:k])