"""
Bereitet ein PDF-Korpus für den Index auf: bessere OCR, wo nötig, dann Docling.

Für jedes PDF:
  0. Doppelseiten-Scans (Querformat) in Einzelseiten trennen (split_spreads.py),
     damit OCR und Layout-Erkennung nie über den Falz lesen.
  1. Textschicht prüfen (pypdf): fehlt sie oder ist die mitgelieferte OCR schlecht,
     wird mit OCRmyPDF (Tesseract, deutsches Modell aus tessdata/) eine neue angelegt —
     und nur behalten, wenn sie nachweislich besser ist (OCR-Fehler plus als Punkt
     gelesene Kommas). Gute Textschichten bleiben unangetastet.
  2. Docling liest die (alte oder neue) Textschicht ohne eigene OCR und erkennt
     Absätze, Überschriften, Fußnoten, Seitenränder. Doclings eigene OCR verliert
     bei diesen Scans Text und wird deshalb nicht benutzt.
  3. Qualitätsprüfung mit grobid_qa.py auf dem bereinigten Ergebnis (wie es der
     Loader liefert), festgehalten in manifest.csv.

Läuft in der Docling-Umgebung:
    conda activate docling
    python prepare_corpus.py                         # PDFs aus data/
    python prepare_corpus.py data_test               # anderer Ordner
    python prepare_corpus.py data --only "Trost;Harzer"  # nur Titel mit diesen Anfängen
    python prepare_corpus.py data --no-split         # Doppelseiten nicht trennen
    python prepare_corpus.py data --layout egret_large --no-headings  # Docling-Varianten

Ausgabe in corpus_processed/<ordner>_split/ (bzw. <ordner>/ mit --no-split):
  split/<pdf>.pdf    getrennte Fassung + .pages.json (Seite im Original je neuer Seite)
  <pdf>.json         Docling-Dokument (Eingabe für docling_loader.py)
  ocr/<pdf>.pdf      neu erkannte PDFs (nur wo nötig)
  qa/<pdf>.md        QA-Bericht mit Belegstellen
  manifest.csv       Entscheidung, Kennzahlen und QA-Urteil pro PDF
Bereits Verarbeitetes wird übersprungen; zum Neuaufbau den Ordner löschen.
"""

import argparse
import csv
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

from docling.datamodel.base_models import InputFormat
from docling.datamodel import pipeline_options as po
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption

import grobid_qa as qa
import split_spreads
from docling_loader import period_anomalies_per_10k

PDF_DIR = Path("data")
OUT_ROOT = Path("corpus_processed")
TESSDATA = Path("tessdata").resolve()        # deu.traineddata + configs
os.environ["TESSDATA_PREFIX"] = str(TESSDATA)

# Ab wann die mitgelieferte Textschicht als schlecht gilt und neu erkannt wird.
# Kalibriert am Vergleich von fünf Titeln: Harzer (13.6) blieb mit alter Schicht
# besser, Beissner (47), Binder (39), Fiechter (63), Trost (264) wurden deutlich besser.
REOCR_OCR_PER_10K = 20       # OCR-Fehlerbilder je 10.000 Wörter im pypdf-Rohtext
REOCR_UNUSABLE_PCT = 10      # Seiten, auf denen pypdf nur Zeichensalat liefert, %
REOCR_SCAN_PCT = 1           # Seiten ganz ohne Textschicht, %

# Layout-Modelle von Docling; heron ist der Standard
LAYOUT_MODELS = {
    "heron": po.DOCLING_LAYOUT_HERON,
    "egret_medium": po.DOCLING_LAYOUT_EGRET_MEDIUM,
    "egret_large": po.DOCLING_LAYOUT_EGRET_LARGE,
    "egret_xlarge": po.DOCLING_LAYOUT_EGRET_XLARGE,
}

MANIFEST_COLS = [
    "datei", "seiten", "getrennt", "entscheidung", "grund", "roh_ocr_per_10k", "roh_unbrauchbar_pct",
    "roh_scan_pct", "roh_punkt_per_10k", "neu_ocr_per_10k", "neu_punkt_per_10k", "sekunden", "qa_urteil", "absaetze", "fussnoten", "giant_para_pct",
    "lost_pct", "ocr_per_10k", "fnmark_per_10k", "contam_per_100k", "broken_per_100p", "order",
]


def assess_text_layer(pdf: Path) -> dict:
    """Kennzahlen der vorhandenen Textschicht, wie pypdf sie liest."""
    pages = qa.read_pdf_pages(pdf)
    raw = "\n".join(pages)
    words = len(raw.split()) or 1
    ocr = len(qa.OCR_RE.findall(qa.URL_RE.sub(" ", raw)))
    return {
        "seiten": len(pages),
        "roh_ocr_per_10k": round(10000 * ocr / words, 1),
        "roh_punkt_per_10k": period_anomalies_per_10k(" ".join(raw.split())),
        "roh_unbrauchbar_pct": round(100 * sum(qa.page_unusable(p) for p in pages) / max(1, len(pages)), 1),
        "roh_scan_pct": round(100 * sum(len(p.strip()) < 200 for p in pages) / max(1, len(pages)), 1),
    }


def decide(m: dict) -> tuple[str, str]:
    reasons = []
    if m["roh_scan_pct"] >= REOCR_SCAN_PCT:
        reasons.append(f"{m['roh_scan_pct']} % Seiten ohne Textschicht")
    if m["roh_ocr_per_10k"] >= REOCR_OCR_PER_10K:
        reasons.append(f"OCR-Fehler {m['roh_ocr_per_10k']} je 10k Wörter")
    if m["roh_unbrauchbar_pct"] >= REOCR_UNUSABLE_PCT:
        reasons.append(f"{m['roh_unbrauchbar_pct']} % unlesbare Seiten")
    return ("neu erkannt", "; ".join(reasons)) if reasons else ("Textschicht behalten", "")


