Book-Reader-Assist/tools/split_epub_book.py

"""Export configured EPUB books into chapter-scoped source-text Markdown."""
import json
import posixpath
import re
import sys
import zipfile
from pathlib import Path
from xml.etree import ElementTree as ET

from lxml import html


ROOT = Path(__file__).resolve().parent.parent / "library"
XLINK = "{http://www.w3.org/1999/xlink}href"
BLOCK_TAGS = {
    "address", "article", "aside", "blockquote", "br", "caption", "div",
    "figcaption", "figure", "footer", "h1", "h2", "h3", "h4", "h5", "h6",
    "header", "hr", "li", "main", "nav", "ol", "p", "section", "table",
    "tbody", "td", "tfoot", "th", "thead", "tr", "ul",
}


def ch(number):
    return f"Chapter-{number:02d}"


def slug(title, limit=60):
    text = re.sub(r"^\d+(\.\d+)?[a-z]?\s*[-)]\s*", "", title.strip())
    text = re.sub(r"[:,.'\"‘’“”?/\\()]", "", text)
    text = re.sub(r"[^\w\s-]", "", text)
    words, output = text.split(), ""
    for word in words:
        candidate = f"{output}-{word}" if output else word
        if len(candidate) > limit:
            break
        output = candidate
    return re.sub(r"-+", "-", output) or "Untitled"


def folder_name(key, title):
    return key if key in ("Front-Matter", "Back-Matter") else f"{key}-{slug(title)}"


def outrageous_sections():
    front = ["titlepage.xhtml", "text/part0000.html", "text/part0001.html"]
    front += [f"text/part{number:04d}.html" for number in range(2, 8)]
    titles = [
        "Take The OUTRAGEOUS Path To OUTRAGEOUS Success",
        "Defining Advertising And The Three Possible Responses",
        "Applying OUTRAGEOUS Advertising",
        "Using The Most Powerful Tool - Direct Mail",
        "Presenting Critical Components",
        "Going The Extra Mile For Extra Money",
        "Harnessing The Power Of Testimonials",
        "Using Premiums To Skyrocket Response",
        "Reminding Sequentially Until They Buy",
        "Swiping Ideas And Deploying Them Into Your Business",
        "Celebrating Holidays & Special Occasions",
        "Sampling Some Really Outrageous Ideas",
        "Putting All The Pieces Together",
    ]
    sections = [("Front-Matter", "Front matter, endorsements, foreword, preface, and contents", front)]
    for number, title in enumerate(titles, start=1):
        sections.append((ch(number), title, [f"text/part{number + 7:04d}.html"]))
    sections.append((
        "Back-Matter",
        "Appendices A-C and bonus material",
        [f"text/part{number:04d}.html" for number in range(21, 25)],
    ))
    return sections


BOOKS = {
    "Outrageous Advertising That's Outrageously Successful": outrageous_sections,
}


def epub_spine(epub):
    container = ET.fromstring(epub.read("META-INF/container.xml"))
    opf_path = container.find(".//{*}rootfile").attrib["full-path"]
    opf = ET.fromstring(epub.read(opf_path))
    base = posixpath.dirname(opf_path)
    manifest = {
        item.attrib["id"]: item.attrib
        for item in opf.findall(".//{*}manifest/{*}item")
    }
    return [
        posixpath.normpath(posixpath.join(base, manifest[item.attrib["idref"]]["href"]))
        for item in opf.findall(".//{*}spine/{*}itemref")
    ]


def normalize_text(parts):
    text = "".join(parts).replace("\u00a0", " ")
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r" *\n *", "\n", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def document_text(epub, document, image_ocr):
    root = html.fromstring(epub.read(document))
    for unwanted in root.xpath("//*[local-name()='style' or local-name()='script']"):
        unwanted.drop_tree()
    parts = []

    def walk(element):
        tag = element.tag.split("}")[-1].lower() if isinstance(element.tag, str) else ""
        if tag in ("img", "image"):
            href = element.get("src") or element.get(XLINK) or element.get("xlink:href") or element.get("href")
            if href:
                ref = posixpath.normpath(posixpath.join(posixpath.dirname(document), href))
                record = image_ocr.get(ref, {})
                asset = record.get("asset", ref)
                parts.append(f"\n\n[Embedded image: {asset}]\n")
                ocr_text = record.get("text", "").strip()
                if ocr_text:
                    parts.append("[OCR from embedded image; verify layout in the extracted image asset.]\n")
                    parts.append(ocr_text + "\n")
                else:
                    parts.append("[No readable text recognized; consult the extracted image asset.]\n")
            return

        if element.text:
            parts.append(element.text)
        for child in element:
            walk(child)
            if child.tail:
                parts.append(child.tail)
        if tag in BLOCK_TAGS:
            parts.append("\n")

    walk(root)
    return normalize_text(parts)


def main(book):
    book_root = ROOT / book
    epub_path = book_root / "source-file.epub"
    ocr_path = book_root / "ocr-embedded-images.json"
    image_ocr = json.loads(ocr_path.read_text(encoding="utf-8")) if ocr_path.exists() else {}
    sections = BOOKS[book]()

    with zipfile.ZipFile(epub_path) as epub:
        spine = epub_spine(epub)
        covered = [item for _, _, items in sections for item in items]
        missing = [item for item in spine if item not in covered]
        extras = [item for item in covered if item not in spine]
        duplicates = sorted({item for item in covered if covered.count(item) > 1})
        if missing or extras or duplicates:
            raise SystemExit(f"spine mismatch missing={missing} extras={extras} duplicates={duplicates}")

        for key, title, documents in sections:
            body = []
            for document in documents:
                body.append(f"<!-- EPUB item {document} -->\n{document_text(epub, document, image_ocr)}\n")
            out = book_root / folder_name(key, title)
            out.mkdir(exist_ok=True)
            header = [
                f"# {book} — {key}: {title}",
                f"- **Source**: /library/{book}/source-file.epub",
                f"- **EPUB items**: {', '.join(documents)}",
                "- **Image note**: Embedded images were extracted to the book's images folder; OCR is included beside image markers when text was recognized.",
            ]
            (out / f"{key}-source-text.md").write_text(
                "\n".join(header) + "\n\n---\n\n" + "\n".join(body),
                encoding="utf-8",
                newline="\n",
            )
            print(f"{key:14} items={documents} {title[:55]}")
        print(f"spine items covered exactly once: {len(spine)}")


if __name__ == "__main__":
    main(sys.argv[1])