"""
Pré-filtre visuel CLIP (ViT-B/32) — triage zero-shot des documents KYC.

Place dans le pipeline : APRES le televersement, AVANT Docling + Gemini.
Objectif : economiser le quota Gemini en ecartant rapidement (CPU, < 1s)
les documents qui ne ressemblent clairement pas au type attendu.

Design retenu (cf. discussion architecture) :
- Zero-shot : aucun dataset a constituer, aucun fine-tuning. Les classes
  sont decrites par des phrases (prompts) comparees aux embeddings d'image.
- Triage, PAS rejet binaire : seuls les desaccords a haute confiance
  (>= SEUIL_BLOCAGE) bloquent. Une basse confiance repart vers le
  pipeline classique (Docling + Gemini) — un document atypique mais
  legitime ne doit jamais etre perdu.
- INT8 dynamique sur CPU : quantification des couches Linear (~4x moins
  de memoire, accelération sur CPU). Fallback FP32 si indisponible.
- PDF : rasterisation de la 1re page via PyMuPDF (aucune dependance
  systeme Poppler — compatible HF Spaces).

Dependances a ajouter dans requirements.txt :
    open_clip_torch
    pymupdf
    torch (deja present)
    Pillow (deja present)
"""

import os
import io

import torch
import open_clip
from PIL import Image

# Seuil de confiance a partir duquel un desaccord type_attendu/detecte
# bloque le document AVANT de consommer du quota Gemini.
# 0.80 = tolerant (recommande pour demarrer, a ajuster apres mesure).
SEUIL_BLOCAGE = 0.80

# Libelles d'affichage (cles = identifiants internes = doc_type du formulaire)
CATEGORY_LABELS = {
    "CNI_PASSEPORT": "CNI / Passeport",
    "JUSTIFICATIF_ADRESSE": "Justificatif d'adresse",
    "BULLETIN_SALAIRE": "Bulletin de paie",
    "RIB": "RIB / Releve bancaire",
    "RELEVE_COMPTE": "Relevé bancaire / Mobile Money (Wave, OM)",
    "AUTRE": "Autre document",
}

# Prompts textuels zero-shot : plusieurs par categorie pour robustesse.
PROMPTS = {
    "CNI_PASSEPORT": [
        "a photo of an identity card or a passport page with a portrait photo and official stamps",
        "a national identity card document",
    ],
    "JUSTIFICATIF_ADRESSE": [
        "a utility bill, rent receipt or official proof of address showing a postal address",
        "an invoice or official letter showing a home address",
    ],
    "BULLETIN_SALAIRE": [
        "a payslip or salary statement document with tables of amounts and deductions",
        "a pay stub with a net salary",
    ],
    "RIB": [
        "a bank account details document with an IBAN number and a bank logo",
        "a bank statement header with bank coordinates",
    ],
    "RELEVE_COMPTE": [
        "a bank statement or mobile money statement showing a list of transactions and account balance",
        "a mobile money account statement from Wave or Orange Money with transaction history",
    ],
    "AUTRE": [
        "a random photo, a screenshot or any unrelated document",
    ],
}


class DocumentClassifier:
    """Singleton — le chargement du modele ne se fait qu'une fois."""

    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(DocumentClassifier, cls).__new__(cls)
            cls._instance._load()
        return cls._instance

    def _load(self):
        print("[PREFILTRE] Chargement de CLIP ViT-B/32...")
        model, _, self.preprocess = open_clip.create_model_and_transforms(
            "ViT-B-32", pretrained="laion2b_s34b_b79k"
        )

        # Quantification INT8 dynamique (CPU) — les deux tours separement,
        # inplace. Si l'un echoue, on continue en FP32 (budget < 1s tenu).
        try:
            torch.quantization.quantize_dynamic(model.visual, {torch.nn.Linear}, dtype=torch.qint8, inplace=True)
        except Exception as e:
            print(f"[PREFILTRE] Quantification 'visual' impossible ({e}) — FP32.")
        try:
            torch.quantization.quantize_dynamic(model.text, {torch.nn.Linear}, dtype=torch.qint8, inplace=True)
        except Exception as e:
            print(f"[PREFILTRE] Quantification 'text' impossible ({e}) — FP32.")

        model.eval()
        self.model = model
        self.tokenizer = open_clip.get_tokenizer("ViT-B-32")

        # Encodage des prompts : fait UNE SEULE FOIS au chargement.
        self.prompt_index = []  # position -> categorie
        all_prompts = []
        for cat, prompts in PROMPTS.items():
            for p in prompts:
                all_prompts.append(p)
                self.prompt_index.append(cat)

        with torch.no_grad():
            tokens = self.tokenizer(all_prompts)
            txt = self.model.encode_text(tokens)
            self.text_features = txt / txt.norm(dim=-1, keepdim=True)
        print("[PREFILTRE] Pret.")

    @staticmethod
    def _to_image(file_bytes, file_name):
        """PDF -> rasterise la 1re page (PyMuPDF, sans Poppler). Sinon PIL."""
        ext = os.path.splitext(file_name)[1].lower()
        if ext == ".pdf":
            import fitz  # pymupdf
            doc = fitz.open(stream=file_bytes, filetype="pdf")
            pix = doc[0].get_pixmap(dpi=144)
            img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")
            doc.close()
            return img
        return Image.open(io.BytesIO(file_bytes)).convert("RGB")

    @torch.no_grad()
    def classify(self, file_bytes, file_name):
        """
        Retourne (categorie_detectee, confiance) avec categorie dans
        CATEGORY_LABELS. Confiance = probabilite softmax agregee par
        categorie (moyenne sur ses prompts).
        """
        img = self.preprocess(self._to_image(file_bytes, file_name)).unsqueeze(0)
        feats = self.model.encode_image(img)
        feats = feats / feats.norm(dim=-1, keepdim=True)

        logits = self.model.logit_scale.exp() * feats @ self.text_features.T
        probs = logits.softmax(dim=-1).cpu().numpy()[0]

        agg = {}
        for p, cat in zip(probs, self.prompt_index):
            agg[cat] = agg.get(cat, 0.0) + float(p)
        total = sum(agg.values()) or 1.0
        best = max(agg, key=agg.get)
        return best, agg[best] / total


# ── Test rapide en ligne de commande ─────────────────────────────────────────
if __name__ == "__main__":
    import sys
    clf = DocumentClassifier()
    for path in sys.argv[1:]:
        with open(path, "rb") as fh:
            cat, conf = clf.classify(fh.read(), os.path.basename(path))
        print(f"{path} -> {CATEGORY_LABELS.get(cat, cat)} [{conf:.1%}]")