split_book.py

#!/usr/bin/env python3
"""Split a book PDF into chapter-sized (or fixed-size) chunks for NotebookLM."""

import argparse
import re
import shutil
import sys
from pathlib import Path

from pypdf import PdfReader, PdfWriter
from pypdf.errors import PdfReadError


def load_reader(path):
    p = Path(path)
    if not p.exists():
        print(f"Error: input file not found: {path}", file=sys.stderr)
        sys.exit(1)

    try:
        reader = PdfReader(str(p))
    except PdfReadError as e:
        print(f"Error: could not read PDF '{path}': {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Error: could not open '{path}': {e}", file=sys.stderr)
        sys.exit(1)

    if reader.is_encrypted:
        try:
            result = reader.decrypt("")
        except Exception:
            result = 0
        if not result:
            stem = p.stem
            print(
                f"Error: '{path}' is encrypted and could not be opened with an "
                f"empty password. Try decrypting it first, e.g.:\n"
                f"  qpdf --decrypt {path} {stem}_decrypted.pdf",
                file=sys.stderr,
            )
            sys.exit(1)

    try:
        num_pages = len(reader.pages)
    except Exception as e:
        print(f"Error: could not read pages from '{path}': {e}", file=sys.stderr)
        sys.exit(1)

    if num_pages == 0:
        print(f"Error: '{path}' has zero pages.", file=sys.stderr)
        sys.exit(1)

    return reader, num_pages


def _walk_outline(items, depth, level, reader, out):
    for item in items:
        if isinstance(item, list):
            # A list immediately following an item is that item's children,
            # one level deeper.
            _walk_outline(item, depth + 1, level, reader, out)
            continue
        if depth == level:
            try:
                page_num = reader.get_destination_page_number(item)
            except Exception:
                continue
            title = item.title if item.title else "Untitled"
            out.append((page_num, title))


def get_chapters_from_outline(reader, level, num_pages):
    try:
        outline = reader.outline
    except Exception:
        outline = None

    if not outline:
        return None

    entries = []
    _walk_outline(outline, 1, level, reader, entries)

    if not entries:
        return None

    entries.sort(key=lambda x: x[0])

    deduped = []
    seen_pages = set()
    for page_num, title in entries:
        if page_num in seen_pages:
            continue
        seen_pages.add(page_num)
        deduped.append((page_num, title))

    chapters = []
    for i, (page_num, title) in enumerate(deduped):
        start = page_num
        end = deduped[i + 1][0] if i + 1 < len(deduped) else num_pages
        if end <= start:
            continue
        chapters.append({"title": title, "start": start, "end": end})

    return chapters if chapters else None


def group_chapters(chapters, target_pages):
    if not chapters:
        return []

    groups = []
    current = None
    current_titles = []

    for ch in chapters:
        ch_pages = ch["end"] - ch["start"]

        if current is None:
            current = {"start": ch["start"], "end": ch["end"]}
            current_titles = [ch["title"]]
        else:
            current_pages = current["end"] - current["start"]
            if current_pages >= target_pages / 2 and ch_pages >= target_pages:
                groups.append(_finalize_group(current, current_titles))
                current = {"start": ch["start"], "end": ch["end"]}
                current_titles = [ch["title"]]
            else:
                current["end"] = ch["end"]
                current_titles.append(ch["title"])

        current_pages = current["end"] - current["start"]
        if current_pages >= target_pages:
            groups.append(_finalize_group(current, current_titles))
            current = None
            current_titles = []

    if current is not None:
        groups.append(_finalize_group(current, current_titles))

    # Fold a too-small trailing group into the previous one rather than
    # leaving a tiny leftover file. Only applies to the final group.
    if len(groups) >= 2:
        last = groups[-1]
        last_pages = last["end"] - last["start"]
        if last_pages < target_pages / 2:
            prev = groups[-2]
            merged_title = f"{prev['title']}_to_{last['title']}"
            prev["end"] = last["end"]
            prev["title"] = merged_title
            groups.pop()

    return groups


def _finalize_group(current, titles):
    title = titles[0] if len(titles) == 1 else f"{titles[0]}_to_{titles[-1]}"
    return {"title": title, "start": current["start"], "end": current["end"]}


def get_chapters_fixed(num_pages, fixed_n):
    if fixed_n <= 0:
        raise ValueError(f"fixed_n must be a positive integer, got {fixed_n}")

    chapters = []
    start = 0
    while start < num_pages:
        end = min(start + fixed_n, num_pages)
        chapters.append({"title": f"pages_{start + 1}-{end}", "start": start, "end": end})
        start = end
    return chapters


