from dataclasses import dataclass
from typing import List, Sequence

from langchain_core.documents import Document
from sentence_transformers import CrossEncoder


@dataclass
class RerankedResult:
    document: Document
    score: float


class CrossEncoderReranker:
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2") -> None:
        self.model_name = model_name
        self.model = CrossEncoder(model_name)

    def rerank(
        self,
        query: str,
        documents: Sequence[Document],
        top_k: int = 5,
    ) -> List[RerankedResult]:
        if not query.strip():
            raise ValueError("Query cannot be empty.")
        if not documents:
            return []
        if top_k <= 0:
            raise ValueError("top_k must be greater than 0.")

        pairs = [(query, doc.page_content) for doc in documents]
        scores = self.model.predict(pairs)

        ranked = sorted(
            [
                RerankedResult(document=doc, score=float(score))
                for doc, score in zip(documents, scores)
            ],
            key=lambda item: item.score,
            reverse=True,
        )

        return ranked[:top_k]
