"""
Lädt mit prepare_corpus.py erzeugte Docling-Dokumente als LangChain-Dokumente.

Bereinigung (clean_blocks, reines Python, auch von grobid_qa.py benutzt):
  - Seitenränder (Kolumnentitel, Seitenzahlen) fallen weg
  - OCR-Zeichensalat vom Buchrand ("£", "X N £", "CE EEE VKU") wird verworfen,
    Salat vor einem echten Absatz abgeschnitten
  - an Spalten- oder Seitenwechseln getrennte Absätze werden wieder zusammengeführt,
    Trennstriche aus Zeilenumbrüchen entfernt ("Veröf- fentlichung" -> "Veröffentlichung")
  - Fußnotenzeichen im Fließtext werden entfernt ("Kafka.1", "Felice 6 , die")
  - als Punkt gelesene Kommas werden repariert ("Einzelheiten. die" -> "Einzelheiten, die"),
    als "©" gelesene Trennstriche am Zeilenrand entfernt ("Atmo- © sphäre")
  - bekannte Kodierungsfehler einiger Verlags-PDFs werden ersetzt (" ¼ " -> " = ")

Jeder Chunk beginnt mit einer Kopfzeile aus Autor, Titel (aus dem Dateinamen) und
Abschnittspfad, z. B. "Lehmann, Hans-Thies: Der Buchstaebliche Koerper — 3. Schrift".
PyPDFLoader-Chunks tragen den Autor oft zufällig über den Kolumnentitel mit sich;
Docling sortiert Kolumnentitel aus und legt Überschriften in die Metadaten. Ohne
Kopfzeile fanden Fragen wie "nach Lehmann" die richtige Stelle deshalb schlechter.

Chunks (load): Fließtext wird pro Abschnitt zusammengefügt und mit demselben
Splitter wie bisher zerlegt, damit Vergleiche mit PyPDFLoader fair bleiben.
Fußnoten, Bildunterschriften und Tabellen werden als eigene Chunks indexiert,
Literaturverzeichnisse standardmäßig nicht. Metadaten je Chunk:
  source        Pfad des Original-PDFs (nicht der neu erkannten Kopie)
  page          Seite ab 0 im Original-PDF, wie bei PyPDFLoader (bei getrennten
                Doppelseiten über <json>.pages.json zurückgerechnet)
  half          "L"/"R" bei getrennten Doppelseiten, sonst ""
  element_type  body | footnote | caption | table | reference
  section       Abschnittspfad aus der Überschriften-Hierarchie ("1. Nachlass › a) Situation 1924")
  extractor     "docling"
"""

import json
import re
import unicodedata
from dataclasses import dataclass, replace
from pathlib import Path

INDEX_REFERENCES = False     # Literaturverzeichnis-Einträge mitindexieren?

LABEL_KIND = {
    "title": "head", "section_header": "head", "text": "body", "paragraph": "body",
    "list_item": "body", "footnote": "footnote", "caption": "caption", "formula": "body",
    "code": "body", "reference": "reference", "document_index": "skip",
    "page_header": "skip", "page_footer": "skip",
}

TERMINAL = tuple('.!?:…"\'“”‘’»«)]')
HYPHEN_KEEP = {"und", "oder", "bis", "sowie", "als", "wie", "noch", "bzw", "and", "or", "to"}
VOWEL = re.compile(r"[aeiouyäöüAEIOUYÄÖÜ]")
REAL_WORD = re.compile(r"[A-Za-zÄÖÜäöüß]{4,}")
NORMAL_TOKEN = re.compile(r"^[(„»\"']?[A-Za-zÄÖÜäöüß][a-zäöüß]{2,}[.,;:!?)»«\"“”']*$|^[(\[]?\d[\d:.,\-–]*[)\]]?[.,;:]?$")

# Kodierungsfehler aus Verlags-PDFs ohne Unicode-Zuordnung (siehe grobid_qa.py)
ENCODING_FIXES = [
    (re.compile(r"(?<=\s)¼(?=\s)"), "="),
    (re.compile(r"(?<=[\s(])À(?=\s?[\d.])"), "−"),
    (re.compile(r"(?<=\w) Â (?=\w)"), " × "),
    (re.compile(r"\bÓ (?=\d{4})"), "© "),
]
# Fußnotenzeichen am Wort ("Kafka1 ") oder nach Satzzeichen ("Kafka.1 ")
FNMARK_GLUED = re.compile(r"(?<=[a-zäöüß]{3})\d{1,3}(?=\s|$)")
FNMARK_AFTER_PUNCT = re.compile(r"(?<=[a-zäöüß][.,;:!?»“\"])\d{1,3}(?=\s|$)")
FNMARK_SPACED = re.compile(
    r"(?<=[A-Za-zÄÖÜäöüß]{3}[.,;:!?»«\"”)]) (\d{1,3})(?: \d)? (?=[,;:]|[A-ZÄÖÜ„»])"
    r"|(?<=[a-zäöüß]{3}) (\d{1,3})(?: \d)? (?=[,;:])")