def slugify(title, max_len=60):
    if title is None:
        title = ""
    cleaned = re.sub(r"[^\w\s-]", "", title, flags=re.UNICODE)
    cleaned = re.sub(r"\s+", "_", cleaned.strip())
    cleaned = cleaned.strip("_")
    cleaned = cleaned[:max_len].strip("_")
    return cleaned if cleaned else "untitled"


def write_chunks(reader, chunks, out_dir):
    out_path = Path(out_dir)
    if out_path.exists():
        # Clear any stale files from a prior run at this path (e.g. a
        # different --level/--target-pages produced a different chunk
        # count/plan) so old files never linger alongside the new ones.
        shutil.rmtree(out_path)
    out_path.mkdir(parents=True, exist_ok=True)

    width = len(str(len(chunks)))
    written = []
    for i, chunk in enumerate(chunks, start=1):
        writer = PdfWriter()
        for page_idx in range(chunk["start"], chunk["end"]):
            writer.add_page(reader.pages[page_idx])

        idx_str = str(i).zfill(width)
        filename = f"{idx_str}_{slugify(chunk['title'])}.pdf"
        file_path = out_path / filename

        with open(file_path, "wb") as f:
            writer.write(f)

        num_pages = chunk["end"] - chunk["start"]
        print(f"Wrote {file_path} ({num_pages} pages)")
        written.append(file_path)

    return written


def print_plan(chapters, num_pages):
    print(f"{'Idx':>4}  {'Pages':>15}  {'Count':>6}  Title")
    total_pages_covered = 0
    for i, ch in enumerate(chapters, start=1):
        page_count = ch["end"] - ch["start"]
        total_pages_covered += page_count
        page_range = f"{ch['start'] + 1}-{ch['end']}"
        print(f"{i:>4}  {page_range:>15}  {page_count:>6}  {ch['title']}")
    print(f"\nTotal chunks: {len(chapters)}, total pages covered: {total_pages_covered}")


def build_arg_parser():
    parser = argparse.ArgumentParser(
        description="Split a book PDF into smaller PDFs sized for NotebookLM sources."
    )
    parser.add_argument("input", help="Path to the source PDF.")
    parser.add_argument("--level", type=int, default=1, help="Outline depth to split at (default 1).")
    parser.add_argument(
        "--target-pages", type=int, default=None,
        help="Grouping mode: merge consecutive chapters until each group is at least N pages.",
    )
    parser.add_argument(
        "--fixed", type=int, default=None,
        help="Fallback mode: ignore the outline and split every N pages.",
    )
    parser.add_argument("--out", type=str, default="./chunks", help="Base output directory (default ./chunks).")
    parser.add_argument("--list", action="store_true", help="Print the plan and exit without writing files.")
    return parser


def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    if args.fixed is not None and args.fixed <= 0:
        print(f"Error: --fixed must be a positive integer, got {args.fixed}.", file=sys.stderr)
        sys.exit(1)
    if args.target_pages is not None and args.target_pages <= 0:
        print(f"Error: --target-pages must be a positive integer, got {args.target_pages}.", file=sys.stderr)
        sys.exit(1)

    input_path = Path(args.input)
    reader, num_pages = load_reader(args.input)
    print(f"Loaded {input_path.name}: {num_pages} pages")

    if args.fixed is not None:
        chapters = get_chapters_fixed(num_pages, args.fixed)
    else:
        chapters = get_chapters_from_outline(reader, args.level, num_pages)
        if chapters is None:
            if args.target_pages is not None:
                chapters = get_chapters_fixed(num_pages, args.target_pages)
            else:
                print(
                    f"Error: no usable outline found at --level {args.level}. "
                    f"Try a different --level (e.g. --level {args.level + 1}), "
                    f"or use --fixed N / --target-pages N to split without an outline.",
                    file=sys.stderr,
                )
                sys.exit(1)
        elif args.target_pages is not None:
            chapters = group_chapters(chapters, args.target_pages)

    print_plan(chapters, num_pages)

    if args.list:
        return

    out_dir = Path(args.out) / input_path.stem
    print(f"\nOutput directory: {out_dir}")
    write_chunks(reader, chapters, out_dir)


if __name__ == "__main__":
    main()