import os
import logging
from typing import Dict, List, Optional, Tuple, Union, Any
import tempfile

# LangChain imports
from langchain_community.document_loaders import PyPDFLoader, UnstructuredPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders.pdf import PDFMinerLoader
from langchain_community.document_loaders import PyPDFDirectoryLoader

# Image processing
import pytesseract
from PIL import Image
pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
# Table extraction
import camelot
import pandas as pd

# For unstructured data
from unstructured.partition.pdf import partition_pdf
# from unstructured.partition.auto import partition
os.environ['OCR_AGENT']=r'C:\Program Files\Tesseract-OCR\tesseract.exe'
# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class AdvancedPDFProcessor:
    """
    Classe pour traiter des documents PDF avec extraction avancée de texte, images et tableaux
    en utilisant LangChain et d'autres bibliothèques modernes.
    """
    
    def __init__(self, 
                 ocr_enabled: bool = True,
                 extract_tables: bool = True,
                 extract_images: bool = True,
                 chunk_size: int = 1000,
                 chunk_overlap: int = 200):
        """
        Initialise le processeur PDF avec les options configurées.
        
        Args:
            ocr_enabled: Si True, applique l'OCR sur les images détectées
            extract_tables: Si True, tente d'extraire les tableaux
            extract_images: Si True, extrait les images du PDF
            chunk_size: Taille des chunks pour la division du texte
            chunk_overlap: Chevauchement entre les chunks
        """
        self.ocr_enabled = ocr_enabled
        self.extract_tables = extract_tables
        self.extract_images = extract_images
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        
        # Configurer pytesseract si OCR est activé
        if ocr_enabled:
            # Chemin vers l'exécutable Tesseract si nécessaire
            pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
            pass
    
    def process_pdf(self, pdf_path: str) -> Dict[str, Any]:
        """
        Traite un fichier PDF et extrait son contenu de manière structurée.
        
        Args:
            pdf_path: Chemin vers le fichier PDF
            
        Returns:
            Dictionnaire contenant le texte extrait, les images, les tableaux et métadonnées
        """
        logger.info(f"Début du traitement du fichier PDF: {pdf_path}")
        
        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"Le fichier {pdf_path} n'existe pas")
        
        result = {
            "text": [],
            "chunks": [],
            "tables": [],
            "images": [],
            "metadata": {
                "filename": os.path.basename(pdf_path),
                "path": pdf_path,
                "size_bytes": os.path.getsize(pdf_path),
            }
        }
        
        # 1. Extraction de texte avec différentes méthodes pour maximiser la couverture
        result["text"] = self._extract_text(pdf_path)
        
        # 2. Chunking du texte pour une meilleure gestion par les LLMs
        result["chunks"] = self._chunk_text(result["text"])
        
        # 3. Extraction des tableaux si activée
        if self.extract_tables:
            result["tables"] = self._extract_tables(pdf_path)
        
        # 4. Extraction et analyse des images si activée
        if self.extract_images:
            result["images"] = self._extract_images(pdf_path)
        
        logger.info("Traitement du PDF terminé: %d chunks, %d tableaux, %d images",
                    len(result['chunks']), len(result['tables']), len(result['images']))
        
        return result
    
    def _extract_text(self, pdf_path: str) -> str:
        """Extrait le texte du PDF en utilisant plusieurs méthodes pour une couverture maximale."""
        text_content = ""
        
        # Méthode 1: PyPDFLoader de LangChain
        try:
            logger.info("Extraction de texte avec PyPDFLoader")
            loader = PyPDFLoader(pdf_path)
            documents = loader.load()
            text_content += "\n".join([doc.page_content for doc in documents])
        except Exception as e:
            logger.warning(f"Erreur avec PyPDFLoader: {e}")
        
        # Méthode 2: Utiliser PDFMinerLoader pour une extraction plus détaillée
        try:
            logger.info("Extraction de texte avec PDFMinerLoader")
            miner_loader = PDFMinerLoader(pdf_path)
            miner_docs = miner_loader.load()
            if not text_content:  # Si la première méthode a échoué
                text_content = "\n".join([doc.page_content for doc in miner_docs])
        except Exception as e:
            logger.warning(f"Erreur avec PDFMinerLoader: {e}")
        
        # Méthode 3: Utiliser Unstructured pour une extraction plus avancée
        try:
            logger.info("Extraction de texte avec Unstructured")
            elements = partition_pdf(pdf_path, extract_images_in_pdf=False, infer_table_structure=False)
            unstructured_text = "\n".join([str(element) for element in elements])
            
            # Si les méthodes précédentes n'ont rien donné ou si Unstructured a trouvé plus de contenu
            if not text_content or len(unstructured_text) > len(text_content):
                text_content = unstructured_text
        except Exception as e:
            logger.warning(f"Erreur avec Unstructured: {e}")
        
        return text_content
    
    def _chunk_text(self, text: str) -> List[str]:
        """Divise le texte en chunks pour un meilleur traitement."""
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            length_function=len,
        )
        chunks = text_splitter.split_text(text)
        return chunks
    
    def _extract_tables(self, pdf_path: str) -> List[Dict[str, Union[str, pd.DataFrame]]]:
        """Extrait les tableaux du PDF en utilisant Camelot."""
        tables_data = []
        
        try:
            logger.info("Extraction des tableaux avec Camelot")
            # Utiliser stream pour les tableaux avec des lignes claires et lattice pour les tableaux avec des bordures
            tables_stream = camelot.read_pdf(pdf_path, pages='all', flavor='stream')
            tables_lattice = camelot.read_pdf(pdf_path, pages='all', flavor='lattice')
            
            # Traiter les tableaux de type 'stream'
            for i, table in enumerate(tables_stream):
                if table.df.size > 0:  # Vérifier que le tableau contient des données
                    tables_data.append({
                        "page": table.page,
                        "type": "stream",
                        "data": table.df,
                        "accuracy": table.accuracy,
                        "description": f"Table {i+1} (Stream) de la page {table.page}"
                    })
            
            # Traiter les tableaux de type 'lattice'
            for i, table in enumerate(tables_lattice):
                if table.df.size > 0:
                    tables_data.append({
                        "page": table.page,
                        "type": "lattice",
                        "data": table.df,
                        "accuracy": table.accuracy,
                        "description": f"Table {i+1} (Lattice) de la page {table.page}"
                    })
                    
        except Exception as e:
            logger.warning(f"Erreur lors de l'extraction des tableaux: {e}")
        
        return tables_data
    
    def _extract_images(self, pdf_path: str) -> List[Dict[str, Any]]:
        """Extrait et analyse les images du PDF."""
        images_data = []
        
        try:
            logger.info("Extraction des images avec Unstructured")
            with tempfile.TemporaryDirectory() as temp_dir:
                elements = partition_pdf(
                    pdf_path, 
                    extract_images_in_pdf=True,
                    images_output_dir=temp_dir
                )
                
                # Collecter les chemins des images extraites
                image_elements = [el for el in elements if hasattr(el, 'image_path') and el.image_path]
                
                for i, img_element in enumerate(image_elements):
                    img_path = img_element.image_path
                    img_data = {
                        "page": getattr(img_element, 'page_number', None),
                        "path": img_path,
                        "position": getattr(img_element, 'coordinates', None),
                        "text": None  # Sera rempli par OCR si activé
                    }
                    
                    # Appliquer OCR si activé
                    if self.ocr_enabled and img_path and os.path.exists(img_path):
                        try:
                            logger.info("Extraction des textes avec ocr avec Unstructured")
                            img = Image.open(img_path)
                            ocr_text = pytesseract.image_to_string(img)
                            img_data["text"] = ocr_text.strip()
                        except Exception as e:
                            logger.warning(f"Erreur OCR sur l'image {i}: {e}")
                    
                    images_data.append(img_data)
                    
        except Exception as e:
            logger.warning(f"Erreur lors de l'extraction des images: {e}")
        
        return images_data