# Abkürzungen, nach denen ein kleingeschriebenes Wort normal ist ("vgl. die", "Hrsg. von")
ABBREVIATIONS = {
    "vgl", "bzw", "usw", "ebd", "sog", "ders", "dies", "hrsg", "hg", "anm", "aufl", "bd", "bde",
    "nr", "verf", "übers", "eigtl", "wörtl", "bes", "ggf", "inkl", "evtl", "insb", "ca", "etc",
    "kap", "abschn", "zit", "ff", "jh", "jhs", "dgl", "ugs", "lat", "griech", "frz", "engl",
    "geb", "gest", "orig", "ausg", "hs", "hss", "fol", "masch", "diss", "phil",
}
PERIOD_BEFORE_LOWER = re.compile(r"\b([A-Za-zÄÖÜäöüß]*[a-zäöüß]{2})\. (?=[a-zäöüß]{2,}\b)")


def source_label(filename: str) -> str:
    """Autor und Titel aus Dateinamen wie 'Lehmann,_Hans-Thies_DerBuchstaeblicheKoerper_Zur….pdf'."""
    stem = Path(filename).stem.strip(" ._")
    tokens = [t for t in stem.split("_") if t]
    authors, rest, i = [], [], 0
    while i < len(tokens) and tokens[i].endswith(","):        # "Nachname," "Vorname"
        first = tokens[i + 1] if i + 1 < len(tokens) else ""
        authors.append(f"{tokens[i].rstrip(',')}, {first}".strip(", "))
        i += 2
    rest = tokens[i:]
    # Binnenmajuskeln trennen: "DerBuchstaeblicheKoerper" -> "Der Buchstaebliche Koerper"
    title = " – ".join(re.sub(r"(?<=[a-zäöüß])(?=[A-ZÄÖÜ])", " ", t) for t in rest)
    if not authors:                                            # "Kafka_Die_Verwandlung_"
        authors, title = [rest[0]] if rest else [], " ".join(rest[1:])
    return "; ".join(authors) + (f": {title}" if title else "")


def chunk_header(source: str, section: str = "") -> str:
    """Kopfzeile, die jedem Chunk vorangestellt wird (Autor, Titel, Abschnitt)."""
    return source_label(Path(source).name) + (f" — {section}" if section else "") + "\n\n"


def _repair_periods(text: str) -> str:
    """OCR liest Kommas oft als Punkte. Im Deutschen folgt auf einen Satzpunkt kein
    kleingeschriebenes Wort, nach Abkürzungen schon ("Einzelheiten. die" -> "Einzelheiten, die")."""
    return PERIOD_BEFORE_LOWER.sub(
        lambda m: m.group(0) if m.group(1).lower() in ABBREVIATIONS else m.group(1) + ", ", text)


def period_anomalies_per_10k(text: str) -> float:
    """Wie oft ein Punkt vor einem kleingeschriebenen Wort steht (Abkürzungen ausgenommen)."""
    hits = sum(m.group(1).lower() not in ABBREVIATIONS for m in PERIOD_BEFORE_LOWER.finditer(text))
    return round(10000 * hits / max(1, len(text.split())), 1)


@dataclass
class Block:
    kind: str          # head, body, footnote, caption, table, reference
    text: str
    page: int | None   # ab 1, wie Docling
    section: str = ""
    # Seitenwechsel innerhalb eines zusammengeführten Absatzes: (Zeichenposition, Seite)
    page_breaks: tuple = ()
    level: int = 1     # Ebene einer Überschrift (Docling, HeadingHierarchyOptions)


def _normalize(text: str) -> str:
    text = unicodedata.normalize("NFC", text).replace("­", "")
    text = " ".join(text.split())
    # "©" ohne Jahreszahl ist ein falsch gelesener Trennstrich oder Randfleck ("Atmo- © sphäre")
    text = re.sub(r"\s[©®](?!\s?\d{4})(?=\s)", "", text)
    # Trennstrich aus Zeilenumbruch innerhalb eines Blocks
    text = re.sub(r"(?<=[a-zäöüß])- (?=([a-zäöüß]+))",
                  lambda m: "- " if m.group(1) in HYPHEN_KEEP else "", text)
    for pattern, repl in ENCODING_FIXES:
        text = pattern.sub(repl, text)
    return _repair_periods(text)


