src/Algorithms/document_validator.py

import os
import tempfile
import json
import time
from docling.document_converter import DocumentConverter
import streamlit as st

from core.llm_client import call_llm, LLMUnavailableError

if "form_start_time" not in st.session_state:
    st.session_state.form_start_time = time.time()


# Signatures binaires (magic bytes) minimales pour verifier qu'un fichier
# televerse correspond reellement a son extension, pas seulement a son nom.
# Pas de dependance externe (python-magic) : juste assez de types pour
# couvrir les formats acceptes par st.file_uploader (pdf/png/jpg) ici.
_SIGNATURES_PAR_EXTENSION = {
    ".pdf": [b"%PDF-"],
    ".png": [b"\x89PNG\r\n\x1a\n"],
    ".jpg": [b"\xff\xd8\xff"],
    ".jpeg": [b"\xff\xd8\xff"],
}


def _signature_correspond_extension(file_bytes, ext) -> bool:
    """True si les premiers octets du fichier correspondent a une signature
    connue pour `ext`, ou si `ext` n'est pas dans la liste surveillee (on ne
    bloque alors rien, par securite contre un faux negatif sur un type non
    couvert)."""
    signatures = _SIGNATURES_PAR_EXTENSION.get(ext)
    if not signatures:
        return True
    return any(file_bytes.startswith(sig) for sig in signatures)

# ============================================================================
# CHAMPS ATTENDUS PAR TYPE DE DOCUMENT
# Le score n'est plus auto-déclaré par le LLM (il recopiait le 0.95 du
# template). Il est maintenant CALCULÉ en Python : pourcentage de champs
# attendus réellement remplis dans les données extraites.
# ============================================================================
CHAMPS_ATTENDUS = {
    "CNI_PASSEPORT": [
        "nom_identifie", "genre", "date_naissance", "type_piece_identite",
        "numero_piece_identite", "date_expiration_identite",
        "pays_naissance", "profession"
    ],
    "JUSTIFICATIF_ADRESSE": [
        "nom_identifie", "adresse_complete", "statut_logement",
        "quartier", "ville", "pays_residence"
    ],
    "BULLETIN_SALAIRE": [
        "nom_identifie", "statut_professionnel", "profession",
        "nom_employeur", "secteur_activite", "revenu_mensuel", "devise"
    ],
    "RIB": ["nom_identifie", "nom_banque"],
    # Relevés bancaires / Wave / Orange Money (champ 5 de l'upload KYC,
    # passés avec expected_type="RELEVE_COMPTE" — voir kyc_form.py).
    # Avant ce fix, ce type était absent d'ici -> fallback "pas de type
    # connu" -> comparé aux 31 champs de SCHEMA_EXTRACTION entier (score
    # quasi toujours < 70%, dossier jamais VERT). Même liste que RIB.
    "RELEVE_COMPTE": ["nom_identifie", "nom_banque"],
}

# Champs qu'un type de document n'a JAMAIS le droit d'émettre.
# Exemple réel du bug corrigé : le RIB mentionnait "Dakar" (ville de la
# banque) et écrasait "Aubervilliers" (ville du client). Maintenant bloqué
# à deux niveaux : dans le prompt ET dans le code de fusion (kyc_form.py).
CHAMPS_INTERDITS = {
    "CNI_PASSEPORT": [
        "quartier", "ville", "pays_residence", "statut_logement",
        "nom_banque", "revenu_mensuel", "devise", "nom_employeur",
        "secteur_activite", "origine_fonds"
    ],
    "JUSTIFICATIF_ADRESSE": [
        "numero_piece_identite", "type_piece_identite", "nom_banque",
        "revenu_mensuel", "devise", "date_naissance", "pays_naissance"
    ],
    "BULLETIN_SALAIRE": [
        "numero_piece_identite", "type_piece_identite", "statut_logement",
        "quartier", "nom_banque", "date_naissance"
    ],
    "RIB": [
        "adresse_complete", "quartier", "ville", "pays_residence",
        "statut_logement", "revenu_mensuel", "statut_professionnel",
        "profession", "nom_employeur", "secteur_activite",
        "numero_piece_identite", "type_piece_identite", "devise",
        "date_naissance", "pays_naissance"
    ],
}

