backend_processor.py

"""
DataClean Pro - Module de traitement avancé
Gestion batch, files d'attente, et traitement parallèle
"""

import os
import json
import time
import logging
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional, Any
from enum import Enum
from datetime import datetime
import hashlib

import google.generativeai as genai
import pandas as pd
import PyPDF2

logger = logging.getLogger(__name__)


class JobStatus(Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"
    PARTIAL = "partial"


class ExtractionType(Enum):
    FACTURE = "facture"
    PRODUIT = "produit"
    CONTACT = "contact"
    GENERIQUE = "generique"


@dataclass
class FileResult:
    filename: str
    status: str
    data: Optional[Dict] = None
    error: Optional[str] = None
    processing_time: float = 0.0


@dataclass
class JobResult:
    job_id: str
    status: JobStatus
    total_files: int
    processed_files: int
    successful_files: int
    failed_files: int
    results: List[FileResult]
    output_file: Optional[str] = None
    created_at: str = ""
    completed_at: Optional[str] = None
    total_processing_time: float = 0.0


class DataProcessor:
    """
    Processeur de données avancé avec support batch et parallélisation
    """
    
    def __init__(self, api_key: str, max_workers: int = 4):
        self.api_key = api_key
        self.max_workers = max_workers
        self.model = None
        self._initialize_model()
        
        # Cache pour éviter de retraiter les mêmes fichiers
        self.cache = {}
        
    def _initialize_model(self):
        """Initialise le modèle Gemini"""
        genai.configure(api_key=self.api_key)
        self.model = genai.GenerativeModel('gemini-1.5-pro')
        logger.info("Modèle Gemini initialisé")
    
    def _get_file_hash(self, file_path: Path) -> str:
        """Calcule le hash d'un fichier pour le cache"""
        hasher = hashlib.md5()
        with open(file_path, 'rb') as f:
            buf = f.read(65536)
            while len(buf) > 0:
                hasher.update(buf)
                buf = f.read(65536)
        return hasher.hexdigest()
    
    def _extract_content(self, file_path: Path) -> str:
        """Extrait le contenu d'un fichier"""
        ext = file_path.suffix.lower()
        
        try:
            if ext == '.pdf':
                return self._extract_pdf(file_path)
            elif ext in ['.xlsx', '.xls']:
                return self._extract_excel(file_path)
            elif ext == '.csv':
                return self._extract_csv(file_path)
            elif ext in ['.txt', '.json']:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    return f.read()
        except Exception as e:
            logger.error(f"Erreur extraction {file_path}: {e}")
        
        return ""
    
    def _extract_pdf(self, file_path: Path) -> str:
        """Extrait le texte d'un PDF"""
        text = ""
        with open(file_path, 'rb') as file:
            reader = PyPDF2.PdfReader(file)
            for page in reader.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n"
        return text
    
    def _extract_excel(self, file_path: Path) -> str:
        """Extrait les données d'un Excel"""
        df = pd.read_excel(file_path)
        return df.to_string()
    
    def _extract_csv(self, file_path: Path) -> str:
        """Extrait les données d'un CSV"""
        for encoding in ['utf-8', 'latin-1', 'cp1252']:
            try:
                df = pd.read_csv(file_path, encoding=encoding)
                return df.to_string()
            except:
                continue
        return ""
    
    def _get_prompt(self, extraction_type: ExtractionType) -> str:
        """Retourne le prompt approprié pour le type d'extraction"""
        prompts = {
            ExtractionType.FACTURE: """
Analyse cette facture et extrais les données en JSON:
{
    "numero_facture": "",
    "date_facture": "YYYY-MM-DD",
    "fournisseur": {"nom": "", "adresse": "", "siret": ""},
    "client": {"nom": "", "adresse": ""},
    "lignes": [{"description": "", "quantite": 0, "prix_unitaire": 0, "total": 0}],
    "montant_ht": 0,
    "tva": {"taux": 0, "montant": 0},
    "montant_ttc": 0
}

CONTENU:
{content}

Réponds UNIQUEMENT avec le JSON, sans explications.
""",
            ExtractionType.PRODUIT: """
Analyse ces données produit et extrais en JSON:
{
    "reference": "",
    "nom": "",
    "description": "",
    "categorie": "",
    "prix_ht": 0,
    "prix_ttc": 0,
    "stock": 0,
    "marque": "",
    "ean": ""
}

CONTENU:
{content}

Réponds UNIQUEMENT avec le JSON, sans explications.
""",
            ExtractionType.CONTACT: """
Analyse ces données de contact et extrais en JSON:
{
    "nom": "",
    "prenom": "",
    "email": "",
    "telephone": "",
    "entreprise": "",
    "poste": "",
    "adresse": "",
    "ville": "",
    "code_postal": "",
    "pays": ""
}

CONTENU:
{content}

Réponds UNIQUEMENT avec le JSON, sans explications.
""",
            ExtractionType.GENERIQUE: """
Analyse ce document et extrais toutes les données structurées en JSON.
Identifie le type de document et organise les informations de manière logique.
Normalise les dates en YYYY-MM-DD et les montants en nombres.

CONTENU:
{content}

Réponds UNIQUEMENT avec le JSON, sans explications.
"""
        }
        return prompts.get(extraction_type, prompts[ExtractionType.GENERIQUE])
    
    def _process_with_gemini(self, content: str, extraction_type: ExtractionType) -> Dict[str, Any]:
        """Traite le contenu avec Gemini AI"""
        prompt = self._get_prompt(extraction_type).format(content=content[:50000])
        
        try:
            response = self.model.generate_content(prompt)
            response_text = response.text.strip()
            
            # Nettoyer les balises markdown
            if response_text.startswith('```'):
                lines = response_text.split('\n')
                response_text = '\n'.join(lines[1:-1] if lines[-1] == '```' else lines[1:])
            
            data = json.loads(response_text)
            return {"success": True, "data": data}
            
        except json.JSONDecodeError as e:
            return {"success": False, "error": f"JSON invalide: {e}"}
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def process_file(self, file_path: Path, extraction_type: ExtractionType) -> FileResult:
        """Traite un seul fichier"""
        start_time = time.time()
        
        # Vérifier le cache
        file_hash = self._get_file_hash(file_path)
        cache_key = f"{file_hash}_{extraction_type.value}"
        
        if cache_key in self.cache:
            logger.info(f"Cache hit pour {file_path.name}")
            cached = self.cache[cache_key]
            return FileResult(
                filename=file_path.name,
                status="success",
                data=cached,
                processing_time=0.0
            )
        
        # Extraire le contenu
        content = self._extract_content(file_path)
        
        if not content:
            return FileResult(
                filename=file_path.name,
                status="error",
                error="Impossible d'extraire le contenu",
                processing_time=time.time() - start_time
            )
        
        # Traiter avec Gemini
        result = self._process_with_gemini(content, extraction_type)
        processing_time = time.time() - start_time
        
        if result["success"]:
            # Mettre en cache
            self.cache[cache_key] = result["data"]
            
            return FileResult(
                filename=file_path.name,
                status="success",
                data=result["data"],
                processing_time=processing_time
            )
        else:
            return FileResult(
                filename=file_path.name,
                status="error",
                error=result["error"],
                processing_time=processing_time
            )
    
    def process_batch(
        self, 
        job_id: str,
        file_paths: List[Path], 
        extraction_type: ExtractionType,
        output_folder: Path
    ) -> JobResult:
        """
        Traite un lot de fichiers en parallèle
        """
        start_time = time.time()
        results: List[FileResult] = []
        
        logger.info(f"Démarrage batch {job_id}: {len(file_paths)} fichiers")
        
        # Traitement parallèle
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_file = {
                executor.submit(self.process_file, fp, extraction_type): fp 
                for fp in file_paths
            }
            
            for future in as_completed(future_to_file):
                file_path = future_to_file[future]
                try:
                    result = future.result()
                    results.append(result)
                    logger.info(f"Traité: {result.filename} - {result.status}")
                except Exception as e:
                    results.append(FileResult(
                        filename=file_path.name,
                        status="error",
                        error=str(e)
                    ))
        
        # Statistiques
        successful = [r for r in results if r.status == "success"]
        failed = [r for r in results if r.status == "error"]
        
        # Créer l'Excel de sortie
        output_file = None
        if successful:
            output_file = self._create_output_excel(
                job_id, 
                [r.data for r in successful],
                [r.filename for r in successful],
                output_folder
            )
        
        # Déterminer le statut final
        if len(failed) == 0:
            status = JobStatus.COMPLETED
        elif len(successful) == 0:
            status = JobStatus.FAILED
        else:
            status = JobStatus.PARTIAL
        
        total_time = time.time() - start_time
        
        return JobResult(
            job_id=job_id,
            status=status,
            total_files=len(file_paths),
            processed_files=len(results),
            successful_files=len(successful),
            failed_files=len(failed),
            results=results,
            output_file=output_file,
            created_at=datetime.now().isoformat(),
            completed_at=datetime.now().isoformat(),
            total_processing_time=total_time
        )
    
    def _create_output_excel(
        self, 
        job_id: str, 
        data_list: List[Dict],
        filenames: List[str],
        output_folder: Path
    ) -> Optional[str]:
        """Crée le fichier Excel de sortie"""
        try:
            # Ajouter le nom du fichier source à chaque entrée
            for data, filename in zip(data_list, filenames):
                if isinstance(data, dict):
                    data['_source'] = filename
            
            # Aplatir les données imbriquées si nécessaire
            flat_data = []
            for item in data_list:
                if isinstance(item, dict):
                    flat_item = self._flatten_dict(item)
                    flat_data.append(flat_item)
            
            df = pd.DataFrame(flat_data)
            
            output_filename = f"export_{job_id}.xlsx"
            output_path = output_folder / output_filename
            
            with pd.ExcelWriter(output_path, engine='openpyxl') as writer:
                df.to_excel(writer, index=False, sheet_name='Données')
                
                # Ajuster les colonnes
                worksheet = writer.sheets['Données']
                for column in worksheet.columns:
                    max_length = max(len(str(cell.value or '')) for cell in column)
                    worksheet.column_dimensions[column[0].column_letter].width = min(max_length + 2, 50)
            
            logger.info(f"Excel créé: {output_path}")
            return output_filename
            
        except Exception as e:
            logger.error(f"Erreur création Excel: {e}")
            return None
    
    def _flatten_dict(self, d: Dict, parent_key: str = '', sep: str = '_') -> Dict:
        """Aplatit un dictionnaire imbriqué"""
        items = []
        for k, v in d.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, dict):
                items.extend(self._flatten_dict(v, new_key, sep).items())
            elif isinstance(v, list):
                # Pour les listes, on les convertit en string JSON
                items.append((new_key, json.dumps(v, ensure_ascii=False)))
            else:
                items.append((new_key, v))
        return dict(items)


# Fonctions utilitaires pour tests
def test_processor():
    """Test du processeur"""
    api_key = os.environ.get('GEMINI_API_KEY')
    if not api_key:
        print("❌ GEMINI_API_KEY non définie")
        return
    
    processor = DataProcessor(api_key)
    
    # Test avec un texte exemple
    test_content = """
    FACTURE N° 2024-TEST-001
    Date: 20 janvier 2024
    
    FOURNISSEUR: Test Company SARL
    123 rue de Test, 75001 Paris
    
    MONTANT HT: 1000.00 EUR
    TVA 20%: 200.00 EUR
    TOTAL TTC: 1200.00 EUR
    """
    
    result = processor._process_with_gemini(test_content, ExtractionType.FACTURE)
    print(f"Résultat: {json.dumps(result, indent=2, ensure_ascii=False)}")


if __name__ == '__main__':
    test_processor()