def is_garbage(text: str, kind: str = "body") -> bool:
    """OCR-Rauschen vom Buchrand oder aus dem Falz: kaum Buchstaben, keine echten Wörter.
    Fußnoten sind oft nur Belege ("11 T S. 297.", "Vgl. Anm. 13.") und werden milder geprüft."""
    nonspace = [c for c in text if not c.isspace()]
    if not nonspace:
        return True
    letters = sum(c.isalpha() for c in nonspace)
    if kind == "footnote":
        return letters < 2 or letters / len(nonspace) < 0.25
    if letters / len(nonspace) < 0.5:
        return True
    real = [w for w in REAL_WORD.findall(text) if VOWEL.search(w)]
    if len(text) < 40 and not real:
        return True
    words = text.split()
    if len(text) < 80 and sum(len(w) for w in words) / len(words) < 2.5:
        return True
    # lange Blöcke aus OCR-Salat (z. B. erkannte Handschrift-Faksimiles): kaum echte Wörter
    return len(words) >= 8 and real_word_share(words) < 0.3


def real_word_share(words: list[str]) -> float:
    """Anteil 'echter' Wörter: mind. 4 Buchstaben, mit Vokal, überwiegend klein geschrieben.
    Deutsche Prosa liegt bei etwa 0,5–0,6, OCR-Salat wie 'Bohran pn SE Hin A At' weit darunter."""
    real = sum(bool(re.fullmatch(r"[„»(\"']?[A-ZÄÖÜa-zäöüß]?[a-zäöüß]{3,}[.,;:!?)«“\"']*", w))
               and bool(VOWEL.search(w)) for w in words)
    return real / max(1, len(words))


def strip_garbage_prefix(text: str) -> str:
    """Schneidet Zeichensalat vor einem echten Absatz ab: 'ET EST EEE "24 MEF = Re Doppler (1973:85) …'"""
    tokens = text.split(" ")
    i = 0
    while i < len(tokens) and not NORMAL_TOKEN.match(tokens[i]):
        i += 1
    prefix = " ".join(tokens[:i])
    if i >= 4 and re.search(r"[\"'=£®&>|~<{}]", prefix):
        return " ".join(tokens[i:]).lstrip()
    return text


def _strip_footnote_marks(text: str) -> str:
    text = FNMARK_AFTER_PUNCT.sub("", FNMARK_GLUED.sub("", text))
    return FNMARK_SPACED.sub(" ", text)


def _join(a: str, b: str) -> str:
    """Zwei Teile eines getrennten Absatzes verbinden, Trennstrich am Ende auflösen."""
    if a.endswith("-") and b[:1].islower():
        first = b.split(" ", 1)[0].strip(".,;:")
        if first in HYPHEN_KEEP:         # "Primär- und Sekundärtexte" bleibt
            return a + " " + b
        return a[:-1] + b                # "Veröf-" + "fentlichung"
    return a + " " + b


def read_blocks(path: Path) -> list[Block]:
    """Docling-JSON in Lesereihenfolge, ohne Seitenränder, Text normalisiert."""
    doc = json.loads(Path(path).read_text(encoding="utf-8"))
    blocks: list[Block] = []

    def resolve(ref):
        _, coll, idx = ref.split("/")
        return coll, doc[coll][int(idx)]

    def walk(node):
        for child in node.get("children", []):
            coll, item = resolve(child["$ref"])
            prov = item.get("prov") or []
            page = prov[0].get("page_no") if prov else None
            if coll == "texts" and item.get("content_layer", "body") == "body":
                kind = LABEL_KIND.get(item.get("label"), "body")
                if kind != "skip":
                    raw = item.get("text", "")
                    text = _normalize(raw)
                    # Docling führt Absätze teils über Seiten hinweg; die Seitenwechsel
                    # stehen in prov (charspan bezogen auf den Rohtext, hier umgerechnet)
                    scale = len(text) / max(1, len(raw))
                    breaks = tuple((round(p["charspan"][0] * scale), p.get("page_no"))
                                   for p in prov[1:] if p.get("charspan") and p.get("page_no") != page)
                    blocks.append(Block(kind, text, page, page_breaks=breaks, level=item.get("level") or 1))
            elif coll == "tables" and item.get("content_layer", "body") == "body":
                cells = (item.get("data") or {}).get("table_cells", [])
                blocks.append(Block("table", _normalize(" ".join(c.get("text", "") for c in cells)), page))
            walk(item)

    walk(doc.get("body", {}))
    return [b for b in blocks if b.text]


