#!/usr/bin/env python3
"""Extract figures from a (mostly vector) datasheet PDF into per-PDF folders.

Most figures in these Renesas datasheets are vector line art, so `pdfimages`
finds nothing. Instead we:

  1. Read the structured text layer (`mutool draw -F stext`) to locate the
     "Figure N. <title>" captions and the surrounding body text.
  2. Derive a coarse band for each figure: bounded below by its caption and
     above by whichever comes last -- the previous row of captions on the page,
     the last body-text/heading line, or the top of the content area.
  3. Where a page places figures side by side (performance-graph pages usually
     do), split the band horizontally at the midpoints between the captions of
     that row, so each figure gets its own column.
  4. Render the page and tighten that band to the actual ink bounding box, so
     the crop hugs the drawing instead of the column.
  4. Write `figure-NN.png` named by the datasheet's own global figure number.

Usage:
    extract-figures.py [--dry-run] [--dpi N] [--outdir DIR] file.pdf [file.pdf ...]
"""

from __future__ import annotations

import argparse
import re
import shutil
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path

# --- page geometry, in PDF points (origin top-left, y increasing downward) ---
HEADER_BOTTOM = 62.0      # running head + its horizontal rule (ends ~60pt)
FOOTER_TOP = 740.0        # doc id / revision / page number sit below this
BODY_MIN_SIZE = 9.5       # body text and section headings are >= 10pt
BODY_MAX_X0 = 60.0        # ...and start at the left margin (~47.5pt)
CAPTION_MAX_SIZE = 13.0   # captions are bold, 9-12pt depending on template
CAPTION_LINE_GAP = 5.0    # max vertical gap between wrapped caption lines
ROW_TOLERANCE = 6.0       # captions starting within this many points of each
                          # other sit side by side in the same row
BODY_GAP = 5.0            # clearance below body text; stext line bboxes clip
                          # descenders, which would otherwise leak into a crop
PAD = 6.0                 # padding around the ink bbox, in points

CAPTION_RE = re.compile(r"^\s*Figure\s+(\d+(?:\.\d+)?)\s*\.\s*(.*)$")

INK_THRESHOLD = 200       # 0-255 grey; anything darker counts as ink
ANALYSIS_DPI = 100        # cheap render used only for ink detection


@dataclass
class Line:
    x0: float
    y0: float
    x1: float
    y1: float
    size: float
    text: str
    bold: bool = False


@dataclass
class Caption:
    number: str
    title: str
    top: float        # y0 of the first caption line
    bottom: float     # y1 of the last wrapped caption line
    x0: float         # left edge across all wrapped lines
    x1: float         # right edge across all wrapped lines


@dataclass
class Figure:
    number: str
    title: str
    page: int
    top: float
    bottom: float
    left: float
    right: float


def run(cmd: list[str]) -> str:
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(f"{cmd[0]} failed: {proc.stderr.strip()}")
    return proc.stdout


def page_lines(pdf: Path, page: int) -> tuple[float, float, list[Line]]:
    xml = run(["mutool", "draw", "-F", "stext", "-o", "-", str(pdf), str(page)])
    root = ET.fromstring(xml)
    pg = root.find("page")
    if pg is None:
        return 0.0, 0.0, []
    lines: list[Line] = []
    for el in pg.iter("line"):
        bbox = el.get("bbox")
        if not bbox:
            continue
        x0, y0, x1, y1 = (float(v) for v in bbox.split())
        font = el.find("font")
        size = float(font.get("size")) if font is not None else 0.0
        bold = bool(font is not None and "bold" in (font.get("name") or "").lower())
        lines.append(Line(x0, y0, x1, y1, size, el.get("text") or "", bold))
    lines.sort(key=lambda l: (l.y0, l.x0))
    return float(pg.get("width")), float(pg.get("height")), lines