def process_pdf_document(pdf_path: str, **kwargs) -> Dict[str, Any]:
    """
    Fonction utilitaire pour traiter un document PDF avec des options configurables.
    
    Args:
        pdf_path: Chemin vers le fichier PDF à traiter
        **kwargs: Options de configuration pour le processeur PDF
        
    Returns:
        Dictionnaire contenant les données extraites du PDF
    """
    processor = AdvancedPDFProcessor(**kwargs)
    return processor.process_pdf(pdf_path)

def process_pdf_with_unstructured_loader(pdf_path: str, **kwargs) -> Dict[str, Any]:
    """
    Fonction qui utilise spécifiquement UnstructuredPDFLoader de LangChain
    pour extraire le contenu d'un PDF.
    
    Args:
        pdf_path: Chemin vers le fichier PDF à traiter
        **kwargs: Options supplémentaires à passer au loader
        
    Returns:
        Dictionnaire contenant le texte extrait et autres données
    """
    logger.info(f"Traitement du PDF avec UnstructuredPDFLoader: {pdf_path}")
    
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"Le fichier {pdf_path} n'existe pas")
    
    result = {
        "text": "",
        "chunks": [],
        "metadata": {
            "filename": os.path.basename(pdf_path),
            "path": pdf_path,
            "size_bytes": os.path.getsize(pdf_path),
        },
        "elements": []
    }
    
    try:
        # Configuration du loader avec les options avancées
        loader = UnstructuredPDFLoader(
            pdf_path,
            mode="elements",  # Pour obtenir une extraction structurée par éléments
            strategy="fast",
            **kwargs
        )
        
        # Chargement et extraction du contenu
        documents = loader.load()
        
        # Extraire le texte brut
        result["text"] = "\n".join([doc.page_content for doc in documents])
        
        # Stocker les documents individuels avec leurs métadonnées
        result["elements"] = [
            {
                "content": doc.page_content,
                "metadata": doc.metadata
            } for doc in documents
        ]
        
        # Chunking du texte si nécessaire
        chunk_size = kwargs.get("chunk_size", 1000)
        chunk_overlap = kwargs.get("chunk_overlap", 200)
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            length_function=len,
        )
        result["chunks"] = text_splitter.split_text(result["text"])
        
        logger.info(f"UnstructuredPDFLoader: extrait {len(documents)} éléments et {len(result['chunks'])} chunks")
        
    except Exception as e:
        logger.warning(f"Erreur lors du traitement avec UnstructuredPDFLoader: {e}")
    
    return result