Aivanov_scan_ogc/pipeline/extract.py

"""Orchestration d'extraction pour un dossier OGC.

Chaîne les étages du pipeline sans connaître leur implémentation interne :

  ingest → routing → OCR page par page → enrichissement page-spécifique → validation ATIH

L'orchestration elle-même ne contient aucune logique métier : elle délègue à
`pipeline.recueil`, `pipeline.validation`, `pipeline.classify`, `pipeline.ocr_qwen`
et `pipeline.prompts`. Cela permet de tester indépendamment chaque étage.
"""
from __future__ import annotations

import time
from pathlib import Path

from .classify import detect_page_type, route_by_index
from .ingest import pdf_to_images
from .json_utils import parse_json_output
from .ocr_qwen import QwenVLOCR
from .prompts import PAGE_TYPES, PROMPT_HEADER
from .recueil import enrich_recueil, resolve_recueil_zones
from .validation import annotate as validate_annotate


def _run_page_ocr(ocr: QwenVLOCR, image_path: Path, ptype: str) -> tuple[dict | None, str, float]:
    """Exécute le prompt principal associé à un type de page et parse le JSON.

    Retourne (parsed_dict_ou_None, ocr_raw, elapsed_s). `parsed=None` quand
    la page n'a pas de prompt structuré associé (concertation_med, hospit.).
    """
    conf = PAGE_TYPES.get(ptype)
    if not conf or conf["prompt"] == PROMPT_HEADER:
        return None, "", 0.0
    res = ocr.run(image_path, conf["prompt"], max_new_tokens=4096)
    parsed = parse_json_output(res["text"])
    return parsed, res["text"], round(res["elapsed_s"], 2)


def _resolve_routing(images: list[Path], ocr: QwenVLOCR,
                     use_standard_routing: bool,
                     verbose: bool) -> tuple[list[str | None], list[str]]:
    """Détermine le type de chaque page, soit par ordre standard (une seule
    vérification sur la page 1), soit par classification OCR page par page.

    Retourne (page_types, headers). `headers[i]` est vide si pas de classify
    effectuée sur la page i.
    """
    page_types: list[str | None] = [None] * len(images)
    headers: list[str] = [""] * len(images)

    if use_standard_routing and images:
        ptype1, header1 = detect_page_type(images[0], ocr)
        if ptype1 == "recueil":
            page_types = list(route_by_index(len(images)))
            headers[0] = header1
            if verbose:
                print("  routing standard (page 1 = recueil OK)")
            return page_types, headers
        if verbose:
            print(f"  page 1 = {ptype1} → fallback classification")

    # Fallback : classify page par page
    for i, img in enumerate(images):
        page_types[i], headers[i] = detect_page_type(img, ocr)
    return page_types, headers


def extract_dossier(pdf_path: str | Path, verbose: bool = True,
                    use_standard_routing: bool = True) -> dict:
    """Pipeline complet d'un dossier : PDF → JSON structuré + annoté ATIH.

    Étages :
    1. `ingest.pdf_to_images`    : PDF → PNG 300 dpi (avec deskew auto, cache)
    2. `_resolve_routing`        : type de chaque page
    3. `_run_page_ocr`           : OCR du schéma structuré par type de page
    4. `recueil.enrich_recueil`  : checkboxes + crop Recodage pour la page recueil
    5. `validation.annotate`     : validation ATIH de tous les codes extraits

    Paramètre `use_standard_routing=True` exploite l'ordre standard des 6 pages
    OGC et économise 5 appels OCR par dossier. Bascule automatique sur la
    classification page-à-page si la page 1 n'est pas le recueil attendu.
    """
    pdf_path = Path(pdf_path)
    ocr = QwenVLOCR()
    if verbose:
        print(f"[{pdf_path.name}] modèle prêt, VRAM={ocr.vram_gb:.2f} Go")

    images = pdf_to_images(str(pdf_path))
    if verbose:
        print(f"[{pdf_path.name}] {len(images)} pages converties")

    page_types, headers = _resolve_routing(images, ocr, use_standard_routing, verbose)

    _, cb_zones = resolve_recueil_zones()

    result: dict = {
        "fichier": pdf_path.stem,
        "pdf_hash": images[0].parent.name if images else "",
        "pages": [],
        "extraction": {},
    }

    for idx, img_path in enumerate(images, 1):
        t0 = time.time()
        ptype = page_types[idx - 1]
        header_text = headers[idx - 1]
        page_info: dict = {
            "page": idx,
            "type": ptype,
            "header": header_text.strip(),
            "elapsed_s": None,
        }
        if verbose:
            print(f"  p{idx}: {ptype}")

        parsed, ocr_raw, elapsed = _run_page_ocr(ocr, img_path, ptype) if ptype else (None, "", 0.0)
        if parsed is not None:
            page_info["ocr_raw"] = ocr_raw
            page_info["parsed"] = parsed
            page_info["elapsed_s"] = elapsed

            if ptype == "recueil" and isinstance(parsed, dict):
                enrich_recueil(parsed, img_path, ocr, cb_zones)
                page_info["parsed"] = parsed

            result["extraction"][ptype] = parsed
        else:
            page_info["elapsed_s"] = round(time.time() - t0, 2)

        result["pages"].append(page_info)

    return validate_annotate(result)