def find_captions(lines: list[Line]) -> list[Caption]:
    """Locate the "Figure N. <title>" captions on a page.

    Captions wrap onto several lines; the bottom of the *last* line is what
    bounds the next figure, otherwise the wrapped remainder of this caption
    leaks into the top of the following crop. The horizontal extent spans all
    the wrapped lines, and is what separates side-by-side figures.
    """
    out: list[Caption] = []
    for i, line in enumerate(lines):
        m = CAPTION_RE.match(line.text)
        if not m or line.size > CAPTION_MAX_SIZE or not line.bold:
            continue
        title = m.group(2).strip()
        # Absorb continuation lines: same size, directly below, horizontally
        # overlapping (a side-by-side caption is not a continuation), and not
        # itself a new caption.
        prev = line
        x0, x1 = line.x0, line.x1
        for nxt in lines[i + 1:]:
            if abs(nxt.size - line.size) > 0.5:
                continue
            # A caption in the neighbouring column is not a continuation, and
            # must not stop the search either.
            if nxt.x1 < x0 or nxt.x0 > x1:
                continue
            if nxt.y0 - prev.y1 > CAPTION_LINE_GAP or nxt.y0 < prev.y1 - 1:
                break
            if CAPTION_RE.match(nxt.text):
                break
            title = f"{title} {nxt.text.strip()}".strip()
            x0, x1 = min(x0, nxt.x0), max(x1, nxt.x1)
            prev = nxt
        out.append(Caption(m.group(1), title, line.y0, prev.y1, x0, x1))
    out.sort(key=lambda c: (c.top, c.x0))
    return out


def caption_rows(captions: list[Caption]) -> list[list[Caption]]:
    """Group captions that sit side by side into rows, ordered top to bottom."""
    rows: list[list[Caption]] = []
    for cap in captions:
        if rows and abs(cap.top - rows[-1][0].top) <= ROW_TOLERANCE:
            rows[-1].append(cap)
        else:
            rows.append([cap])
    for row in rows:
        row.sort(key=lambda c: c.x0)
    return rows


def figures_on_page(pdf: Path, page: int) -> list[Figure]:
    width, _, lines = page_lines(pdf, page)
    captions = find_captions(lines)
    if not captions:
        return []

    rows = caption_rows(captions)
    figures: list[Figure] = []
    for r, row in enumerate(rows):
        row_top = min(c.top for c in row)
        upper = HEADER_BOTTOM
        # The previous row of captions bounds this one from above -- use the
        # last wrapped line of that row, not its first.
        if r > 0:
            upper = max(upper, max(c.bottom for c in rows[r - 1]) + BODY_GAP)
        # So does the last body-text / heading line above the row.
        for line in lines:
            if line.y1 >= row_top or line.y1 <= HEADER_BOTTOM:
                continue
            if line.y0 >= FOOTER_TOP:
                continue
            is_body = line.size >= BODY_MIN_SIZE and line.x0 <= BODY_MAX_X0
            if is_body and line.y1 + BODY_GAP > upper:
                # Do not step over an earlier row's figures.
                if any(upper < c.top < line.y1 for pr in rows[:r] for c in pr):
                    continue
                upper = line.y1 + BODY_GAP
        # Split the row horizontally at the midpoints between its captions.
        for i, cap in enumerate(row):
            left = 0.0 if i == 0 else (row[i - 1].x1 + cap.x0) / 2.0
            right = width if i == len(row) - 1 else (cap.x1 + row[i + 1].x0) / 2.0
            figures.append(
                Figure(cap.number, cap.title, page, upper, cap.top, left, right)
            )
    return figures


def figures_on_page_above(pdf: Path, page: int) -> list[Figure]:
    """Same as figures_on_page, but for templates where the caption sits
    ABOVE the figure (e.g. USB PD spec) instead of below it (Renesas-style).
    """
    width, _, lines = page_lines(pdf, page)
    captions = find_captions(lines)
    if not captions:
        return []

    rows = caption_rows(captions)
    figures: list[Figure] = []
    for r, row in enumerate(rows):
        row_bottom = max(c.bottom for c in row)
        upper = row_bottom + BODY_GAP
        lower = FOOTER_TOP
        # The next row of captions bounds this one from below -- use its
        # topmost caption line.
        if r + 1 < len(rows):
            lower = min(lower, min(c.top for c in rows[r + 1]) - BODY_GAP)
        # So does the first body-text / heading line below the row.
        for line in lines:
            if line.y0 <= row_bottom or line.y0 >= lower:
                continue
            if line.y1 <= HEADER_BOTTOM:
                continue
            if line.bold and CAPTION_RE.match(line.text):
                continue  # a real next caption; already bounded via rows[r+1]
            is_body = line.size >= BODY_MIN_SIZE and line.x0 <= BODY_MAX_X0
            if is_body and line.y0 - BODY_GAP < lower:
                lower = line.y0 - BODY_GAP
        # Split the row horizontally at the midpoints between its captions.
        for i, cap in enumerate(row):
            left = 0.0 if i == 0 else (row[i - 1].x1 + cap.x0) / 2.0
            right = width if i == len(row) - 1 else (cap.x1 + row[i + 1].x0) / 2.0
            figures.append(
                Figure(cap.number, cap.title, page, upper, lower, left, right)
            )
    return figures


