from typing import List
from src.interface.base_datastore import DataItem
from src.interface.base_indexer import BaseIndexer
from docling.document_converter import DocumentConverter 
from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
import os
from docling_core.transforms.chunker.tokenizer.huggingface import HuggingFaceTokenizer
from transformers import AutoTokenizer
from langchain_text_splitters import RecursiveCharacterTextSplitter
# pip install langchain langchain-text-splitters


EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
MAX_TOKENS = 256  # set to a small number for illustrative purposes
#précedemment 512

class Indexer(BaseIndexer):
    
    def __init__(self):
        
        self.converter = DocumentConverter()
        self.tokenizer = HuggingFaceTokenizer(
            tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_ID),
            max_tokens = MAX_TOKENS,  # optional, by default derived from `tokenizer` for HF case
        )
        self.chunker = HybridChunker(
            tokenizer = self.tokenizer,
            max_tokens = MAX_TOKENS,
            #merge_peers = True,
            #handle_tables = "separate",
            #handle_pictures = "separate"
        )
        """self.text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
            tokenizer = self.tokenizer,
            chunk_size = MAX_TOKENS,
            chunk_overlap = MAX_TOKENS // 10,
            separators=["\n\n", "\n", ". ", " ", ""],
        )"""

    def index(self, document_paths: List[str]) -> List[DataItem]:
        """Convertit les documents en format docling puis les découpe en morceaux.
        Les morceaux sont ensuite convertis sous forme de DataItem en y ajoutant des métadonnées

        Args:
            document_paths (List[str]): Liste des documents à indexer

        Returns:
            List[DataItem]: Liste des DataItems indexés 
        """
        items = []
        for document_path in document_paths:
            try:
                document = self.converter.convert(document_path).document
                chunks = list(self.chunker.chunk(document))
                #chunks = self.text_splitter.split_text(document)
                
                item = self._convert_to_DataItem(chunks, document_path)
                
                items.extend(item)

            except Exception as e:
                print(f"Erreur lors du traitement de {document_path}: {e}")
                continue
            
        return items
    
    def _convert_to_DataItem(self, chunks, document_path: str) -> List[DataItem]:
        """Crée une liste de DataItems en ajoutant les éventuels headers au début du contenu 
        et en récupérant la source à l'aide du module os

        Args:
            chunks : liste de chunks de documents à traiter
            document_path: chemin vers le document, permet d'utiliser os pour récupérer directement le nom

        Returns:
            Lits[DataItem]: liste de DataItem contenant les métadonnées et séparant le contenu et la source
        """
        items = []
        for i, chunk in enumerate(chunks):
            
            try:
                headings = ""
                
                if(hasattr(chunk, 'meta') and hasattr(chunk.meta, 'headings') and chunk.meta.headings):
                    headings = "## " + ", ".join(chunk.meta.headings) + "\n"
                
                text = chunk.text if(hasattr(chunk, 'text')) else str(chunk)
                    
                content = f"{headings}{text}"
                
                filename = os.path.basename(document_path)
                source = f"{filename}:chunk {i}"
                
                item = DataItem(content = content, source = source)
                items.append(item)
            
            except Exception as e:
                print(f" Erreur sur le chunk {i}: {e}")
                continue
           
        return items