def clean_blocks(path: Path) -> list[Block]:
    """Bereinigte Blöcke mit Abschnittszuordnung, siehe Moduldoku."""
    out: list[Block] = []
    section, stack = "", []              # stack: (Ebene, Überschrift) der offenen Abschnitte
    last_body = None                     # Index des letzten Fließtextblocks in out
    for b in read_blocks(path):
        if b.kind != "table":
            b = replace(b, text=strip_garbage_prefix(b.text))
            if is_garbage(b.text, b.kind):
                continue
        if b.kind == "head":
            stack = [(lvl, txt) for lvl, txt in stack if lvl < b.level] + [(b.level, b.text)]
            section = " › ".join(txt for _, txt in stack)
            last_body = None             # über Überschriften hinweg nie zusammenführen
            out.append(replace(b, section=section))
            continue
        if b.kind == "body":
            text = _strip_footnote_marks(b.text)
            prev = out[last_body] if last_body is not None else None
            if prev is not None and not prev.text.endswith(TERMINAL):
                merged = _join(prev.text, text)
                start = len(merged) - len(text)
                breaks = prev.page_breaks + ((start, b.page),) + tuple((start + o, p) for o, p in b.page_breaks)
                out[last_body] = replace(prev, text=merged, page_breaks=breaks)
                continue
            out.append(replace(b, text=text, section=section))
            last_body = len(out) - 1
            continue
        out.append(replace(b, section=section))
    return out


def load(json_path: Path, source: str, splitter, header: bool = True) -> list:
    """LangChain-Dokumente, gechunkt mit `splitter` (RecursiveCharacterTextSplitter).
    header=True stellt jedem Chunk Autor, Titel und Abschnitt voran (chunk_header)."""
    from langchain_core.documents import Document

    blocks = clean_blocks(json_path)
    # Getrennte Doppelseiten: neue Seite -> Seite im Original-PDF (split_spreads.py)
    pages_file = Path(json_path).with_name(Path(json_path).stem + ".pages.json")
    page_map = json.loads(pages_file.read_text(encoding="utf-8")) if pages_file.exists() else None
    kinds = ["body", "footnote", "caption", "table"] + (["reference"] if INDEX_REFERENCES else [])
    docs = []
    for kind in kinds:
        # Blöcke einer Art pro Abschnitt zusammenfassen, Seitenanfänge merken
        groups: list[list[Block]] = []
        for b in blocks:
            if b.kind != kind:
                continue
            if groups and groups[-1][-1].section == b.section:
                groups[-1].append(b)
            else:
                groups.append([b])
        for group in groups:
            text, starts = "", []
            for b in group:
                if text:
                    text += "\n\n"
                starts.append((len(text), b.page))
                starts += [(len(text) + offset, page) for offset, page in b.page_breaks]
                text += b.text
            def original_page(page):
                """Docling-Seite (ab 1) -> (Seite im Original-PDF ab 0, Hälfte)."""
                if page_map and page and page <= len(page_map):
                    return page_map[page - 1]["original_page"], page_map[page - 1]["half"]
                return ((page - 1) if page else None), ""

            # Abschnitte des Texts mit ihrer Seite: [(Anfang, Ende, Seite), ...]
            spans = [(s, e, p) for (s, p), (e, _) in zip(starts, starts[1:] + [(len(text), None)])]
            cursor = 0
            for chunk in splitter.split_text(text):
                pos = text.find(chunk[:200], cursor)
                pos = pos if pos >= 0 else max(text.find(chunk[:200]), 0)
                cursor = pos + 1
                # Seite, auf der der größte Teil des Chunks steht (nicht nur sein Anfang)
                share = {}
                for s, e, p in spans:
                    overlap = min(e, pos + len(chunk)) - max(s, pos)
                    if overlap > 0:
                        key = original_page(p)
                        share[key] = share.get(key, 0) + overlap
                original, half = max(share, key=share.get) if share else original_page(group[0].page)
                content = (chunk_header(source, group[0].section) if header else "") + chunk
                docs.append(Document(page_content=content, metadata={
                    "source": source,
                    "page": original,
                    "half": half,
                    "element_type": kind,
                    "section": group[0].section,
                    "extractor": "docling",
                }))
    return docs