def read_pgm(path: Path) -> tuple[int, int, bytes]:
    data = path.read_bytes()
    if not data.startswith(b"P5"):
        raise ValueError(f"{path} is not a binary PGM")
    fields, pos = [], 2
    while len(fields) < 3:
        while pos < len(data) and data[pos : pos + 1].isspace():
            pos += 1
        if data[pos : pos + 1] == b"#":
            while data[pos : pos + 1] not in (b"\n", b""):
                pos += 1
            continue
        start = pos
        while pos < len(data) and not data[pos : pos + 1].isspace():
            pos += 1
        fields.append(int(data[start:pos]))
    return fields[0], fields[1], data[pos + 1 :]


def ink_bbox(
    pgm: tuple[int, int, bytes],
    top_pt: float,
    bottom_pt: float,
    scale: float,
    left_pt: float = 0.0,
    right_pt: float | None = None,
) -> tuple[float, float, float, float] | None:
    """Tightest ink bbox (in points) within the given band."""
    w, h, px = pgm
    r0 = max(0, int(top_pt * scale))
    r1 = min(h, int(bottom_pt * scale) + 1)
    c0 = max(0, int(left_pt * scale))
    c1 = min(w, int(right_pt * scale) + 1) if right_pt is not None else w
    if r1 <= r0 or c1 <= c0:
        return None

    rows = [r for r in range(r0, r1)
            if min(px[r * w + c0 : r * w + c1]) < INK_THRESHOLD]
    if not rows:
        return None
    top_px, bottom_px = rows[0], rows[-1] + 1

    span = c1 - c0
    col_min = bytearray(b"\xff" * span)
    for r in rows:
        row = px[r * w + c0 : r * w + c1]
        for c in range(span):
            if row[c] < col_min[c]:
                col_min[c] = row[c]
    cols = [c for c in range(span) if col_min[c] < INK_THRESHOLD]
    if not cols:
        return None
    left_px, right_px = c0 + cols[0], c0 + cols[-1] + 1

    return (left_px / scale, top_px / scale, right_px / scale, bottom_px / scale)


