#!/usr/bin/env python3
"""Stratta — ingest pre-pass.

Reads a norm PDF with PyMuPDF, builds the hierarchical TreeRAG skeleton (chapters
from PDF bookmarks + sub-sections via heading regex, down to X.Y.Z.W), gives
each node the text between its heading and the next one on cleaned pages (no
running headers, no page numbers, hyphenation resolved, clause numbers joined
to their paragraph), and rasterizes each page that contains a figure caption.
Output is a single JSON consumed by the `ingest-norm` skill, which then
enriches sections (formulas, tables, cross-refs, summaries) and uploads the
figures. `warnings[]` lists what the pre-pass could not decide on its own.

Usage:
    python scripts/ingest-prepass.py --pdf <path> --output <dir> [--language fr]

Output (under <dir>/):
    prepass.json       full extracted tree + figure manifest
    figures/figure-<N>.png  rasterized full-page renders for each caption
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from collections import defaultdict
from pathlib import Path

try:
    import fitz  # PyMuPDF
except ImportError:
    sys.stderr.write(
        "ERROR: PyMuPDF not installed. Run: pip install --user pymupdf\n"
    )
    sys.exit(2)


BM_NUM_RE = re.compile(r"^\s*(\d+)\s+(.+?)\s*$")
BM_ANNEX_RE = re.compile(
    r"^\s*(?:ANNEXE|Annexe)\s+([A-Z])(?:\s*\(([^)]+)\))?\s*(.*?)\s*$"
)

SUB_RE = re.compile(
    r"^(\d+(?:\.\d+){1,3})\s+([A-Za-zÀ-ÿ][^\n]{1,140})$", re.MULTILINE
)
TOC_DOT_RE = re.compile(r"\.\s*\.\s*\.")
CLEAN_DOTS = re.compile(r"\s*(\.\s*){3,}.*$")

CAPTION_RE = re.compile(
    r"^\s*(?:Figure|Fig\.|Bild|Abbildung|Figura)\s+(\d+[a-z]?)\b[\s:.\-—–]*(.*)$",
    re.IGNORECASE | re.MULTILINE,
)


def clean(s: str) -> str:
    return re.sub(r"\s+", " ", CLEAN_DOTS.sub("", s)).strip()


# --- page lines ----------------------------------------------------------------
#
# PyMuPDF's plain text follows the PDF's own block order. On a SIA norm the
# clause numbers sit in a column of their own, and on an OCR'd copy (the form
# a bureau actually holds: 60 of the 80 PDFs of the first client corpus are
# scans) that column is read as one block BEFORE the paragraphs, so a page
# arrives as "4.4 / 4.4.1 / 4.4.1.1 / ... / Résistances du terrain /
# Généralités / Sont considérés ..." and no heading regex can see a heading.
#
# Lines are therefore rebuilt from the words and their positions: words on the
# same baseline form a line, left to right. That puts "4.4 Résistances du
# terrain" back together on both the native and the OCR'd copy. A genuine
# two-column page (two text columns, not a numbering column) would be
# scrambled by that rule, so such pages are detected and read in block order.

_LINES_CACHE: dict[tuple[int, int], list[str]] = {}
NUMBER_TOKEN_RE = re.compile(r"^(?:\d{1,2}(?:\.\d{1,3}){0,4}\.?|[A-Z](?:\.\d{1,2}){1,3}|—|-|•)$")
OCR_DOTTED_NUMBER_RE = re.compile(r"^(\d+(?:\.\d+)*) \.(\d)")


def _rows(page: "fitz.Page") -> list[list[tuple]]:
    words = page.get_text("words")  # x0, y0, x1, y1, text, block, line, word
    if not words:
        return []
    heights = sorted(w[3] - w[1] for w in words)
    median = heights[len(heights) // 2]
    # A word set vertically (a library watermark such as "Ecole Polytechnique
    # Fédérale de Lausanne" running up the margin) has a box far taller than
    # the text. Left in, its words land on every line they cross.
    words = [w for w in words if (w[3] - w[1]) <= 2.5 * median]
    if not words:
        return []
    tol = max(2.0, 0.45 * median)
    rows: list[list[tuple]] = []
    for w in sorted(words, key=lambda w: (w[1], w[0])):
        if rows and abs(rows[-1][0][1] - w[1]) <= tol:
            rows[-1].append(w)
        else:
            rows.append([w])
    for row in rows:
        row.sort(key=lambda w: w[0])
    return rows


def _segments(row: list[tuple], gap: float) -> list[list[tuple]]:
    """Split a row where the horizontal gap between words is a column gutter."""
    out: list[list[tuple]] = [[row[0]]]
    for prev, w in zip(row, row[1:]):
        if w[0] - prev[2] > gap:
            out.append([w])
        else:
            out[-1].append(w)
    return out


def _two_column(rows: list[list[list[tuple]]]) -> bool:
    """Two text columns: many rows carry two segments that both read as prose."""
    if len(rows) < 12:
        return False
    both = 0
    for segs in rows:
        texts = [" ".join(w[4] for w in s) for s in segs]
        prose = [t for t in texts if len(t) > 24 and not NUMBER_TOKEN_RE.match(t.split()[0])]
        if len(prose) >= 2:
            both += 1
    return both / len(rows) > 0.3


def page_lines(doc: "fitz.Document", pageno: int) -> list[str]:
    """The lines of a page, in reading order, numbers joined to their text."""
    key = (id(doc), pageno)
    cached = _LINES_CACHE.get(key)
    if cached is not None:
        return cached
    page = doc[pageno]
    rows = _rows(page)
    gap = 0.12 * page.rect.width
    segmented = [_segments(row, gap) for row in rows]
    if _two_column(segmented):
        lines = join_split_headings(page.get_text("text")).splitlines()
    else:
        lines = []
        for segs in segmented:
            text = " ".join(" ".join(w[4] for w in s) for s in segs)
            text = OCR_DOTTED_NUMBER_RE.sub(r"\1.\2", text)
            lines.append(text)
    lines = [l.rstrip() for l in lines]
    _LINES_CACHE[key] = lines
    return lines


def page_text(doc: "fitz.Document", pageno: int) -> str:
    return "\n".join(page_lines(doc, pageno))


def detect_toc_pages(doc: "fitz.Document") -> set[int]:
    out = set()
    for i in range(doc.page_count):
        lines = [l for l in page_lines(doc, i) if l.strip()]
        leader = sum(1 for l in lines if TOC_DOT_RE.search(l))
        if leader >= 5:
            out.add(i)
    return out


# Standalone integers only: the "9" and "1" of a heading "9.1 Délimitation"
# stay, or every "Généralités" heading would read as one running line.
FURNITURE_NUMBER_RE = re.compile(r"(?<![\d.])\d{1,4}(?![\d.])")
FURNITURE_EDGE_RE = re.compile(r"^[\s/|\-–—·.]+|[\s/|\-–—·.]+$")


def furniture_key(line: str) -> str:
    """A running line without the page number it carries.

    `SIA 267, Copyright © 2013 by SIA Zurich 59` on an odd page and
    `64 SIA 267, Copyright © 2013 by SIA Zurich` on an even one are the same
    footer; so are the copies a watermark decorates with a stray slash.
    """
    s = FURNITURE_NUMBER_RE.sub(" ", line)
    s = re.sub(r"\s+", " ", s).strip()
    return FURNITURE_EDGE_RE.sub("", s).strip()


def detect_running_text(doc: "fitz.Document") -> set[str]:
    """Lines appearing on >=5 pages near the top/bottom (running headers/footers),
    keyed without their page number."""
    counter: dict[str, int] = defaultdict(int)
    for i in range(doc.page_count):
        lines = [l.strip() for l in page_lines(doc, i) if l.strip()]
        for l in lines[:3] + lines[-3:]:
            counter[furniture_key(l)] += 1
    return {l for l, c in counter.items() if c >= 5 and len(l) > 8}


CHAP_UPPER_RE = re.compile(
    r"^(\d+)\s+([A-ZÉÈÀÂÔÎÛÇ][^a-z\n]{3,120})\s*$", re.MULTILINE
)
# Norms adopted from CEN (SIA 262.6xx, SIA 267.1xx) set their headings in
# sentence case, which CHAP_UPPER_RE rejects by design. Used only as a second
# pass, when the uppercase form found almost nothing — on a native SIA norm it
# would promote body sentences to chapters.
CHAP_MIXED_RE = re.compile(
    r"^(\d{1,2})\s+([A-ZÀ-Þ][A-Za-zÀ-ÿ][^\n]{2,90})\s*$", re.MULTILINE
)
ANNEX_BODY_RE = re.compile(
    r"^(?:ANNEXE|Annexe)\s+([A-Z])(?:\s*\(([^)]+)\))?\s*(.*?)$", re.MULTILINE
)

# A heading split across two lines: the number alone, the title underneath.
# Comes from the numbering column of CEN-style layouts, where PyMuPDF reads the
# column before the text. Left as-is, every regex below misses the heading.
SPLIT_NUM_RE = re.compile(r"^\s*(\d{1,2}(?:\.\d{1,2}){0,3})\s*$")
SPLIT_TITLE_RE = re.compile(r"^\s*([A-ZÀ-Þ][A-Za-zÀ-ÿ'’\-][^\n]{2,90})\s*$")
# Words that open a continuing sentence, never a heading.
SPLIT_STOP_RE = re.compile(
    r"^(Le |La |Les |Il |Elle |Dans |Pour |Selon |Cette |Ce |Ces |Si |En |Au |Aux |De |Des |Du |Un |Une )",
)


def join_split_headings(text: str) -> str:
    """Rewrite `12\\nTitre` as `12 Titre` so the heading regexes can see it.

    Conservative on purpose: the title line must look like a title (starts
    uppercase, no sentence-ending punctuation, not a sentence opener). A false
    join invents a chapter, which is worse than missing one.
    """
    lines = text.splitlines()
    out: list[str] = []
    i = 0
    while i < len(lines):
        m = SPLIT_NUM_RE.match(lines[i])
        if m and i + 1 < len(lines):
            nxt = lines[i + 1]
            t = SPLIT_TITLE_RE.match(nxt)
            if t and not nxt.rstrip().endswith((".", ",", ";", ":")) and not SPLIT_STOP_RE.match(t.group(1)):
                out.append(f"{m.group(1)} {t.group(1).strip()}")
                i += 2
                continue
        out.append(lines[i])
        i += 1
    return "\n".join(out)


def heading_text(doc: "fitz.Document", pageno: int) -> str:
    """Page text prepared for heading detection."""
    return join_split_headings(page_text(doc, pageno))


NOT_A_TITLE_RE = re.compile(
    r"^(?:EN|SN|ISO|DIN|SIA|NOTE|Tableau|Figure|Table|Bild)\b|^\W|\.$", re.IGNORECASE
)
# A heading cut mid-phrase by the column break: "Résistance à la flexion au".
TRUNCATED_RE = re.compile(
    r"\b(?:ou|et|au|aux|de|des|du|le|la|les|un|une|dans|pour|par|sur|avec|sans|selon|entre)$",
    re.IGNORECASE,
)


def plausible_title(title: str) -> bool:
    """Reject what a numbered line can be other than a heading.

    A stray table cell ("I ~ ~"), a normative reference ("EN 12063:2024 (F)") or
    a wrapped sentence all match the heading shape. Promoting one invents a
    chapter, and an invented chapter swallows the page range of a real one.
    """
    t = title.strip()
    if not 3 <= len(t) <= 90:
        return False
    # A normative heading is capitalised. A lowercase one is a table row that
    # happens to sit behind a number ("2.3 retrait des obstacles").
    if not (t[0].isupper() or t[0].isdigit()):
        return False
    if NOT_A_TITLE_RE.search(t):
        return False
    if SPLIT_STOP_RE.match(t) or TRUNCATED_RE.search(t):
        return False
    letters = sum(1 for c in t if c.isalpha() or c.isspace())
    return letters / len(t) >= 0.7


def dense_heading_pages(doc: "fitz.Document", toc_pages: set[int]) -> set[int]:
    """Pages listing many numbered headings: a contents page, whatever its
    typography. Detecting it by leader dots alone misses the ones set without
    them, and every heading read there points at the contents page instead of
    the section it names."""
    dense: set[int] = set()
    for i in range(doc.page_count):
        if i in toc_pages:
            continue
        found = {m.group(1) for m in SUB_RE.finditer(heading_text(doc, i))}
        # A page of the body drills into one chapter; a contents page walks
        # across several. Counting headings alone would drop a dense page of
        # definitions (3.1 … 3.8), which is real content.
        if len(found) >= 5 and len({n.split(".", 1)[0] for n in found}) >= 3:
            dense.add(i)
    return dense


def prune_chapters(chapters: dict[str, dict]) -> dict[str, dict]:
    """Keep the run of chapter numbers that reads like a table of contents.

    Real chapters are consecutive and move forward through the document. A lone
    "25" between chapters 1 and 2, or a chapter that starts thirty pages before
    the one preceding it, is a detection artefact.
    """
    numbered = sorted(((int(k), k) for k in chapters if k.isdigit()))
    if not numbered:
        return chapters

    # A norm with N detected chapters does not have a chapter 40. Missing a few
    # headings is normal; a number far past the count is a table cell.
    ceiling = 2 * len(numbered) + 3
    numbered = [(n, k) for n, k in numbered if n <= ceiling]
    if not numbered:
        return chapters

    # Longest run whose pages move forward, so one bad page does not discard
    # every chapter after it.
    best = [1] * len(numbered)
    prev = [-1] * len(numbered)
    for i in range(len(numbered)):
        page_i = chapters[numbered[i][1]]["pageStart"]
        for j in range(i):
            if chapters[numbered[j][1]]["pageStart"] <= page_i and best[j] + 1 > best[i]:
                best[i], prev[i] = best[j] + 1, j
    idx = best.index(max(best))
    chain = []
    while idx != -1:
        chain.append(numbered[idx][1])
        idx = prev[idx]

    # Missing a heading or two leaves a small gap; doubling (10 → 20 → 30) means
    # the numbers stopped being chapters and started being table rows. Only
    # from 5 upwards: an OCR'd copy that lost chapters 3 to 6 jumps from 2 to
    # 7, and that is a hole, not a doubling.
    kept: dict[str, dict] = {}
    previous = None
    for key in reversed(chain):
        n = int(key)
        if previous is not None and previous >= 5 and n - previous > 3 and n >= 2 * previous:
            break
        kept[key] = chapters[key]
        previous = n

    # Annexes close a norm, in order. One that lands before the last chapter was
    # read off the contents page, and its page range would swallow the document.
    last_page = max((m["pageStart"] for m in kept.values()), default=0)
    for key, meta in sorted(
        ((k, m) for k, m in chapters.items() if not k.isdigit()), key=lambda kv: kv[0]
    ):
        if meta["pageStart"] < last_page:
            continue
        kept[key] = meta
        last_page = meta["pageStart"]
    return kept


def extract_chapters(
    doc: "fitz.Document", toc_pages: set[int] | None = None
) -> dict[str, dict]:
    """Top-level chapters: prefer PDF bookmarks; fallback to body-text regex
    (uppercase heading on its own line) when the PDF has no bookmarks."""
    chapters: dict[str, dict] = {}
    for _depth, title, page in doc.get_toc(simple=True):
        m = BM_ANNEX_RE.match(title)
        if m:
            letter, kind, rest = m.group(1), m.group(2) or "", (m.group(3) or "").strip()
            path = f"Annexe {letter}"
            t = path + (f" ({kind})" if kind else "")
            if rest and rest.lower() != path.lower():
                t += f" — {rest}"
            chapters[path] = {"title": t, "pageStart": page}
            continue
        m = BM_NUM_RE.match(title)
        if m:
            num, rest = m.group(1), m.group(2).strip()
            chapters[num] = {"title": rest, "pageStart": page}
    if chapters:
        return chapters

    # Fallback: no bookmarks → scan body for chapter headings.
    skip = set(toc_pages or set())

    def scan(pattern: "re.Pattern[str]", skip_pages: set[int]) -> dict[str, dict]:
        found: dict[str, dict] = {}
        for pageno in range(doc.page_count):
            if pageno in skip_pages:
                continue
            text = heading_text(doc, pageno)
            for m in pattern.finditer(text):
                num, title = m.group(1), clean(m.group(2))
                if int(num) > 50 or num in found or not plausible_title(title):
                    continue
                found[num] = {"title": title, "pageStart": pageno + 1}
            for m in ANNEX_BODY_RE.finditer(text):
                letter, kind = m.group(1), m.group(2) or ""
                rest = clean(m.group(3) or "")
                path = f"Annexe {letter}"
                if path in found:
                    continue
                t = path + (f" ({kind})" if kind else "")
                if rest and rest.lower() != path.lower():
                    t += f" — {rest}"
                found[path] = {"title": t, "pageStart": pageno + 1}
        return found

    def toc_like_pages(found: dict[str, dict]) -> set[int]:
        """Pages holding three or more chapter headings are the table of
        contents, whatever the typography. Without this every chapter of such a
        norm starts on the contents page, and every section body is wrong."""
        per_page: dict[int, int] = defaultdict(int)
        for meta in found.values():
            per_page[meta["pageStart"]] += 1
        return {p - 1 for p, n in per_page.items() if n >= 3}

    # A règlement (SIA 144) is cut in articles, "Art. 7 Titre", a Eurocode
    # in "Section 7 Titre". The passes below would take table rows for
    # chapters and lose the whole text outside any section; a labelled
    # layout is unambiguous, so it is read first and wins outright.
    labelled = infer_labelled_chapters(doc, skip)
    if len(labelled) >= 3:
        return recover_missing_chapters(doc, prune_chapters(labelled), skip)

    for pattern in (CHAP_UPPER_RE, CHAP_MIXED_RE):
        chapters = scan(pattern, skip)
        extra = toc_like_pages(chapters)
        if extra:
            chapters = scan(pattern, skip | extra)
        if len([k for k in chapters if k.isdigit()]) >= 3:
            break

    for num, meta in infer_unnumbered_chapters(doc, skip).items():
        chapters.setdefault(num, meta)

    chapters = prune_chapters(chapters)
    return recover_missing_chapters(doc, chapters, skip)


LABEL_RE = re.compile(
    r"^(Art\.?|Section|Chapitre|Kapitel|Abschnitt)\s*(\d{1,3})\b[\s:_.\u2013-]*(.*?)(?:\s+(\d{1,3}))?$"
)


def infer_labelled_chapters(doc: "fitz.Document", skip: set[int]) -> dict[str, dict]:
    """Chapters from labelled headings: "Art. 7 Titre" (a règlement, SIA 144)
    or "Section 7 Titre" (a Eurocode adopted as SIA 267.001).

    Two sources, because neither is complete on its own:

    - the **contents page** (three or more labelled lines on one page) gives
      clean titles, and for a règlement the printed page number
      ("Art. 1 But et modalités de l'appel d'offres 6");
    - the **body** gives where each heading really starts. In a règlement the
      OCR merges the margin column with the text ("Art. 1 1.1 L'enjeu…"), so
      the body knows the page but not the title; in a Eurocode the body line
      is "Section 2 _ Bases du calcul geotechnique", usable but often without
      its accents.

    A number seen in both takes its title from the contents and its page from
    the body. A number seen only in the contents takes the printed page plus
    the offset the pairs establish; only in the body, the body's title.
    """
    contents: dict[str, tuple[str, int | None]] = {}
    body: dict[str, tuple[str, int]] = {}
    for pageno in range(doc.page_count):
        if pageno in skip:
            continue
        hits: list[tuple[str, str, int | None]] = []
        for line in (l.strip() for l in page_lines(doc, pageno)):
            m = LABEL_RE.match(line)
            if not m:
                continue
            title = clean(m.group(3))
            printed = int(m.group(4)) if m.group(4) else None
            hits.append((m.group(2), title, printed))
        if len({h[0] for h in hits}) >= 3:
            for num, title, printed in hits:
                if title and plausible_title(title):
                    contents.setdefault(num, (title, printed))
        else:
            for num, title, printed in hits:
                body.setdefault(num, (title, pageno + 1))

    offsets = sorted(
        body[num][1] - printed
        for num, (_t, printed) in contents.items()
        if printed is not None and num in body and body[num][1] > printed
    )
    offset = offsets[len(offsets) // 2] if offsets else None

    found: dict[str, dict] = {}
    for num in sorted(set(contents) | set(body), key=int):
        c_title, printed = contents.get(num, ("", None))
        b_title, b_page = body.get(num, ("", None))
        title = c_title or b_title
        page = b_page
        if page is None and printed is not None and offset is not None:
            page = printed + offset
        if not title or page is None or not (1 <= page <= doc.page_count):
            continue
        found[num] = {"title": title, "pageStart": page}
    return found


SCOPE_TITLE_RE = re.compile(
    r"^(DOMAINE D.APPLICATION|GELTUNGSBEREICH|CAMPO D.APPLICAZIONE|SCOPE)\b", re.I
)


def _page_top_titles(doc: "fitz.Document", pages: range, skip: set[int]) -> list[tuple[int, str]]:
    """Uppercase, plausible titles among the first lines of each page: where a
    chapter of a SIA norm starts. Deeper in a page a capitalised line is a
    table header or a note."""
    out: list[tuple[int, str]] = []
    for pageno in pages:
        if pageno in skip:
            continue
        for line in [l.strip() for l in page_lines(doc, pageno)][:3]:
            m = UNNUMBERED_UPPER_RE.match(line)
            if m and plausible_title(m.group(1)) and not re.match(r"^(TABLEAU|TABELLE|TABELLA|TABLE|FIGURE|FIG\.|BILD|ANNEXE|ANHANG|ANNEX)\b", line):
                out.append((pageno + 1, clean(m.group(1))))
                break
    return out


def recover_missing_chapters(doc: "fitz.Document", chapters: dict[str, dict], skip: set[int]) -> dict[str, dict]:
    """Fill holes in the chapter sequence by position.

    On the OCR'd copies the numbering column is sometimes lost on a whole run
    of pages: the chapter title survives, and so does the first sub-heading,
    but neither carries its number, so `infer_unnumbered_chapters` has nothing
    to read. When the sequence goes 1, 2, 6 and exactly three title-only pages
    lie between chapter 2 and chapter 6, those are chapters 3, 4 and 5, in
    page order. Anything less certain is left as a hole for the agent.

    A leading "Domaine d'application" before chapter 1 is chapter 0, which is
    how SIA numbers it.
    """
    numbered = sorted((int(k), k) for k in chapters if k.isdigit())
    if not numbered:
        return chapters
    used = {meta["title"].upper() for meta in chapters.values()}
    recovered: dict[str, dict] = {}

    for (a, ka), (b, kb) in zip(numbered, numbered[1:]):
        holes = list(range(a + 1, b))
        if not holes:
            continue
        start, end = chapters[ka]["pageStart"], chapters[kb]["pageStart"]
        candidates = [
            (page, title)
            for page, title in _page_top_titles(doc, range(start, end - 1), skip)
            if title.upper() not in used
        ]
        if len(candidates) != len(holes):
            continue
        for n, (page, title) in zip(holes, candidates):
            recovered[str(n)] = {"title": title, "pageStart": page, "recovered": True}

    first_n, first_k = numbered[0]
    if first_n == 1 and "0" not in chapters:
        before = _page_top_titles(doc, range(0, chapters[first_k]["pageStart"] - 1), skip)
        scope = [(p, t) for p, t in before if SCOPE_TITLE_RE.match(t)]
        if len(scope) == 1:
            recovered["0"] = {"title": scope[0][1], "pageStart": scope[0][0], "recovered": True}

    return {**chapters, **recovered}


UNNUMBERED_UPPER_RE = re.compile(r"^([A-ZÉÈÀÂÔÎÛÇ][^a-z\n]{5,120})$")
FIRST_SUB_RE = re.compile(r"^(\d{1,2})\.\d{1,2}\b")


def infer_unnumbered_chapters(doc: "fitz.Document", skip: set[int]) -> dict[str, dict]:
    """Chapters whose number the OCR lost.

    On a scanned copy the large chapter number is often the one glyph the OCR
    does not recognise, so the page reads "FONDATIONS SUR PIEUX" followed by
    "9.1 Délimitation". The number is then taken from the first sub-heading
    under the title. Only fills holes: a chapter found with its number wins.
    """
    found: dict[str, dict] = {}
    for pageno in range(doc.page_count):
        if pageno in skip:
            continue
        lines = [l.strip() for l in page_lines(doc, pageno)]
        for i, line in enumerate(lines):
            m = UNNUMBERED_UPPER_RE.match(line)
            if not m or not plausible_title(m.group(1)):
                continue
            for follow in lines[i + 1 : i + 9]:
                sub = FIRST_SUB_RE.match(follow)
                if sub:
                    num = sub.group(1)
                    if num not in found:
                        found[num] = {"title": clean(m.group(1)), "pageStart": pageno + 1}
                    break
    return found


def chapter_page_spans(chapters: dict[str, dict], n_pages: int) -> dict[str, tuple[int, int]]:
    """First and last page of each numbered chapter, from where the next starts."""
    ordered = sorted(
        ((k, v["pageStart"]) for k, v in chapters.items()),
        key=lambda kv: kv[1],
    )
    spans: dict[str, tuple[int, int]] = {}
    for idx, (key, start) in enumerate(ordered):
        end = ordered[idx + 1][1] - 1 if idx + 1 < len(ordered) else n_pages
        if key.isdigit():
            spans[key] = (start, max(start, end))
    return spans


def extract_subsections(
    doc: "fitz.Document",
    chap_prefixes: set[str],
    toc_pages: set[int],
    running: set[str],
    chapter_spans: dict[str, tuple[int, int]] | None = None,
) -> dict[str, tuple[str, int]]:
    """Numeric sub-sections (depths 2-4) detected in body pages."""
    out: dict[str, tuple[str, int]] = {}
    skip = set(toc_pages) | dense_heading_pages(doc, toc_pages)
    for i in range(doc.page_count):
        if i in skip:
            continue
        text = heading_text(doc, i)
        for m in SUB_RE.finditer(text):
            num = m.group(1)
            title = clean(m.group(2))
            if not title or title in running:
                continue
            top = num.split(".", 1)[0]
            if top not in chap_prefixes:
                continue
            # "9.522" is "9.5.2.2" with a dot the OCR dropped. No norm has a
            # fortieth sub-section; keeping it would hang a page of 9.5 under
            # a sibling of 9.7 and knock the real 9.6 and 9.7 out of sequence.
            if any(int(part) > 40 for part in num.split(".")[1:]):
                continue
            # A sub-section sits inside its chapter. "1.1" found sixty pages
            # after chapter 1 ended is a numbered line in an annexe, and keeping
            # it hangs an unrelated page under the wrong parent.
            span = chapter_spans.get(top) if chapter_spans else None
            if span and not span[0] <= i + 1 <= span[1]:
                continue
            if not plausible_title(title):
                continue
            if num in out:
                continue
            out[num] = (title, i + 1)
    return out


def build_tree(chapters: dict, subsections: dict, n_pages: int) -> list[dict]:
    """Returns flat list of nodes with depth 0-indexed (chapter=0, section=1, ...)."""
    nodes: list[dict] = []
    for path, info in chapters.items():
        nodes.append(
            {"depth": 0, "path": path, "title": info["title"], "pageStart": info["pageStart"]}
        )
    for path, (title, pg) in subsections.items():
        nodes.append(
            {"depth": path.count("."), "path": path, "title": title, "pageStart": pg}
        )

    def sort_key(n):
        p = n["path"]
        if p.startswith("Annexe"):
            return (10**9, ord(p[-1]), [])
        parts = [int(x) for x in p.split(".")]
        return (parts[0], 0, parts)

    nodes.sort(key=sort_key)

    # pageEnd via next sibling/ancestor
    for idx, node in enumerate(nodes):
        d, pg = node["depth"], node["pageStart"]
        next_pg = n_pages
        for j in range(idx + 1, len(nodes)):
            if nodes[j]["depth"] <= d:
                next_pg = max(nodes[j]["pageStart"] - 1, pg)
                break
        node["pageEnd"] = next_pg

    # Parent links + nodeId + orderIndex
    parent_stack: list[tuple[int, str]] = []  # (depth, nodeId)
    for idx, node in enumerate(nodes):
        d = node["depth"]
        while parent_stack and parent_stack[-1][0] >= d:
            parent_stack.pop()
        node["nodeId"] = f"s-{slugify(node['path'])}"
        node["parentNodeId"] = parent_stack[-1][1] if parent_stack else None
        node["orderIndex"] = idx
        parent_stack.append((d, node["nodeId"]))

    return nodes


def slugify(path: str) -> str:
    return re.sub(r"[^a-zA-Z0-9]+", "-", path).strip("-").lower() or "root"


# --- section text ------------------------------------------------------------
#
# Until 2026-09-04 a node's text was the concatenation of the pages it spanned,
# which meant three things the audit measured on the production corpus
# (ADR 22 § 2.1): a chapter carried the full text of all its sub-sections
# (61 000 characters twice for SIA 267.153 § 11 and § 11.3), the running header
# and the page number of every page sat in the middle of the prose, and words
# hyphenated at a line end stayed broken ("ter- rain"). The text of a node is
# now what lies between its heading and the next heading, on cleaned pages.

PAGE_NUMBER_RE = re.compile(r"^\s*\d{1,4}\s*$")
CLAUSE_NUMBER_RE = re.compile(r"^\s*(\d{1,2}(?:\.\d{1,3}){1,4})\s*$")
HYPHEN_BREAK_RE = re.compile(r"([a-zà-ÿ])[-—]\n([a-zà-ÿ])")


def dehyphenate(text: str) -> str:
    """Join a word broken at a line end: `tra-\\nvaux` becomes `travaux`.

    A compound word split at its hyphen loses it too (`pieux-\\nradier` reads
    `pieuxradier`); syllable breaks outnumber compounds at line ends by far,
    and a broken word is worse for search than a missing hyphen.
    """
    return HYPHEN_BREAK_RE.sub(r"\1\2", text)


def clean_page(text: str, running: set[str], page_count: int) -> str:
    """Drop what is page furniture rather than norm text.

    - the running header and footer lines `detect_running_text` found (they
      were detected before, and never removed);
    - a bare page number in the three first or three last lines of the page;
    - trailing whitespace on every line.
    """
    lines = [l.rstrip() for l in text.splitlines()]
    kept: list[str] = []
    n = len(lines)
    for i, line in enumerate(lines):
        s = line.strip()
        if s and (s in running or furniture_key(s) in running):
            continue
        # Printed page numbers run a little past the PDF's page count when
        # the front matter is numbered separately; well past it, a bare
        # number in the margin is a table value.
        near_edge = i < 3 or i >= n - 3
        if near_edge and PAGE_NUMBER_RE.match(s) and int(s) <= page_count + 20:
            continue
        kept.append(line)
    return "\n".join(kept)


def join_clause_numbers(text: str) -> str:
    """Rewrite `9.5.1.1\\nLes pieux…` as `9.5.1.1 Les pieux…`.

    PyMuPDF reads the numbering column before the paragraph, so every numbered
    clause of a SIA norm arrives as a number on its own line. Joined, the
    number becomes an anchor the reader can search for; alone, it is noise.
    Content only: heading detection keeps its stricter rule
    (`join_split_headings`), because a clause number followed by a sentence is
    exactly what must NOT become a heading.
    """
    lines = text.splitlines()
    out: list[str] = []
    i = 0
    while i < len(lines):
        m = CLAUSE_NUMBER_RE.match(lines[i])
        if m and i + 1 < len(lines) and lines[i + 1].strip():
            out.append(f"{m.group(1)} {lines[i + 1].strip()}")
            i += 2
            continue
        out.append(lines[i])
        i += 1
    return "\n".join(out)


def tidy(text: str) -> str:
    """Collapse runs of blank lines; strip the ends."""
    return re.sub(r"\n{3,}", "\n\n", text).strip()


def locate_heading(lines: list[str], node: dict) -> int:
    """Index of the line carrying this node's heading on its start page, or -1.

    Matches on the number first (`9.5 ` at the start of a line), then on the
    first words of the title, so a title PyMuPDF wrapped differently from the
    contents page is still found.
    """
    path = node["path"]
    title = node["title"].split(" — ")[-1].strip()
    stem = re.sub(r"\s+", " ", title[:24]).lower()
    if path.startswith("Annexe"):
        needle = path.lower()
        for i, line in enumerate(lines):
            if line.strip().lower().startswith(needle):
                return i
        return -1
    if path.isdigit():
        labelled = re.compile(
            r"^(?:Art\.?|Section|Chapitre|Kapitel|Abschnitt)\s*" + path + r"\b"
        )
        for i, line in enumerate(lines):
            if labelled.match(line.strip()):
                return i
    prefix = path + " "
    for i, line in enumerate(lines):
        s = line.strip()
        if not s.startswith(prefix):
            continue
        rest = re.sub(r"\s+", " ", s[len(prefix):]).lower()
        if not stem or rest.startswith(stem[: min(len(stem), 12)]):
            return i
    for i, line in enumerate(lines):
        if stem and len(stem) >= 8 and re.sub(r"\s+", " ", line.strip()).lower().startswith(stem):
            return i
    return -1


def extract_section_text(
    doc: "fitz.Document", nodes: list[dict], running: set[str]
) -> list[str]:
    """Give each node the text between its heading and the next one.

    Pages are cleaned first (`clean_page`), then read as one stream. Every
    node's heading is located on its start page; the node's text runs from the
    line after its heading to the next located heading, in document order,
    whatever the depth. A heading that cannot be located falls back to the
    start of its page, which is the previous behaviour for that node only, and
    is reported in the warnings so the skill can look at it.

    Returns the warnings.
    """
    warnings: list[str] = []
    pages = [
        clean_page(join_split_headings(page_text(doc, i)), running, doc.page_count)
        for i in range(doc.page_count)
    ]
    lines_per_page = [p.split("\n") for p in pages]

    # Absolute line index of the first line of each page in the stream.
    page_offsets: list[int] = []
    total = 0
    for lines in lines_per_page:
        page_offsets.append(total)
        total += len(lines)
    stream = [line for lines in lines_per_page for line in lines]

    located: list[tuple[int, int, dict]] = []  # (start_line, body_line, node)
    for node in nodes:
        pageno = node["pageStart"] - 1
        at = locate_heading(lines_per_page[pageno], node)
        if at == -1:
            warnings.append(f"heading not located: {node['path']} {node['title'][:40]!r} (p. {node['pageStart']})")
            start = page_offsets[pageno]
            located.append((start, start, node))
        else:
            start = page_offsets[pageno] + at
            located.append((start, start + 1, node))

    # Document order, not path order: annexes sort last by path but sit last
    # in the document anyway; a heading not located sorts at its page start.
    located.sort(key=lambda t: (t[0], -t[2]["depth"]))
    for idx, (start, body, node) in enumerate(located):
        end = located[idx + 1][0] if idx + 1 < len(located) else len(stream)
        text = "\n".join(stream[body:end])
        node["rawText"] = tidy(dehyphenate(join_clause_numbers(text)))
        node["textSource"] = "heading" if body > start else "page-start"

    # A chapter whose text lives in its sections is normal; a leaf with
    # nothing under its heading is a heading the text extraction lost.
    parents = {n["parentNodeId"] for n in nodes if n.get("parentNodeId")}
    empty = [n["path"] for n in nodes if len(n["rawText"]) < 40 and n["nodeId"] not in parents]
    if empty:
        warnings.append(f"{len(empty)} leaf node(s) with no text: {', '.join(empty[:8])}")
    return warnings


def missing_chapter_numbers(chapters: dict[str, dict]) -> list[str]:
    """Holes in the chapter sequence, for the skill to look at the contents page."""
    numbers = sorted(int(k) for k in chapters if k.isdigit())
    if len(numbers) < 2:
        return []
    return [str(n) for n in range(numbers[0], numbers[-1]) if n not in set(numbers)]


def prune_subsections(
    subsections: dict[str, tuple[str, int]], chapters: dict[str, dict]
) -> tuple[dict[str, tuple[str, int]], list[str]]:
    """Keep, under each parent, the children whose numbers advance with the pages.

    The same idea as `prune_chapters`, one level down: a titled "4.2" found on
    page 63 inside chapter 9 is a table row, and keeping it would hang a page
    of chapter 9 under chapter 4. Children are checked against their parent's
    first page and against each other's order.
    """
    dropped: list[str] = []
    by_parent: dict[str, list[str]] = defaultdict(list)
    for path in subsections:
        parent = path.rsplit(".", 1)[0]
        by_parent[parent].append(path)

    kept: dict[str, tuple[str, int]] = {}
    for parent, children in by_parent.items():
        parent_page = (
            chapters.get(parent, {}).get("pageStart")
            if parent in chapters
            else (subsections[parent][1] if parent in subsections else None)
        ) or 0
        ordered = [
            p
            for p in sorted(children, key=lambda p: [int(x) for x in p.split(".")])
            if subsections[p][1] >= parent_page
        ]
        for path in children:
            if path not in ordered:
                title, page = subsections[path]
                dropped.append(f"{path} {title[:30]!r} (p. {page}, before its parent)")

        # Longest run whose pages never go backwards, like `prune_chapters`:
        # a greedy walk would keep the one stray row and drop every real
        # child after it.
        pages = [subsections[p][1] for p in ordered]
        best = [1] * len(ordered)
        prev = [-1] * len(ordered)
        for i in range(len(ordered)):
            for j in range(i):
                if pages[j] <= pages[i] and best[j] + 1 > best[i]:
                    best[i], prev[i] = best[j] + 1, j
        chain: set[int] = set()
        if ordered:
            # On a tie, the run that ends earliest in the document: a stray
            # table row sits far past the real children.
            top = max(best)
            idx = min((i for i in range(len(ordered)) if best[i] == top), key=lambda i: pages[i])
            while idx != -1:
                chain.add(idx)
                idx = prev[idx]
        for i, path in enumerate(ordered):
            title, page = subsections[path]
            if i in chain:
                kept[path] = (title, page)
            else:
                dropped.append(f"{path} {title[:30]!r} (p. {page}, out of sequence)")
    return kept, dropped


def extract_figures(
    doc: "fitz.Document",
    toc_pages: set[int],
    out_dir: Path,
    dpi: int = 150,
) -> list[dict]:
    """One figure manifest entry per caption detected. Page is rendered to PNG."""
    figs_dir = out_dir / "figures"
    figs_dir.mkdir(parents=True, exist_ok=True)
    seen: dict[str, dict] = {}
    matrix = fitz.Matrix(dpi / 72, dpi / 72)
    for pageno in range(doc.page_count):
        if pageno in toc_pages:
            continue
        text = page_text(doc, pageno)
        rendered = False
        page_png: Path | None = None
        for m in CAPTION_RE.finditer(text):
            num = m.group(1)
            label_rest = m.group(2).strip()
            if num in seen:
                continue
            caption = f"Figure {num}" + (f" — {label_rest}" if label_rest else "")
            # render the page once if not done
            if not rendered:
                pix = doc[pageno].get_pixmap(matrix=matrix, alpha=False)
                page_png = figs_dir / f"page-{pageno + 1:03d}.png"
                pix.save(str(page_png))
                rendered = True
            seen[num] = {
                "figureNumber": num,
                "caption": caption[:240],
                "page": pageno + 1,
                "fileName": f"figures/{page_png.name}" if page_png else None,
                "renderDpi": dpi,
                "mimeType": "image/png",
            }
    return list(seen.values())


def guess_language(doc: "fitz.Document") -> str:
    """Best-effort language guess based on common French/German/Italian markers."""
    sample = " ".join(page_text(doc, i) for i in range(min(5, doc.page_count))).lower()
    scores = {
        "fr": sum(sample.count(w) for w in (" la ", " les ", " sont ", " avec ", " selon ")),
        "de": sum(sample.count(w) for w in (" der ", " die ", " sind ", " mit ", " nach ")),
        "it": sum(sample.count(w) for w in (" la ", " sono ", " con ", " per ", " della ")),
        "en": sum(sample.count(w) for w in (" the ", " are ", " with ", " of ", " shall ")),
    }
    return max(scores, key=scores.get)


def main() -> None:
    ap = argparse.ArgumentParser(description="Stratta ingest pre-pass (PyMuPDF).")
    ap.add_argument("--pdf", required=True, help="Path to the norm PDF.")
    ap.add_argument("--output", required=True, help="Output directory.")
    ap.add_argument(
        "--language",
        choices=["fr", "de", "it", "en"],
        help="Override detected language.",
    )
    ap.add_argument("--dpi", type=int, default=150, help="Figure render DPI.")
    args = ap.parse_args()

    pdf_path = Path(args.pdf).resolve()
    out_dir = Path(args.output).resolve()
    if not pdf_path.exists():
        sys.stderr.write(f"PDF not found: {pdf_path}\n")
        sys.exit(1)
    out_dir.mkdir(parents=True, exist_ok=True)

    doc = fitz.open(str(pdf_path))
    toc_pages = detect_toc_pages(doc)
    running = detect_running_text(doc)
    chapters = extract_chapters(doc, toc_pages)
    chap_prefixes = {k for k in chapters if k.isdigit()}
    spans = chapter_page_spans(chapters, doc.page_count)
    subsections = extract_subsections(doc, chap_prefixes, toc_pages, running, spans)
    subsections, pruned = prune_subsections(subsections, chapters)
    nodes = build_tree(chapters, subsections, doc.page_count)
    warnings = [f"sub-section dropped, out of sequence: {p}" for p in pruned]
    holes = missing_chapter_numbers(chapters)
    if holes:
        warnings.append(
            f"chapters missing from the sequence: {', '.join(holes)} (check the contents page; an OCR'd copy often loses the chapter number)"
        )
    recovered = sorted(
        (int(k) for k, meta in chapters.items() if k.isdigit() and meta.get("recovered"))
    )
    if recovered:
        warnings.append(
            f"chapter numbers recovered by position: {', '.join(map(str, recovered))} (their sub-headings are unnumbered in the text; check the contents page)"
        )
    warnings += extract_section_text(doc, nodes, running)
    figures = extract_figures(doc, toc_pages, out_dir, dpi=args.dpi)

    lang = args.language or guess_language(doc)

    by_depth: dict[int, int] = defaultdict(int)
    for n in nodes:
        by_depth[n["depth"]] += 1
    chars = sum(len(n["rawText"]) for n in nodes)
    parents = {n["parentNodeId"] for n in nodes if n.get("parentNodeId")}
    empty_nodes = sum(
        1 for n in nodes if len(n["rawText"]) < 40 and n["nodeId"] not in parents
    )
    if not chapters:
        warnings.append("no chapter detected: neither bookmarks nor headings")

    manifest = {
        "doc": {
            "sourcePath": str(pdf_path),
            "pageCount": doc.page_count,
            "language": lang,
            "tocSource": "bookmarks+regex" if doc.get_toc() else "regex-only",
            "metadata": {k: v for k, v in doc.metadata.items() if v},
            "tocPagesDetected": sorted(toc_pages),
            "runningHeadersDetected": sorted(running)[:10],
        },
        "stats": {
            "sectionCount": len(nodes),
            "byDepth": dict(by_depth),
            "figureCount": len(figures),
            # What the server's coverage score will see (`corpusQuality.ts`),
            # so the skill can judge the pre-pass before writing anything.
            "chars": chars,
            "charsPerPage": round(chars / max(1, doc.page_count)),
            "emptyLeaves": empty_nodes,
        },
        # Anything the pre-pass could not decide on its own. The skill reads
        # this list and looks at the pages it names; an empty list is the
        # normal case on a native SIA norm.
        "warnings": warnings,
        "sections": nodes,
        "figures": figures,
    }

    manifest_path = out_dir / "prepass.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")

    sys.stdout.write(
        json.dumps(
            {
                "ok": True,
                "manifest": str(manifest_path),
                "sectionCount": len(nodes),
                "byDepth": dict(by_depth),
                "figureCount": len(figures),
                "charsPerPage": round(chars / max(1, doc.page_count)),
                "warnings": len(warnings),
                "language": lang,
                "pageCount": doc.page_count,
            },
            ensure_ascii=False,
        )
        + "\n"
    )


if __name__ == "__main__":
    main()