from typing import List
from ..interface.base_datastore import BaseDatastore, DataItem
import lancedb
from lancedb.table import Table
from typing import List
import pyarrow as pa
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
# from concurrent.futures import ThreadPoolExecutor


class Datastore(BaseDatastore):
    
    DB_PATH = "data/sample-lancedb"
    DB_TABLE_NAME = "rag-table"
    
    def __init__(self):
        """Constructeur par défaut, initialise les dimensions des vecteurs pour l'embedding
        (actuellement 384 par défaut pour le modèle all-MiniLm-L6-v2), charge le modèle SentenceTransformer,
        connecte la base de données et récupère la table.
        #Model's maximum sequence length = 256
        """
        self.vector_dimensions = 384 # all-MiniLm-L6-v2 a une dimension fixe de 384
        self.model = SentenceTransformer("all-MiniLM-L6-v2")
        self.vector_db = lancedb.connect(self.DB_PATH)
        self.table: Table = self._get_table()
        
        
    
    def reset_table(self) -> Table:
        """Drop la table si elle existe puis crée une table selon le schéma
        vector (liste de float32 de dimension définie dans la classe), content et source
        et l'ouvre.
        
        Returns:
            Table: La table crée
        """
        try:
            self.vector_db.drop_table(self.DB_TABLE_NAME)
        except Exception as e:
            print("Unable to drop the table, assuming it does not exist.")
        
        schema = pa.schema(
            [
                pa.field("vector", pa.list_(pa.float32(),self.vector_dimensions)),
                pa.field("content", pa.utf8()),
                pa.field("source", pa.utf8()),
            ]
        )
        
        self.table = self.vector_db.create_table(self.DB_TABLE_NAME, schema = schema)
        #self.table = self.vector_db.open_table(self.DB_TABLE_NAME)
        print(f"Table was reset/created: {self.DB_TABLE_NAME} in {self.DB_PATH}")
        return self.table
            
        
    def _get_table(self) -> Table:
        """Ouvre la table ou la reset en cas d'échec

        Returns:
            Table: la table ouverte
        """
        try:
            return self.vector_db.open_table(self.DB_TABLE_NAME)
        except Exception as e:
            print(f"Error opening the table {e}. Trying to reset it.")
            return self.reset_table()
    
    def add_items(self, items: List[DataItem]) -> None:
        """Ajoute les items en entrées dans le dataset (nécessite un embedding)
        opération network bound donc à paralléliser
        Args:
            items (List[DataItem]): Liste de DataItems à ajouter 

        Returns:
            _type_: _description_
        """
        if not items:
            return
        
        contents = [item.content for item in items]
        sources = [item.source for item in items]
        
        "embedding du contenu de chaque entrée par batch de 32"
        print(f"Génération des embeddings pour {len(items)} items...")
        vectors = self.model.encode(
            contents,
            batch_size= 32,
        )
        
        "conversion en dictionnaires pour stocker les documents dans la BDD"
        entries = [
            {
                "vector": vector,
                "content": content,
                "source": source
            }
            for vector, content, source in zip(vectors, contents, sources)
        ]
        
        #self.table.merge_insert("source").when_matched_update_all().when_not_matched_insert_all().execute(entries)
        self.table.add(entries)
        print(f"{len(entries)} items ajoutés")
        
    #deprecated
    def _convert_items_to_entry(self, item: DataItem) -> dict:
        """Convertir un DataItem en dictionnaire correspondant au schéma du Datastore

        Args:
            item (DataItem): item à convertir

        Returns:
            dict: Dictionnaire contenant le vecteur, le contenu et la source 
        """
        
        vector = self.create_vector(item.content)
        
        return{
            "vector": vector,
            "content": item.content,
            "source": item.source,
        }
        
    def create_vector(self, content: str) -> List[float]:
        """Utilise le modèle d'embedding pour convertir le str en vecteur (list[float]) 

        Args:
            content (str): contenu de l'entrée à vectoriser

        Returns:
            List[float]: vecteur renvoyé
        """
        response = self.model.encode(content)
        return response.tolist()
    
    def search_datastore(self, query: str, top_k: int = 5) -> List[str]:
        """Embedde la query et lance une recherche

        Args:
            query (str): requête
            top_k (int, optional): Nombre de documents à retourner au maximum. Defaults to 5.

        Returns:
            List[str]: Tableau contenant le contenu des documents retrouvés
        """                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                             
        
        vector = self.model.encode(query)
        results = (
            self.table
            .search(vector, vector_column_name="vector")
            .select(["content", "source"])
            .limit(top_k)
            .to_list()
        )
        
        result_content = [
            result["content"] 
            for result in results 
            if "content" in result and result["content"] is not None]
        return result_content
    
    