def crop_region(pdf: Path, outdir: Path, dpi: int, spec: str, dry_run: bool) -> str:
    """Crop an explicitly given region: PAGE:x0,y0,x1,y1:NAME (points).

    For figures the datasheet never captioned, so caption detection cannot
    find them. The given box is still tightened to the ink inside it.
    """
    page_s, box_s, name = spec.split(":", 2)
    page = int(page_s)
    bx0, by0, bx1, by1 = (float(v) for v in box_s.split(","))
    tmp = Path(tempfile.mkdtemp(prefix="figregion-"))
    try:
        analysis = tmp / "a"
        run(["pdftoppm", "-gray", "-r", str(ANALYSIS_DPI), "-f", str(page),
             "-l", str(page), "-singlefile", str(pdf), str(analysis)])
        pgm = read_pgm(analysis.with_suffix(".pgm"))
        box = ink_bbox(pgm, by0, by1, ANALYSIS_DPI / 72.0, bx0, bx1)
        if box is None:
            return f"SKIP: no ink in region {box_s} on p{page}"
        x0, y0, x1, y1 = box
        x0, y0 = max(bx0, x0 - PAD), max(by0, y0 - PAD)
        x1, y1 = min(bx1, x1 + PAD), min(by1, y1 + PAD)
        if not dry_run:
            full = tmp / "f"
            run(["pdftoppm", "-png", "-r", str(dpi), "-f", str(page),
                 "-l", str(page), "-singlefile", str(pdf), str(full)])
            s = dpi / 72.0
            geom = (f"{round((x1 - x0) * s)}x{round((y1 - y0) * s)}"
                    f"+{round(x0 * s)}+{round(y0 * s)}")
            outdir.mkdir(parents=True, exist_ok=True)
            run(["magick", str(full.with_suffix(".png")), "-crop", geom,
                 "+repage", str(outdir / name)])
        return (f"p{page} region -> ink {x0:.0f},{y0:.0f} "
                f"{x1 - x0:.0f}x{y1 - y0:.0f}pt")
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def extract(pdf: Path, outdir: Path, dpi: int, dry_run: bool,
            caption_position: str = "below") -> list[tuple[Figure, str]]:
    finder = figures_on_page_above if caption_position == "above" else figures_on_page
    pages = int(
        re.search(r"Pages:\s+(\d+)", run(["pdfinfo", str(pdf)])).group(1)
    )
    results: list[tuple[Figure, str]] = []
    tmp = Path(tempfile.mkdtemp(prefix="figextract-"))
    try:
        for page in range(1, pages + 1):
            figs = finder(pdf, page)
            if not figs:
                continue

            analysis = tmp / f"a{page}"
            run(["pdftoppm", "-gray", "-r", str(ANALYSIS_DPI), "-f", str(page),
                 "-l", str(page), "-singlefile", str(pdf), str(analysis)])
            pgm = read_pgm(analysis.with_suffix(".pgm"))
            scale = ANALYSIS_DPI / 72.0

            full = tmp / f"f{page}"
            if not dry_run:
                run(["pdftoppm", "-png", "-r", str(dpi), "-f", str(page),
                     "-l", str(page), "-singlefile", str(pdf), str(full)])

            for fig in figs:
                box = ink_bbox(pgm, fig.top, fig.bottom, scale,
                               fig.left, fig.right)
                if box is None:
                    results.append((fig, "SKIP: no ink found in band"))
                    continue
                x0, y0, x1, y1 = box
                # Pad, but never past the band edges: the caption sits just
                # below `fig.bottom` and body text just above `fig.top`.
                x0, x1 = max(fig.left, x0 - PAD), min(fig.right, x1 + PAD)
                y0 = max(y0 - PAD, fig.top)
                y1 = min(y1 + PAD, fig.bottom - 2.0)
                name = f"figure-{fig.number}.png"
                note = (f"p{fig.page} band {fig.top:.0f}-{fig.bottom:.0f}pt "
                        f"x {fig.left:.0f}-{fig.right:.0f}pt "
                        f"-> ink {x0:.0f},{y0:.0f} {x1 - x0:.0f}x{y1 - y0:.0f}pt")
                if not dry_run:
                    s = dpi / 72.0
                    geom = (f"{round((x1 - x0) * s)}x{round((y1 - y0) * s)}"
                            f"+{round(x0 * s)}+{round(y0 * s)}")
                    outdir.mkdir(parents=True, exist_ok=True)
                    run(["magick", str(full.with_suffix(".png")), "-crop", geom,
                         "+repage", str(outdir / name)])
                results.append((fig, note))
    finally:
        shutil.rmtree(tmp, ignore_errors=True)
    return results


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pdfs", nargs="+", type=Path)
    ap.add_argument("--dpi", type=int, default=300, help="output render DPI")
    ap.add_argument("--outdir", type=Path, default=None,
                    help="output root (default: alongside each PDF)")
    ap.add_argument("--dry-run", action="store_true",
                    help="report detected figures without writing images")
    ap.add_argument("--region", action="append", default=[],
                    metavar="PAGE:x0,y0,x1,y1:NAME",
                    help="also crop an explicit region (points) from the single "
                         "given PDF; for figures with no 'Figure N.' caption")
    ap.add_argument("--caption-position", choices=["below", "above"], default="below",
                    help="whether 'Figure N. Title' sits below the figure "
                         "(Renesas-style, default) or above it (e.g. USB PD spec)")
    args = ap.parse_args()

    if args.region and len(args.pdfs) != 1:
        ap.error("--region applies to exactly one PDF")

    for pdf in args.pdfs:
        root = args.outdir if args.outdir else pdf.parent
        outdir = root / pdf.stem
        print(f"\n=== {pdf.name} -> {outdir}/")
        for fig, note in extract(pdf, outdir, args.dpi, args.dry_run, args.caption_position):
            print(f"  figure-{fig.number}  {note}")
            print(f"              {fig.title[:88]}")
        for spec in args.region:
            note = crop_region(pdf, outdir, args.dpi, spec, args.dry_run)
            print(f"  {spec.split(':', 2)[2]}  {note}")
    return 0


if __name__ == "__main__":
    sys.exit(main())