# Schéma canonique enrichi — nouvelles clés : personnes_a_charge,
# anciennete_emploi_mois, anciennete_revenus_mois, charges_mensuelles,
# reseaux_sociaux, whatsapp, personne_urgence_*
SCHEMA_EXTRACTION = {
    "nom_identifie": "",
    "genre": "",
    "date_naissance": "",
    "adresse_complete": "",
    "telephone": "",
    "email": "",
    "etat_civil": "",
    "type_piece_identite": "",
    "numero_piece_identite": "",
    "date_expiration_identite": "",
    "statut_professionnel": "",
    "profession": "",
    "nom_employeur": "",
    "secteur_activite": "",
    "revenu_mensuel": "",
    "devise": "",
    "statut_logement": "",
    "quartier": "",
    "ville": "",
    "pays_residence": "",
    "pays_naissance": "",
    "nom_banque": "",
    "origine_fonds": "",
    "personnes_a_charge": "",
    "anciennete_emploi_mois": "",
    "anciennete_revenus_mois": "",
    "charges_mensuelles": "",
    "reseaux_sociaux": "",
    "whatsapp": "",
    "personne_urgence_nom": "",
    "personne_urgence_tel": "",
}

def _construire_prompt(expected_type, extracted_text):
    schema_json = json.dumps({"correspond_au_type": True, "donnees_extraites": SCHEMA_EXTRACTION},
                             ensure_ascii=False, indent=4)
    interdits = CHAMPS_INTERDITS.get(expected_type, [])
    attendus = CHAMPS_ATTENDUS.get(expected_type, [])

    return f"""Tu es un expert en extraction de données KYC. Analyse le texte markdown extrait de ce document (type attendu : {expected_type}).

TEXTE DU DOCUMENT :
{extracted_text}

INSTRUCTIONS STRICTES :
1. REMPLIS le maximum de champs du schéma JSON ci-dessous avec les informations réelles du document. Ta mission est de ne laisser AUCUN champ vide tant que l'information est présente dans le texte, même partiellement ou implicite.
2. CORRESPONDANCE INTELLIGENTE DES VALEURS :
   - "statut_professionnel" : normalise (ex: "Collaborateur en formation", "Stagiaire", "Apprenti(e)" -> "ETUDIANT" ; "CDI", "Employé", "Technicien" -> "SALARIE" ; titulaire d'une entreprise -> "ENTREPRENEUR").
   - "statut_logement" : un reçu de loyer implique "LOCATAIRE" ; un avis d'implication/attestation propriété implique "PROPRIETAIRE" ; "chez mes parents / chez un tiers" implique "HEBERGE".
   - "secteur_activite" : déduis-le de l'activité de l'employeur (ex: BNP Paribas -> "Banque").
   - "origine_fonds" : "Salaires" -> "Salaire", "Virement familial" -> "Contribution Familiale".
3. FORMATS DE SORTIE :
   - "date_naissance", "date_expiration_identite" : JJ.MM.AAAA (ex: 07.02.2002).
   - "revenu_mensuel", "charges_mensuelles" : nombre UNIQUEMENT, sans symbole ni espace, point décimal (ex: 2237.18). Si plusieurs montants, prends le NET À PAYER mensuel.
   - "devise" : code ISO 3 lettres (EUR, XOF, USD...) trouvé à côté du montant.
   - "telephone", "whatsapp" : format international complet (+221...).
   - "genre" : "M" ou "F" uniquement.
   - "personnes_a_charge", "anciennete_emploi_mois", "anciennete_revenus_mois" : entiers. Si une durée est en années, CONVERTIS en mois (2 ans = 24).
4. "nom_employeur" = l'émetteur du bulletin de paie. "nom_banque" = la banque du RIB (le nom commercial, ex: "UBA Senegal SA").
5. INTERDICTIONS ABSOLUES pour ce type de document — laisse TOUJOURS "" pour : {", ".join(interdits) if interdits else "aucune"}.
6. "reseaux_sociaux" : recopie TOUTE url de réseau social trouvée (Facebook, Instagram, LinkedIn, X, TikTok, WhatsApp), séparées par des espaces.
7. "correspond_au_type" : true si le document ressemble réellement à un {expected_type}, false sinon. Ne réponds JAMAIS avec des champs déjà remplis fictivement.
8. Ne modifie pas les noms de clés. Ne mets pas de commentaire dans le JSON.

Réponds UNIQUEMENT avec ce JSON valide :
{schema_json}"""

