from typing import List
from ..interface.base_retriever import BaseRetriever
from ..interface.base_datastore import BaseDatastore
from sentence_transformers import CrossEncoder
import numpy as np

class Retriever(BaseRetriever):
    
    def __init__(self, datastore: BaseDatastore):
        self.datastore = datastore
        self.model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
        
    def search_retriever(self, query: str, top_k: int = 5) -> List[str]:
        """Cherche dans le datastore et classe les résultats par recherche sémantique

        Args:
            query (str): Requête
            top_k (int, optional): Nombre de résultats à retourner. Defaults to 5.

        Returns:
            List[str]: Liste de content de docs classée
        """
        search_results = self.datastore.search_datastore(query, top_k = top_k *5)
        reranked_results = self._rerank(query, search_results, top_k)
        return reranked_results
        
    
    def _rerank(self, query: str, search_results: List[str], top_k: int=10) -> List[str]:
        """Rerank le contenu des documents en fonction de la similarité avec la query

        Args:
            query (str): requête
            search_results (List[str]): liste de documents retrieved
            top_k (int, optional): Nombre de documents à retourner. Defaults to 10.

        Returns:
            List[str]: Liste de documents classée
        """
        pairs =[[query, doc]for doc in search_results]
        scores = self.model.predict(pairs)
        ranked_indices = np.argsort(scores)[::-1]
        results = [search_results[idx] for idx in ranked_indices[:top_k]]
        return results