from llama_index.tools import FunctionTool
from llama_index.embeddings.base import BaseEmbedding
from llama_index.schema import TextNode
import numpy as np

class ToolRouter:
    """
    Paramètres
    --------------
    tools : list[FunctionTool]
        Liste des outils disponibles, chacun avec une description dans ses métadonnées.
    embed_model : BaseEmbedding
        Modèle d'embedding utilisé pour vectoriser les descriptions et les requêtes.
    """
    def __init__(self, tools: list[FunctionTool], embed_model: BaseEmbedding):
        self.tools = tools
        self.embed_model = embed_model
        self.tool_nodes = [
            TextNode(
                text=tool.metadata.description or tool.name,
                metadata={"name": tool.metadata.name}
            )
            for tool in tools
        ]
        self.embeddings = self.embed_model.get_text_embedding_batch([n.text for n in self.tool_nodes])


    def route(self, query: str, top_k: int = 3, min_score: float = 0.5):
    """
        Sélectionne dynamiquement les outils les plus pertinents en fonction d'une requête utilisateur,
        en utilisant la similarité cosinus entre l'embedding de la requête et ceux des outils disponibles.

        Parameters
        ----------
        query : str
            La requête utilisateur ou le "thought" à router vers les outils les plus adaptés.
        top_k : int, optional
            Nombre maximum d'outils à retourner, triés par score décroissant (par défaut : 3).
        min_score : float, optional
            Seuil minimal de similarité cosinus pour qu'un outil soit retenu (par défaut : 0.5).

        Returns
        -------
        List[Tuple[FunctionTool, float]]
            Une liste triée des outils pertinents, chacun accompagné de son score de similarité.

        Usage
        -----
            selected_tools = router.route("extraire les entités d'un texte")
        """

        query_vec = self.embed_model.get_text_embedding(query)
        scored = [
            (tool, float(np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec))))
            for tool, vec in zip(self.tools, self.embeddings)
        ]
        filtered = [(tool, score) for tool, score in scored if score >= min_score]
        return sorted(filtered, key=lambda x: x[1], reverse=True)[:top_k]