class DocumentValidator:
    """
    Module d'Extraction Documentaire (Docling + Gemini).
    Le score de cohérence est CALCULÉ (champs trouvés / champs attendus),
    plus auto-déclaré par le LLM.
    """

    _instance = None

    def __new__(cls):
        if cls._instance is None:
            instance = super(DocumentValidator, cls).__new__(cls)
            try:
                instance._initialize_models()
            except Exception:
                # Ne JAMAIS garder une instance à moitié initialisée (sans
                # doc_converter) : sinon chaque analyse plante en silence
                # jusqu'au redémarrage du process. On laisse _instance à
                # None pour qu'un prochain appel retente le chargement.
                cls._instance = None
                raise
            cls._instance = instance
        return cls._instance

    def _initialize_models(self):
        print("Initialisation du pipeline d'extraction (Docling + IA)...")
        print("-> Chargement de Docling...")
        self.doc_converter = DocumentConverter()
        print("Pipeline prêt !")

    def _calculer_score(self, expected_type, donnees):
        """Score réel = % de champs attendus réellement remplis."""
        attendus = CHAMPS_ATTENDUS.get(expected_type)
        if not attendus:
            # Document optionnel : score basé sur tous les champs du schéma
            attendus = list(SCHEMA_EXTRACTION.keys())
        trouves = [k for k in attendus if str(donnees.get(k, "")).strip()]
        pct = round(100 * len(trouves) / len(attendus), 2) if attendus else 0.0
        return pct, len(trouves), len(attendus)

    def process_document(self, file_bytes, file_name, expected_type):
        if len(file_bytes) > 15 * 1024 * 1024:
            return {"statut_global": "ERREUR_SYSTEME", "erreur": "Fichier trop volumineux (> 15MB)."}

        ext = os.path.splitext(file_name)[1].lower()
        if not ext: ext = ".pdf"

        # Verification du contenu reel (magic bytes), pas seulement de
        # l'extension du nom de fichier : st.file_uploader(type=[...]) ne
        # filtre que cote navigateur, un fichier renomme passe sans probleme.
        # Echoue vite ici plutot que de gaspiller un appel Docling + LLM sur
        # un fichier qui ne correspond pas a son extension.
        if not _signature_correspond_extension(file_bytes, ext):
            return {
                "statut_global": "ERREUR_SYSTEME",
                "erreur": f"Le contenu du fichier ne correspond pas a son extension ({ext}).",
            }

        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp_file:
            tmp_file.write(file_bytes)
            tmp_file_path = tmp_file.name

        try:
            # 1. Extraction Docling
            result = self.doc_converter.convert(tmp_file_path)
            extracted_text = result.document.export_to_markdown()

            # 2. Analyse Sémantique IA (prompt contextuel par type de doc)
            prompt = _construire_prompt(expected_type, extracted_text)

            try:
                response_text = call_llm([{"role": "user", "content": prompt}], json_mode=True).strip()
            except LLMUnavailableError as e:
                raise Exception(f"Echec total : tous les modèles de repli ont échoué. Détail : {e}")

            if response_text.startswith("```json"):
                response_text = response_text[7:-3].strip()
            elif response_text.startswith("```"):
                response_text = response_text[3:-3].strip()

            llm_result = json.loads(response_text)

            # 3. Nettoyage : ne garder QUE les clés du schéma + appliquer
            #    les champs interdits (double sécurité côté serveur)
            donnees_brutes = llm_result.get("donnees_extraites", {}) or {}
            interdits = CHAMPS_INTERDITS.get(expected_type, [])
            donnees = {}
            for cle in SCHEMA_EXTRACTION:
                val = donnees_brutes.get(cle, "")
                val = "" if val is None else str(val).strip()
                if cle in interdits:
                    val = ""
                donnees[cle] = val

            is_valid = bool(llm_result.get("correspond_au_type", False))

            # 4. Score RÉEL calculé en Python
            score, nb_trouves, nb_attendus = self._calculer_score(expected_type, donnees)

            # Seuil : il faut que le document corresponde au type ET qu'au
            # moins 70% des champs attendus soient trouvés pour être VERT
            if not is_valid:
                statut_global = "ROUGE_REJET"
            elif score >= 70:
                statut_global = "VERT_AUTOMATIQUE"
            else:
                statut_global = "ORANGE_TOLERANCE"

            os.remove(tmp_file_path)
            return {
                "statut_global": statut_global,
                "score_semantique_llm": score,
                "champs_trouves": nb_trouves,
                "champs_attendus": nb_attendus,
                "resultat_analyse": {
                    "correspond_au_type": is_valid,
                    "donnees_extraites": donnees
                }
            }

        except Exception as e:
            if os.path.exists(tmp_file_path): os.remove(tmp_file_path)
            return {"statut_global": "ERREUR_SYSTEME", "erreur": str(e)}