"""
PDFs mit Docling extrahieren, zum Vergleich mit GROBID und PyPDFLoader.

Vier Varianten:
  docling       nutzt die vorhandene Textschicht des PDFs (bei Scans: deren OCR)
  docling_ocr   erkennt jede Seite neu per OCR (Apple Vision über ocrmac, Deutsch),
                ignoriert also die mitgelieferte, oft schlechte OCR-Schicht.
                Achtung: lässt bei diesen Scans ganze Absätze aus.
  docling_tess  wie docling_ocr, aber mit Tesseract und dem deutschen Modell
                aus tessdata/deu.traineddata (tessdata_best). Verliert ebenfalls Text:
                Doclings eigener OCR-Ablauf ist bei diesen Doppelseiten-Scans unzuverlässig.
  docling_ocrmypdf
                Arbeitsteilung: OCRmyPDF (Tesseract, deutsch, mit Geraderichten) legt
                eine neue Textschicht ins PDF (extract_out/ocrmypdf/), Docling liest
                diese dann ohne eigene OCR und erkennt die Struktur.

Läuft in einer eigenen Umgebung, weil Docling PyTorch und Layout-Modelle mitbringt:
    conda activate docling
    python extract_docling.py                 # Variante docling_tess, Titel aus TITLES
    python extract_docling.py docling         # andere Variante

Ausgabe: extract_out/<variante>/<pdf>.json (verlustfrei, für grobid_qa.py)
und .md (zum Lesen). Vorhandene Ergebnisse werden übersprungen.
"""

import os
import subprocess
import sys
import time
from pathlib import Path

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import (
    OcrMacOptions, OcrMode, PdfPipelineOptions, TesseractCliOcrOptions)
from docling.document_converter import DocumentConverter, PdfFormatOption

PDF_DIR = Path("data")
OUT_ROOT = Path("extract_out")
TESSDATA = Path("tessdata").resolve()      # enthält deu.traineddata
# Docling prüft die verfügbaren Tesseract-Sprachen ohne Rücksicht auf `path`,
# deshalb das Modellverzeichnis auch über die Umgebungsvariable bekanntmachen.
os.environ["TESSDATA_PREFIX"] = str(TESSDATA)
VARIANTS = ("docling", "docling_ocr", "docling_tess", "docling_ocrmypdf")
OCRMYPDF_DIR = OUT_ROOT / "ocrmypdf"         # neu erkannte PDFs für docling_ocrmypdf
# Auswahl für den Vergleich: Dateinamen beginnen mit diesen Präfixen
TITLES = [
    "Beissner,",                                  # Aufsatz, viele Fußnoten
    "Fiechter,",                                  # GROBID: alles ein Absatz
    "Trost,",                                     # sehr schlechte OCR-Schicht
    "Harzer,",                                    # GROBID: Fließtext als Abbildung
    "Binder,_Hartmut_KafkaHandbuch_Band2_AText",  # großes Buch
]


def converter(variant: str) -> DocumentConverter:
    opts = PdfPipelineOptions()
    opts.do_table_structure = True
    if variant == "docling_ocr":
        opts.do_ocr = True
        opts.ocr_options = OcrMacOptions(mode=OcrMode.FULL_PAGE, lang=["de-DE"])
    elif variant == "docling_tess":
        opts.do_ocr = True
        opts.ocr_options = TesseractCliOcrOptions(mode=OcrMode.FULL_PAGE, lang=["deu"], path=str(TESSDATA))
    else:                                      # docling, docling_ocrmypdf: Textschicht lesen
        opts.do_ocr = False
    return DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=opts)})


def ocrmypdf(pdf: Path) -> Path:
    """Legt mit OCRmyPDF eine neue deutsche Textschicht an (überspringt Vorhandenes)."""
    OCRMYPDF_DIR.mkdir(parents=True, exist_ok=True)
    target = OCRMYPDF_DIR / pdf.name
    if not target.exists():
        subprocess.run(
            [str(Path(sys.executable).parent / "ocrmypdf"), "--force-ocr", "-l", "deu", "--deskew",
             "--output-type", "pdf", "--jobs", "4", "--quiet", str(pdf), str(target)],
            check=True)
    return target


def main():
    variant = sys.argv[1] if len(sys.argv) > 1 else "docling_tess"
    if variant not in VARIANTS:
        sys.exit(f"Variante muss eine von {', '.join(VARIANTS)} sein")
    out_dir = OUT_ROOT / variant
    out_dir.mkdir(parents=True, exist_ok=True)

    pdfs = [p for p in sorted(PDF_DIR.glob("*.pdf")) if any(p.name.startswith(t) for t in TITLES)]
    conv = converter(variant)
    for i, pdf in enumerate(pdfs, 1):
        target = out_dir / f"{pdf.stem}.json"
        if target.exists():
            print(f"[{i}/{len(pdfs)}] {pdf.name}: vorhanden, übersprungen", flush=True)
            continue
        print(f"[{i}/{len(pdfs)}] {pdf.name} ...", end=" ", flush=True)
        t0 = time.time()
        try:
            source = ocrmypdf(pdf) if variant == "docling_ocrmypdf" else pdf
            doc = conv.convert(source).document
        except Exception as e:      # ein kaputtes PDF soll den Lauf nicht abbrechen
            print(f"FEHLER: {e}", flush=True)
            continue
        doc.save_as_json(target)
        doc.save_as_markdown(out_dir / f"{pdf.stem}.md")
        print(f"fertig ({time.time() - t0:.0f} s)", flush=True)


if __name__ == "__main__":
    main()