#!/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) ---
# The first five are template-specific; each has a command-line override of the
# same name (--header-bottom, --footer-top, ...) so a different vendor template
# can be tuned without editing the defaults.
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)
BODY_MIN_WIDTH = float("inf")  # ...or, for indented prose (bullets, numbered
                          # lists), are at least this wide. Off by default; set
                          # it between the widest *in-figure* label and the
                          # narrowest full line of prose. A right-edge test is
                          # not enough: a right-aligned figure label also ends
                          # at the right margin.
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

# Two printed caption conventions are supported:
#   "Figure 12. Title"  / "Figure 4.1. Title"   -- number, period, title
#   "Figure 3-8  Title" / "Figure A-14  Title"  -- chapter/appendix-hyphen-number,
#                                                  no period (USB Type-C spec)
# Group "num" is the figure number as printed, and is used verbatim in the
# output filename ("figure-4.1.png", "figure-3-8.png", "figure-A-14.png").
CAPTION_RE = re.compile(
    r"^\s*Figure\s+(?:"
    r"(?P<dotted>\d+(?:\.\d+)?)\s*\.\s*"        # Figure 12.  /  Figure 4.1.
    r"|"
    r"(?P<hyphen>(?:\d+|[A-Z])-\d+)\s*\.?(?:\s+|$)"   # Figure 3-8 / Figure A-14
    r")(?P<title>.*)$"
)


# A bold "Table N ..." caption is body furniture, not figure ink, and bounds a
# figure band exactly like a body-text line does. Same two number conventions.
TABLE_CAPTION_RE = re.compile(
    r"^\s*Table\s+(?:\d+(?:\.\d+)?\s*\.|(?:\d+|[A-Z])-\d+)(?:\s|$)"
)


def caption_number(m: re.Match) -> str:
    return m.group("dotted") or m.group("hyphen")

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
    name: str = ""    # output filename, assigned during extraction


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("title").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(caption_number(m), 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 is_body_line(line: Line) -> bool:
    """Body text / heading / table caption, i.e. page furniture that bounds a
    figure band -- as opposed to a label drawn inside the figure itself.

    Three shapes count:
      * a line at the left margin (plain paragraphs, section headings);
      * an indented but full-width line (bullets and numbered lists) -- only
        recognised once --body-min-width is set, because the cut-off has to sit
        above the widest label inside any figure. A right-edge test would not
        do: a right-aligned label inside a figure also ends at the margin;
      * a bold "Table N ..." caption, which is often centred and short.
    """
    if line.size < BODY_MIN_SIZE:
        return False
    if line.x0 <= BODY_MAX_X0 or line.x1 - line.x0 >= BODY_MIN_WIDTH:
        return True
    return bool(line.bold and TABLE_CAPTION_RE.match(line.text))


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 = is_body_line(line)
            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 = is_body_line(line)
            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]] = []
    # Multi-page drawings repeat the same caption verbatim (with ", cont."
    # appended) on every sheet. Keep them all: the first sheet keeps the plain
    # name, later ones get -cont1, -cont2, ... in page order.
    seen: dict[str, int] = {}
    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:
                n = seen.get(fig.number, 0)
                seen[fig.number] = n + 1
                fig.name = f"figure-{fig.number}{'' if n == 0 else f'-cont{n}'}.png"
                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 = fig.name
                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:
    global HEADER_BOTTOM, FOOTER_TOP, BODY_MAX_X0, BODY_MIN_WIDTH, BODY_GAP, PAD
    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)")
    # Page-layout tuning. Defaults match the Renesas template; override rather
    # than edit for a different vendor's page furniture.
    ap.add_argument("--header-bottom", type=float, default=HEADER_BOTTOM,
                    metavar="PT", help=f"y below the running head "
                                       f"(default {HEADER_BOTTOM:g})")
    ap.add_argument("--footer-top", type=float, default=FOOTER_TOP,
                    metavar="PT", help=f"y above the page footer "
                                       f"(default {FOOTER_TOP:g})")
    ap.add_argument("--body-max-x0", type=float, default=BODY_MAX_X0,
                    metavar="PT", help="max left edge of a body-text/heading "
                                       f"line (default {BODY_MAX_X0:g}; the USB "
                                       "Type-C template indents to 72)")
    ap.add_argument("--body-min-width", type=float, default=BODY_MIN_WIDTH,
                    metavar="PT", help="treat an indented line (bullet, numbered "
                                       "list) as body text when it is at least "
                                       "this wide (default: off)")
    ap.add_argument("--body-gap", type=float, default=BODY_GAP,
                    metavar="PT", help=f"clearance around body text "
                                       f"(default {BODY_GAP:g})")
    ap.add_argument("--pad", type=float, default=PAD,
                    metavar="PT", help=f"padding around the ink bbox "
                                       f"(default {PAD:g})")
    args = ap.parse_args()

    HEADER_BOTTOM = args.header_bottom
    FOOTER_TOP = args.footer_top
    BODY_MAX_X0 = args.body_max_x0
    BODY_MIN_WIDTH = args.body_min_width
    BODY_GAP = args.body_gap
    PAD = args.pad

    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"  {fig.name[:-4]}  {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())