def layer_score(ocr_per_10k: float, punkt_per_10k: float) -> float:
    """Kleiner ist besser. Zerstörte Wörter schaden der Suche mehr als falsche Satzzeichen."""
    return ocr_per_10k + 0.5 * punkt_per_10k


def ocrmypdf(pdf: Path, target: Path):
    target.parent.mkdir(parents=True, exist_ok=True)
    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)


def read_manifest(path: Path) -> dict:
    if not path.exists():
        return {}
    with open(path, encoding="utf-8") as f:
        return {r["datei"]: r for r in csv.DictReader(f, delimiter=";")}


def write_manifest(path: Path, rows: dict):
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=MANIFEST_COLS, extrasaction="ignore", delimiter=";")
        w.writeheader()
        w.writerows(rows[k] for k in sorted(rows))


def main():
    parser = argparse.ArgumentParser(description="Korpus für den Index aufbereiten")
    parser.add_argument("pdf_dir", nargs="?", type=Path, default=PDF_DIR)
    parser.add_argument("--only", default="", help="Anfänge von Dateinamen, durch ; getrennt")
    parser.add_argument("--no-split", action="store_true", help="Doppelseiten nicht trennen")
    parser.add_argument("--layout", default="heron", choices=LAYOUT_MODELS, help="Docling-Layout-Modell")
    parser.add_argument("--no-headings", action="store_true",
                        help="keine Überschriften-Hierarchie (alle Überschriften auf Ebene 1)")
    args = parser.parse_args()
    pdf_dir, split = args.pdf_dir, not args.no_split
    # Getrennte und neu erkannte PDFs liegen im Basisordner und werden von allen
    # Docling-Varianten geteilt; Varianten bekommen einen eigenen Unterordner.
    base_dir = OUT_ROOT / (pdf_dir.name + ("_split" if split else ""))
    variant = "_".join(([args.layout] if args.layout != "heron" else []) + (["flat"] if args.no_headings else []))
    out_dir = base_dir / "varianten" / variant if variant else base_dir
    (out_dir / "qa").mkdir(parents=True, exist_ok=True)
    manifest_path = out_dir / "manifest.csv"
    manifest = read_manifest(manifest_path)

    opts = PdfPipelineOptions()
    opts.do_ocr = False                      # Textschicht lesen, OCR macht OCRmyPDF
    opts.do_table_structure = True
    opts.layout_options = po.LayoutOptions(model_spec=LAYOUT_MODELS[args.layout])
    opts.heading_hierarchy_options = po.HeadingHierarchyOptions(enabled=not args.no_headings)
    conv = DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=opts)})

    pdfs = sorted(pdf_dir.glob("*.pdf"))
    if args.only:
        prefixes = [p.strip() for p in args.only.split(";") if p.strip()]
        pdfs = [p for p in pdfs if any(p.name.startswith(x) for x in prefixes)]
    for i, pdf in enumerate(pdfs, 1):
        target = out_dir / f"{pdf.stem}.json"
        if target.exists() and pdf.name in manifest:
            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:
            work, row = pdf, {"datei": pdf.name, "getrennt": "nein"}
            if split and split_spreads.is_spread_pdf(pdf):
                work = base_dir / "split" / pdf.name
                if not work.exists():
                    split_spreads.split_pdf(pdf, work)
                row["getrennt"] = "ja"
            row.update(assess_text_layer(work))
            row["entscheidung"], row["grund"] = decide(row)
            source = work
            if row["entscheidung"] == "neu erkannt":
                candidate = base_dir / "ocr" / pdf.name
                if not candidate.exists():
                    ocrmypdf(work, candidate)
                if row["getrennt"] == "ja":   # Seitenzuordnung gilt auch für die OCR-Fassung
                    shutil.copy(split_spreads.pages_json(work), split_spreads.pages_json(candidate))
                new = assess_text_layer(candidate)
                row["neu_ocr_per_10k"], row["neu_punkt_per_10k"] = new["roh_ocr_per_10k"], new["roh_punkt_per_10k"]
                if layer_score(new["roh_ocr_per_10k"], new["roh_punkt_per_10k"]) < \
                        layer_score(row["roh_ocr_per_10k"], row["roh_punkt_per_10k"]):
                    source = candidate
                else:
                    row["entscheidung"] = "Textschicht behalten"
                    row["grund"] += "; neue OCR nicht besser"
            conv.convert(source).document.save_as_json(target)
            if row["getrennt"] == "ja":       # für docling_loader.py: Seiten im Original
                shutil.copy(split_spreads.pages_json(work), out_dir / f"{pdf.stem}.pages.json")
            # QA auf dem bereinigten Ergebnis; Referenz ist die tatsächlich gelesene Textschicht
            m, ex = qa.check(source, target, cleaned=True)
            qa.write_report(m, ex, out_dir / "qa" / f"{pdf.stem}.md")
        except Exception as e:      # ein kaputtes PDF soll den Lauf nicht abbrechen
            print(f"FEHLER: {e}", flush=True)
            continue
        row.update({k: m.get(k) for k in MANIFEST_COLS if k in m})
        row["qa_urteil"] = m["grobid"]
        row["sekunden"] = round(time.time() - t0)
        manifest[pdf.name] = row
        write_manifest(manifest_path, manifest)
        print(f"{row['entscheidung']}, QA {row['qa_urteil']} ({row['sekunden']} s)", flush=True)


if __name__ == "__main__":
    main()