from typing import List, Dict
import numpy as np
from sentence_transformers import CrossEncoder

class RerankerService:
    def __init__(self, model_name: str) -> None:
        self.model = CrossEncoder(model_name)

    def rerank(self, query: str, chunks: List[Dict[str, str]], top_n: int = 4) -> List[Dict[str, str]]:
        if not chunks:
            return []
        
        # Prepare pairs for cross-encoder
        pairs = [[query, chunk["text"]] for chunk in chunks]
        
        # Get scores
        scores = self.model.predict(pairs)
        
        # Add scores to chunks and sort
        for i, chunk in enumerate(chunks):
            chunk["rerank_score"] = float(scores[i])
            
        # Sort by rerank_score descending
        sorted_chunks = sorted(chunks, key=lambda x: x["rerank_score"], reverse=True)
        
        return sorted_chunks[:top_n]