src/guichetoi/inference.py

"""
guichetoi.inference — classify a document and extract fields with LayoutLMv3.

Two entry points:

    CLI mode (single document, prints JSON to stdout, saves a copy):
        python -m guichetoi.inference --image path/to/doc.pdf
        python -m guichetoi.inference --image path/to/doc.png --ocr "optional pre-extracted text"

    Library mode (for FastAPI / web app — load models once, reuse for every request):
        from guichetoi.inference import GuichetOIPipeline
        pipeline = GuichetOIPipeline()                   # load once at startup
        result   = pipeline.run("path/to/doc.pdf")       # call per request

Output: structured dict with doc_class, per-field values, and per-field confidence.

Author: Aziz Mohamed Miladi · GuichetOI ML
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import re
import sys
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import torch
from PIL import Image

try:
    import fitz  # PyMuPDF
except ImportError:
    fitz = None

try:
    import pytesseract
except ImportError:
    pytesseract = None

from transformers import (
    LayoutLMv3ForSequenceClassification,
    LayoutLMv3ForTokenClassification,
    LayoutLMv3Processor,
)

# ────────────────────────────────────────────────────────────────────────────
# Logging
# ────────────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-7s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("guichetoi.inference")


# ────────────────────────────────────────────────────────────────────────────
# Configuration
# ────────────────────────────────────────────────────────────────────────────
# Paths default to repo-root-relative locations so the pipeline works
# regardless of the caller's CWD (Streamlit, FastAPI, CLI from any folder).
# Override any of them via env vars (useful in Docker / k8s where models
# might be mounted from a volume):
#
#   GUICHETOI_CLASSIFIER_DIR
#   GUICHETOI_EXTRACTOR_DIR
#   GUICHETOI_MAPPINGS_PATH
#   GUICHETOI_OUTPUT_DIR
#
# Anchor: src/guichetoi/inference.py  →  parents[2] = repo root
REPO_ROOT = Path(__file__).resolve().parents[2]


def _env_path(env_var: str, default: Path) -> str:
    return os.environ.get(env_var) or str(default)


@dataclass(frozen=True)
class Config:
    """All inference-time configuration in one place."""
    classifier_dir:   str = field(default_factory=lambda: _env_path(
        "GUICHETOI_CLASSIFIER_DIR", REPO_ROOT / "models" / "classifier"))
    extractor_dir:    str = field(default_factory=lambda: _env_path(
        "GUICHETOI_EXTRACTOR_DIR", REPO_ROOT / "models" / "extractor_v3_backup_v2"))
    mappings_path:    str = field(default_factory=lambda: _env_path(
        "GUICHETOI_MAPPINGS_PATH", REPO_ROOT / "assets" / "label_mappings.json"))
    base_processor:   str = "microsoft/layoutlmv3-base"

    max_seq_length:   int = 512        # WordPiece tokens (LayoutLMv3 limit)
    max_words:        int = 1024       # OCR words; processor will truncate to 512 tokens
    ocr_min_conf:     int = 20         # Match training-time filter (Audit Defect 2)

    needs_extraction: frozenset = frozenset({"fiche", "Autorisation", "Mandat", "Certificat"})
    pdf_render_zoom:  float = 2.0      # 2× DPI uplift for OCR quality

    output_dir:       str = field(default_factory=lambda: _env_path(
        "GUICHETOI_OUTPUT_DIR", REPO_ROOT / "outputs"))


# ────────────────────────────────────────────────────────────────────────────
# Data classes for clean return values
# ────────────────────────────────────────────────────────────────────────────
@dataclass
class FieldExtraction:
    """A single extracted field with its confidence."""
    value:      str
    confidence: float

@dataclass
class InferenceResult:
    """Full result of one document inference."""
    image:           str
    doc_class:       str
    doc_confidence:  float
    pages_processed: int
    ocr_source:      str
    fields:          dict = field(default_factory=dict)   # name → FieldExtraction

    def to_dict(self) -> dict:
        d = asdict(self)
        d["fields"] = {k: asdict(v) for k, v in self.fields.items()}
        return d


# ────────────────────────────────────────────────────────────────────────────
# Path resolution — handles local dirs, HF Trainer checkpoints, HF Hub IDs
# ────────────────────────────────────────────────────────────────────────────
def resolve_model_path(model_dir: str) -> str:
    """
    Returns a string suitable for ``from_pretrained()``. Decision logic:

    1. If the path exists locally → resolve as a local dir (with optional
       checkpoint-N subdirectory selection for HF Trainer outputs).
    2. If the path does NOT exist locally AND has exactly one "/" (owner/repo)
       → treat as an HF Hub repo ID and let ``from_pretrained`` download it.
       Set HF_TOKEN in the environment for private repos.
    3. Otherwise → FileNotFoundError with a helpful tip.
    """
    p = Path(model_dir)

    # ── Local path ──────────────────────────────────────────────────────────
    if p.exists():
        for marker in ("config.json", "model.safetensors", "pytorch_model.bin"):
            if (p / marker).exists():
                return str(p)

        checkpoints = [c for c in p.glob("checkpoint-*") if c.is_dir()]
        if checkpoints:
            latest = max(checkpoints, key=lambda c: int(c.name.split("-")[-1]))
            log.info(f"Using checkpoint: {latest.name}")
            return str(latest)

        raise FileNotFoundError(
            f"No model artifacts in {p}. Expected one of: "
            "config.json, model.safetensors, pytorch_model.bin, or checkpoint-*/"
        )

    # ── HF Hub repo ID (owner/repo, path doesn't exist locally) ─────────────
    parts = model_dir.split("/")
    if len(parts) == 2 and all(parts):  # exactly "owner/repo", both non-empty
        log.info(f"Resolved as HF Hub repo: {model_dir}")
        return model_dir

    raise FileNotFoundError(
        f"Model not found: '{model_dir}'\n"
        "  - For a local path: make sure the directory exists and contains "
        "model.safetensors or checkpoint-N/ subdirs.\n"
        "  - For HuggingFace Spaces: set GUICHETOI_CLASSIFIER_DIR / "
        "GUICHETOI_EXTRACTOR_DIR to an HF Hub repo ID "
        "(e.g. 'medaziz012/guichetoi-classifier')."
    )


# ────────────────────────────────────────────────────────────────────────────
# Image / PDF loading
# ────────────────────────────────────────────────────────────────────────────
def load_pages(file_path: Path, cfg: Config) -> list[Image.Image]:
    """
    Load all pages of a document as PIL Images.
    Returns a list of one image for non-PDF inputs, or N images for PDFs.
    """
    suffix = file_path.suffix.lower()

    if suffix == ".pdf":
        if fitz is None:
            raise RuntimeError("PyMuPDF not installed — cannot read PDFs. pip install pymupdf")
        pages = []
        with fitz.open(file_path) as doc:
            matrix = fitz.Matrix(cfg.pdf_render_zoom, cfg.pdf_render_zoom)
            for page in doc:
                pix = page.get_pixmap(matrix=matrix)
                pages.append(Image.frombytes("RGB", (pix.width, pix.height), pix.samples))
        return pages

    return [Image.open(file_path).convert("RGB")]


# ────────────────────────────────────────────────────────────────────────────
# OCR — single pass, uses confidence filter that matches training
# ────────────────────────────────────────────────────────────────────────────
@dataclass
class OCRResult:
    words:  list[str]
    boxes:  list[list[int]]   # normalised to [0, 1000]
    text:   str
    source: str               # "pdf_text", "pytesseract", or "fallback"


def _normalize_text(text: str) -> str:
    return re.sub(r"\s+", " ", (text or "").strip())


def _vertical_fallback_boxes(n_words: int) -> list[list[int]]:
    """Last-resort uniform vertical strip boxes when no real OCR is available."""
    if n_words <= 0:
        return []
    h = max(1000 // n_words, 1)
    return [[0, i * h, 1000, min((i + 1) * h, 1000)] for i in range(n_words)]


# ────────────────────────────────────────────────────────────────────────────
# Per-class field allowlists
# Each document class has only a handful of relevant fields. The model and
# regex fallbacks can produce extractions for fields that don't belong to
# the predicted class (e.g. `Representant_Email` on a fiche-de-renseignement).
# We filter those out after extraction so demo output only shows fields that
# actually make sense for the document type.
# ────────────────────────────────────────────────────────────────────────────
CLASS_FIELDS: dict[str, frozenset[str]] = {
    "fiche": frozenset({
        "Reference_Urbanisme", "DLPI", "cabinet_conseil",
        "Disposition_Mandat", "Batiment_Adresse",
        "nb_log_totale", "Nb_log_pro", "Nb_log_res",
        "Nombre_Logement_Lot_MacroLot",
    }),
    "Mandat": frozenset({
        "Representant_Email", "Representant_Nom_Complet",
        "Representant_Telephone", "Disposition_Mandat",
        "cabinet_conseil",
    }),
    "Autorisation": frozenset({
        "Reference_Urbanisme", "Batiment_Adresse", "DLPI",
        "nb_log_totale",
    }),
    "Certificat": frozenset({
        "Reference_Urbanisme", "Batiment_Adresse",
    }),
}


# ────────────────────────────────────────────────────────────────────────────
# Post-processing — clean noisy model outputs with field-specific validators
# ────────────────────────────────────────────────────────────────────────────
_RE_EMAIL    = re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+")
_RE_PHONE_FR = re.compile(r"(?<!\d)(0[1-9](?:[ .-]?\d){8})(?!\d)")
_RE_REFURB   = re.compile(
    # Urbanism reference codes: PC / PA / DP / CU + immediate digit + body of
    # digits, whitespace, dashes or UPPERCASE letters. Prefix is case-insensitive
    # via `(?i:…)` so "Pc0440…" matches, but the BODY must be uppercase/digits —
    # otherwise the regex catches French words like "rue", "Parcelle" (where the
    # `RU`/`PA` substring trips a too-permissive case-insensitive match).
    r"\b(?i:PC|PA|DP|CU)[\s\-]*\d[\d\sA-Z\-]{4,28}"
)
_RE_INTEGER  = re.compile(r"\b(\d{1,4})\b")

# French postal address — anchored on a street-type keyword so we don't
# match arbitrary "<digit> <text>" sequences. Optional 5-digit postcode +
# city at the end.
_RE_ADDR_FR  = re.compile(
    r"\b\d{1,4}\s*(?:BIS|TER|QUATER|QUINQUIES)?\s+"
    r"(?:rue|avenue|av\.?|boulevard|bd\.?|route|chemin|place|"
    r"all[ée]e|impasse|cours|quai|esplanade|cit[ée]|square|voie|sentier)"
    # Street body excludes digits → the postal code can't be swallowed into
    # the street name. Also excludes the form-label characters °, |, and
    # newline/comma/semicolon so we don't gobble trailing form text like
    # "N° Rue Code Postal Ville".
    r"\s+[^\n,;\d°|]{3,50}"
    # Body is greedy and includes the trailing space → the postal-code
    # separator must accept ZERO chars (`*` not `+`) so the optional group
    # can still latch onto the digit directly.
    r"(?:[,\s]*(\d{5})\s+[\w\-' ]{3,40})?",
    re.IGNORECASE,
)

_NAME_STOPWORDS = re.compile(
    r"\b(Conseiller|Neuf|Mobile|Mail|Email|T[ée]l(?:[ée]phone)?|Adresse|"
    r"Soci[ée]t[ée]|Bureau|Cabinet|Conseil)\b",
    re.IGNORECASE,
)
_ADDRESS_STOPWORDS = re.compile(
    # OCR commonly mis-renders the ligature "Œ" as "OE" (two ASCII letters),
    # so we accept both spellings for "D'ŒUVRE" / "D'OEUVRE".
    r"\b(FICHE|DESCRIPTION|MAITRE|D[’']?OUVRAGE|D[’']?(?:[OŒ]|OE)UVRE|"
    r"CABINET|CONSEIL|BUREAU|OPERATION|RENSEIGNEMENT|PROPRIETAIRE)\b",
    re.IGNORECASE,
)

# Trailing form-field labels / boilerplate that often comes RIGHT AFTER a
# valid address in OCR'd documents — we trim them so the address stays
# clean. Includes OCR mis-readings of `N°` (rendered as `ne`, `nw`, `No`).
_ADDR_TRAIL_TRIM = re.compile(
    r"\s+"
    r"(?:N[°oewé]{0,2}|No|Ne|Nw|Code(?:\s+Postal)?|Postal|Ville|Pays|"
    r"Adresse|Tel|T[ée]l|Email|Je\s+soussign[ée]?|Travaux|Construction|"
    r"Parcelle|Nb\s+de|Lot|CERTIFICAT|PERMIS|Surface)"
    r"\b.*$",
    re.IGNORECASE,
)


def _clean_address_value(addr: str) -> str:
    """Single source of truth for Batiment_Adresse cleanup. Applied to both
    the model's raw output AND the OCR backstop, so the same trimming runs
    regardless of which source produced the address."""
    if not addr:
        return ""
    a = re.sub(r"\s+", " ", addr).strip(" ,.-/")
    a = _ADDRESS_STOPWORDS.sub(" ", a)
    a = _ADDR_TRAIL_TRIM.sub("", a)
    # Trim parenthesized boilerplate (e.g. "(emprise au sol) ...")
    a = re.sub(r"\s*\([^)]*\).*$", "", a)
    # Trim trailing 1-2-char tokens — almost always the first letter of the
    # next form field caught by the regex.
    a = re.sub(r"\s+\S{1,2}\s*$", "", a)
    a = re.sub(r"\s+", " ", a).strip(" ,.-/:;")
    return a
_CABINET_STOPWORDS = re.compile(
    r"\b(OUI|NON|D[eé]nomination|sociale|si\s*oui|si\s*non|mobile|Adresse)\b",
    re.IGNORECASE,
)


_MANDAT_CTX_KEYWORDS = ("ouvrage", "mandat", "dispose", "représ", "repr�s", "represent")


def _mandat_checkbox_score(marker: str) -> int:
    """
    Strict 'is this an X-marked checkbox?' score for an OCR-rendered marker.

    The heuristic only counts STRONG signals — patterns that almost never
    appear in an empty `[]` box. A single ambiguous character like `!`,
    `:`, `D`, `si` is NOT a strong signal: empty boxes degenerate into all
    sorts of one-character garble (Tesseract reads `[]` as `D`, `O`, `Q`,
    `I`, `!`, `|`, …), so we'd be guessing.

    Strong signals (in order of confidence):
      - Explicit X / check-mark glyph (X, ✓, ✗, …) → 5
      - A digit inside the marker (Tesseract often reads an X as 1 or 9)
        wrapped in a small token → 3
      - Multi-character mark pattern like `**`, `*[]`, `[X]`, `[*]` → 3
      - An 'orphan' bracket — one of `[` or `]` but not both — which is
        the classic OCR fragment of `[X]` after the X disappeared → 2

    Anything else returns 0. Better to return None from the detector than
    to commit on noise.
    """
    s = (marker or "").strip()
    if not s:
        return 0

    # X / check glyphs — the strongest signal
    if re.search(r"[Xx✓✔✗✘]", s):
        return 5

    # Digit inside a short marker token — Tesseract often reads `[X]` as `[1]`
    if re.search(r"[1-9]", s):
        return 3

    # Multi-character mark patterns (e.g. `**`, `**[]`)
    if re.search(r"[*#]{2,}", s):
        return 3

    # Orphan bracket — `]` without a matching `[`, or vice versa
    if ("[" in s) != ("]" in s):
        return 2

    # Everything else (single punctuation, single letter, short word) is
    # too weak to claim a checkbox is marked.
    return 0


def _detect_mandat_checkbox(ocr_text: str) -> str | None:
    """
    Decide which checkbox is X-marked next to 'Je dispose d'un mandat de
    représentation du Maître d'ouvrage' on the fiche form.

    Strategy: scan every OUI<m1>/NON<m2> pair in the OCR. For each, look at
    the 200 characters immediately before to see whether it sits in the
    mandat context (keywords: ouvrage, mandat, dispose, …). Pick the first
    matching pair and decide which marker is heavier (= more likely X).
    """
    norm = re.sub(r"\s+", " ", ocr_text)

    pair_re = re.compile(
        r"OUI\s*([^/]{0,15}?)\s*/\s*(?:NON|Non|non)\s*(\S{0,15})",
        re.IGNORECASE,
    )

    for m in pair_re.finditer(norm):
        before = norm[max(0, m.start() - 200): m.start()].lower()
        if not any(k in before for k in _MANDAT_CTX_KEYWORDS):
            continue
        o = _mandat_checkbox_score(m.group(1))
        n = _mandat_checkbox_score(m.group(2))
        if o > n:
            return "OUI"
        if n > o:
            return "NON"
        return None  # ambiguous

    return None


def _clean_field_extractions(
    raw_fields: dict[str, FieldExtraction],
    ocr_text: str,
) -> dict[str, FieldExtraction]:
    """
    Apply per-field validators + regex fallbacks to the model's raw outputs.

    The token-classifier sometimes catches form-label words ("NOM", "Adresse:",
    "OUI/NON", "DESCRIPTION") instead of the actual value cell, because the
    training annotations themselves landed on those words when Tesseract
    missed the small digits/text in the value cells. Without this cleanup the
    raw extractions are noisy enough to look amateurish in a demo.

    Strategy per field:
      - Try to extract a valid-format value from the model's noisy span.
      - If that fails AND the field has a reliable OCR-text pattern, fall
        back to regex against the full OCR text.
      - If still nothing, DROP the field rather than emit garbage.
    """
    cleaned: dict[str, FieldExtraction] = {}

    # Minimum confidence below which we won't trust the model output unless
    # a downstream regex validator can pull a well-formed value out of it.
    # Set conservatively — better to drop than to publish low-confidence noise.
    MIN_TRUSTED_CONF = 0.40

    for name, extr in raw_fields.items():
        raw = (extr.value or "").strip()
        conf = extr.confidence

        # For free-text fields (not regex-extractable), require minimum confidence
        if name in ("cabinet_conseil", "Batiment_Adresse", "Representant_Nom_Complet") and conf < MIN_TRUSTED_CONF:
            continue

        if name == "Representant_Email":
            m = _RE_EMAIL.search(raw)
            if m:
                cleaned[name] = FieldExtraction(m.group(0), conf)

        elif name == "Representant_Telephone":
            m = _RE_PHONE_FR.search(raw)
            if m:
                phone = re.sub(r"\s+", " ", m.group(1)).strip()
                cleaned[name] = FieldExtraction(phone, conf)

        elif name == "Reference_Urbanisme":
            m = _RE_REFURB.search(raw)
            if m:
                ref = re.sub(r"\s+", " ", m.group(0)).strip()
                cleaned[name] = FieldExtraction(ref, conf)

        elif name == "Representant_Nom_Complet":
            value = _NAME_STOPWORDS.split(raw)[0].strip()
            value = re.sub(r"[,;:]+$", "", value).strip()
            if 3 <= len(value) <= 60 and not re.search(r"[<>{}]", value):
                cleaned[name] = FieldExtraction(value, conf)

        elif name in ("nb_log_totale", "Nb_log_pro", "Nb_log_res", "Nombre_Logement_Lot_MacroLot"):
            m = _RE_INTEGER.search(raw)
            if m:
                n = int(m.group(1))
                if 0 <= n <= 9999:
                    cleaned[name] = FieldExtraction(str(n), conf)

        elif name == "DLPI":
            if _ADDRESS_STOPWORDS.search(raw):
                continue  # form text, not a DLPI
            if re.match(r"^\d{1,2}\s*/\s*\d{1,2}\s*/\s*\d{2,4}$", raw):
                cleaned[name] = FieldExtraction(raw, conf)
            elif re.match(r"^[A-Z0-9][\w/.\- ]{1,30}$", raw):
                cleaned[name] = FieldExtraction(raw[:30].strip(), conf)

        elif name == "Disposition_Mandat":
            # Use the checkbox detector on the full OCR text. The previous
            # fallback that picked the first OUI/NON word from the model's
            # noisy span was unreliable — it routinely answered "OUI" just
            # because OUI happens to appear before NON in the form text.
            # If the detector can't reach a confident decision, DROP the
            # field and let the recommendation engine flag the case for
            # manual review rather than committing on a coin flip.
            detected = _detect_mandat_checkbox(ocr_text)
            if detected:
                cleaned[name] = FieldExtraction(detected, max(conf, 0.85))

        elif name == "cabinet_conseil":
            if _CABINET_STOPWORDS.search(raw):
                continue
            if 2 <= len(raw) <= 60:
                cleaned[name] = FieldExtraction(raw, conf)

        elif name == "Batiment_Adresse":
            # Address values from any doc class (model output) get the full
            # cleanup pass — strip form headers AND trailing form labels.
            # Threshold 8 chars: shortest meaningful address is ~"1 rue X" =
            # 7 chars, anything below is a fragment ("1 rue", "rue X").
            stripped = _clean_address_value(raw)
            if 8 <= len(stripped) <= 200:
                cleaned[name] = FieldExtraction(stripped, conf)

        else:
            cleaned[name] = extr

    # ── Backstop: fields the model missed entirely, but OCR has the answer ──
    if "Representant_Email" not in cleaned:
        m = _RE_EMAIL.search(ocr_text)
        if m:
            cleaned["Representant_Email"] = FieldExtraction(m.group(0), 0.6)
    if "Representant_Telephone" not in cleaned:
        m = _RE_PHONE_FR.search(ocr_text)
        if m:
            phone = re.sub(r"\s+", " ", m.group(1)).strip()
            cleaned["Representant_Telephone"] = FieldExtraction(phone, 0.6)
    if "Reference_Urbanisme" not in cleaned:
        m = _RE_REFURB.search(ocr_text)
        if m:
            cleaned["Reference_Urbanisme"] = FieldExtraction(
                re.sub(r"\s+", " ", m.group(0)).strip(), 0.6
            )
    if "Batiment_Adresse" not in cleaned:
        # Most fiches don't reliably extract the address via the model.
        # The OCR text often contains the address verbatim — grab it with
        # a street-type-anchored regex and run the same cleanup as the
        # model-output path so behaviour is consistent.
        m = _RE_ADDR_FR.search(ocr_text)
        if m:
            addr = _clean_address_value(m.group(0))
            if 8 <= len(addr) <= 200:
                cleaned["Batiment_Adresse"] = FieldExtraction(addr, 0.6)

    # ── Disposition_Mandat: checkbox detection backstop ──────────────────
    if "Disposition_Mandat" not in cleaned:
        detected = _detect_mandat_checkbox(ocr_text)
        if detected:
            cleaned["Disposition_Mandat"] = FieldExtraction(detected, 0.85)

    # ── Logement total: regex backstop against the full OCR text ─────────
    # `nb_log_totale` (= total = residential + professional buildings) is
    # the only logement field where the form label maps cleanly to an
    # OCR-extractable pattern. The macrolot threshold lines (<= 3 / > 3
    # logements) on the form refer to MACROLOT counts, not residential vs
    # professional building counts — extracting them as Nb_log_res /
    # Nb_log_pro would mis-label the field. So those two are left to the
    # model (with its known limitations) and the regex backstop only fills
    # in nb_log_totale.
    if "nb_log_totale" not in cleaned:
        norm_ocr = re.sub(r"\s+", " ", ocr_text)
        for pat in (
            r"Nb\s+total\s+de\s+logements\b[^:]*?:\s*(\d+)",
            r"logements\s*/\s*locaux\s*/\s*lots\b[^:]*?:\s*(\d+)",
        ):
            m = re.search(pat, norm_ocr, re.IGNORECASE)
            if m:
                cleaned["nb_log_totale"] = FieldExtraction(m.group(1), 0.7)
                break

    return cleaned


def run_ocr(image: Image.Image, cfg: Config) -> OCRResult:
    """
    Single-pass OCR using pytesseract, returning words + normalised boxes
    using the SAME confidence filter as the training pipeline.
    """
    if pytesseract is None:
        log.warning("pytesseract not installed — falling back to vertical strips")
        return OCRResult([], [], "", "fallback")

    img_w, img_h = image.size
    data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)

    words, boxes = [], []
    for i, raw_token in enumerate(data.get("text", [])):
        token = (raw_token or "").strip()
        if not token:
            continue

        # Confidence filter — MUST match training. Drops -1 sentinels AND low-confidence tokens.
        try:
            conf = float(data.get("conf", ["-1"])[i])
        except (ValueError, TypeError):
            conf = -1
        if conf < cfg.ocr_min_conf:
            continue

        left   = int(data["left"][i])
        top    = int(data["top"][i])
        width  = int(data["width"][i])
        height = int(data["height"][i])
        if width <= 0 or height <= 0:
            continue

        # Normalise to [0, 1000] — LayoutLMv3 contract
        boxes.append([
            max(0, min(1000, int(left            / img_w * 1000))),
            max(0, min(1000, int(top             / img_h * 1000))),
            max(0, min(1000, int((left + width)  / img_w * 1000))),
            max(0, min(1000, int((top + height)  / img_h * 1000))),
        ])
        words.append(token)

        if len(words) >= cfg.max_words:
            log.info(f"Reached max_words={cfg.max_words}; truncating OCR")
            break

    if not words:
        log.warning("OCR returned no usable words — using vertical fallback")
        return OCRResult(["[PAD]"], _vertical_fallback_boxes(1), "", "fallback")

    return OCRResult(words, boxes, " ".join(words), "pytesseract")


def extract_pdf_text(file_path: Path) -> str | None:
    """Quick path: pull embedded text from a PDF without OCR. Returns None if no text or fails."""
    if file_path.suffix.lower() != ".pdf" or fitz is None:
        return None
    try:
        with fitz.open(file_path) as doc:
            text = "\n".join(page.get_text("text") for page in doc)
        text = _normalize_text(text)
        return text or None
    except Exception as e:
        log.debug(f"PDF text extraction failed: {e}")
        return None


# ────────────────────────────────────────────────────────────────────────────
# Pipeline — load once, reuse for every request
# ────────────────────────────────────────────────────────────────────────────
class GuichetOIPipeline:
    """
    Loads classifier + extractor + processor once.
    Call .run(image_path) for each document — no model reloading.

    Use this from the FastAPI service:
        pipeline = GuichetOIPipeline()    # at app startup
        result   = pipeline.run(path)     # in your /predict endpoint
    """

    def __init__(self, cfg: Config = Config(), device: str | None = None):
        self.cfg = cfg
        self.device = torch.device(
            device or ("cuda" if torch.cuda.is_available() else "cpu")
        )
        log.info(f"Loading models on device: {self.device}")

        # Label mappings
        with open(cfg.mappings_path, encoding="utf-8") as f:
            self.mappings = json.load(f)
        self.doc_classes  = self.mappings["doc_classes"]
        self.field_labels = self.mappings["field_labels"]

        # Processor (no internal OCR — we feed our own words+boxes)
        self.processor = LayoutLMv3Processor.from_pretrained(
            cfg.base_processor, apply_ocr=False,
        )

        # HF_TOKEN — required for private HF Hub repos (HuggingFace Spaces deploy).
        # Not needed when loading from local paths.
        hf_token = os.environ.get("HF_TOKEN") or None

        # Models — moved to device, set to eval mode.
        # (mypy ignores below: transformers stubs mis-type from_pretrained's
        #  return so .to(device) appears to take the wrong arg.)
        self.classifier = LayoutLMv3ForSequenceClassification.from_pretrained(
            resolve_model_path(cfg.classifier_dir), token=hf_token,
        ).to(self.device).eval()  # type: ignore[arg-type, unused-ignore]

        self.extractor = LayoutLMv3ForTokenClassification.from_pretrained(
            resolve_model_path(cfg.extractor_dir), token=hf_token,
        ).to(self.device).eval()  # type: ignore[arg-type, unused-ignore]

        log.info(
            f"Pipeline ready · {len(self.doc_classes)} document classes · "
            f"{len(self.field_labels)} field labels"
        )

    # ────────────────────────────────────────────────────────────────────
    # Inference primitives
    # ────────────────────────────────────────────────────────────────────
    def _encode(self, image: Image.Image, words: list[str], boxes: list[list[int]]) -> Any:
        # Returns a transformers BatchEncoding (dict-like, with .word_ids()).
        # Annotated Any because the transformers stubs are ignored in mypy.ini.
        return self.processor(
            image, words, boxes=boxes,
            max_length=self.cfg.max_seq_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        ).to(self.device)

    @torch.no_grad()
    def classify(self, image: Image.Image, words: list[str], boxes: list[list[int]]) -> tuple[str, float]:
        encoding = self._encode(image, words, boxes)
        logits   = self.classifier(**encoding).logits
        probs    = torch.softmax(logits, dim=-1)[0]
        pred_id  = int(probs.argmax())
        return self.doc_classes[pred_id], float(probs[pred_id])

    @torch.no_grad()
    def extract(self, image: Image.Image, words: list[str], boxes: list[list[int]]) -> dict[str, FieldExtraction]:
        """
        Run the BIO extractor and reconstruct spans.

        A span:
          - opens on a B-X tag
          - extends through consecutive I-X tags with the SAME field name
          - closes on O, on a different B-, or on an I- with a different field name
          - rejects orphan I- tags (I- without a matching B- → ignored, prevents phantom spans)
        """
        encoding = self._encode(image, words, boxes)
        outputs  = self.extractor(**encoding)
        logits   = outputs.logits[0]                             # (T, n_labels)
        probs    = torch.softmax(logits, dim=-1)                 # per-token probabilities
        pred_ids = logits.argmax(dim=-1).tolist()
        word_ids = encoding.word_ids(batch_index=0)
        id2label = self.extractor.config.id2label

        spans: list[dict] = []
        cur: dict | None = None
        prev_word = None

        for pos, w_idx in enumerate(word_ids):
            # Skip special tokens and continuation sub-words (only score head sub-word per word)
            if w_idx is None or w_idx == prev_word:
                continue
            prev_word = w_idx

            # Out of bounds (truncation safety)
            if w_idx >= len(words):
                continue

            label = id2label.get(pred_ids[pos], "O")
            conf  = float(probs[pos, pred_ids[pos]])

            if label == "O":
                if cur is not None:
                    spans.append(cur)
                    cur = None
                continue

            tag, _, name = label.partition("-")

            if tag == "B":
                # Close any open span and start a new one
                if cur is not None:
                    spans.append(cur)
                cur = {"name": name, "words": [words[w_idx]], "confs": [conf]}

            elif tag == "I":
                # Continue current span if names match; otherwise drop the orphan I-
                if cur is not None and cur["name"] == name:
                    cur["words"].append(words[w_idx])
                    cur["confs"].append(conf)
                # else: orphan I- without a matching B- → IGNORE (do not start a new span)

        # Don't forget the trailing span
        if cur is not None:
            spans.append(cur)

        # Aggregate spans of the same field name (e.g. multi-line addresses)
        result: dict[str, FieldExtraction] = {}
        for span in spans:
            text = " ".join(span["words"])
            mean_conf = sum(span["confs"]) / len(span["confs"])
            if span["name"] in result:
                # Concatenate multi-span fields, average confidence weighted by length
                prev = result[span["name"]]
                combined_text = f"{prev.value} {text}".strip()
                combined_conf = (prev.confidence + mean_conf) / 2
                result[span["name"]] = FieldExtraction(combined_text, round(combined_conf, 4))
            else:
                result[span["name"]] = FieldExtraction(text, round(mean_conf, 4))

        return result

    # ────────────────────────────────────────────────────────────────────
    # Public entry point
    # ────────────────────────────────────────────────────────────────────
    def run(self, image_path: str | Path, ocr_text: str = "") -> InferenceResult:
        image_path = Path(image_path)
        if not image_path.exists():
            raise FileNotFoundError(image_path)

        log.info(f"Processing: {image_path.name}")

        # Multi-page support: process every page, aggregate at the end
        pages = load_pages(image_path, self.cfg)
        log.info(f"Loaded {len(pages)} page(s)")

        # Decide OCR source ONCE per document — no double OCR
        if ocr_text:
            ocr_source_label = "user_provided"
        else:
            embedded = extract_pdf_text(image_path)
            ocr_source_label = "pdf_embedded_text" if embedded else "pytesseract"
            ocr_text = embedded or ""

        # Classify on the FIRST page only — class is dossier-level, not per-page
        first_page_ocr = run_ocr(pages[0], self.cfg)
        doc_class, doc_conf = self.classify(pages[0], first_page_ocr.words, first_page_ocr.boxes)
        log.info(f"Class: {doc_class}  (confidence: {doc_conf:.1%})")

        result = InferenceResult(
            image=str(image_path),
            doc_class=doc_class,
            doc_confidence=round(doc_conf, 4),
            pages_processed=len(pages),
            ocr_source=ocr_source_label,
        )

        # Extract fields from EVERY page; merge at the end
        if doc_class not in self.cfg.needs_extraction:
            log.info(f"No field extraction needed for class '{doc_class}'")
            return result

        all_fields: dict[str, FieldExtraction] = {}
        ocr_text_by_page: list[str] = []
        for page_idx, page_img in enumerate(pages):
            page_ocr = first_page_ocr if page_idx == 0 else run_ocr(page_img, self.cfg)
            if not page_ocr.words or page_ocr.source == "fallback":
                log.warning(f"Page {page_idx + 1}: no usable OCR, skipping")
                continue
            ocr_text_by_page.append(page_ocr.text)
            page_fields = self.extract(page_img, page_ocr.words, page_ocr.boxes)

            # Keep highest-confidence value when the same field appears on multiple pages
            for name, extraction in page_fields.items():
                if name not in all_fields or extraction.confidence > all_fields[name].confidence:
                    all_fields[name] = extraction

        # Post-process: strip form-label noise, validate formats, fill gaps via OCR-regex
        full_ocr_text = " ".join(ocr_text_by_page)
        result.fields = _clean_field_extractions(all_fields, full_ocr_text)

        # Per-class allowlist: drop fields that don't belong to this document type
        if doc_class in CLASS_FIELDS:
            allowed = CLASS_FIELDS[doc_class]
            result.fields = {k: v for k, v in result.fields.items() if k in allowed}
        if result.fields:
            log.info(f"Extracted {len(result.fields)} field(s):")
            for name, ext in result.fields.items():
                log.info(f"   · {name}: {ext.value!r}  (conf: {ext.confidence:.1%})")
        else:
            log.info("No fields extracted")

        return result


# ────────────────────────────────────────────────────────────────────────────
# CLI entry point
# ────────────────────────────────────────────────────────────────────────────
def _save_result(result: InferenceResult, image_path: Path, cfg: Config) -> Path:
    out_dir = Path(cfg.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"{image_path.stem}_result.json"
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(result.to_dict(), f, ensure_ascii=False, indent=2)
    return out_path


def _prompt_for_image_path() -> str | None:
    """GUI fallback ONLY when running interactively. Skipped on headless servers."""
    if not sys.stdin.isatty():
        return None
    try:
        from tkinter import Tk, filedialog
        root = Tk()
        root.withdraw()
        root.attributes("-topmost", True)
        path = filedialog.askopenfilename(
            title="Select a document",
            filetypes=[
                ("Documents", "*.png *.jpg *.jpeg *.pdf *.bmp *.tif *.tiff"),
                ("All files", "*.*"),
            ],
        )
        root.destroy()
        return path or None
    except Exception as e:
        log.debug(f"GUI prompt unavailable: {e}")
        return None


def main():
    parser = argparse.ArgumentParser(description="GuichetOI ML — document classification + field extraction")
    parser.add_argument("image", nargs="?", help="Path to document (image or PDF)")
    parser.add_argument("--image", dest="image_flag", help="Path to document (alternative to positional arg)")
    parser.add_argument("--ocr",   default="", help="Pre-extracted OCR text (skips Tesseract)")
    parser.add_argument("--device", default=None, choices=[None, "cpu", "cuda"], help="Force device")
    args = parser.parse_args()

    image_path = args.image_flag or args.image or _prompt_for_image_path()
    if not image_path:
        parser.error("No image path provided. Use --image PATH or run interactively.")

    try:
        cfg = Config()
        pipeline = GuichetOIPipeline(cfg=cfg, device=args.device)
        result = pipeline.run(image_path, args.ocr)
        out_path = _save_result(result, Path(image_path), cfg)
        log.info(f"Saved: {out_path}")
        return 0
    except FileNotFoundError as e:
        log.error(f"File not found: {e}")
        return 2
    except Exception as e:
        log.exception(f"Inference failed: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())