bindery/scripts/import_unlinked_epub_chapters.py

#!/usr/bin/env python3
"""Fill a gap in a tracked story with chapters from an ePub that has no chapter links.

Bindery can import chapters from an ePub that FanFicFare made, because that
ePub records the site address (the link) of each chapter. Other ePubs, such as
the author's own ebook of the story, have no chapter links, so Bindery cannot
tell where their chapters go. Use this script when the site no longer has some
chapters, and you have an ePub without chapter links that has them.

Usage:

    CONFIG_DIR=./config uv run python scripts/import_unlinked_epub_chapters.py <story-url> <path-to-epub>
    CONFIG_DIR=./config uv run python scripts/import_unlinked_epub_chapters.py <story-url> --remove

Run the script on the host, in the project folder, with CONFIG_DIR set to the
folder that holds bindery.db. The container image does not include the
scripts folder. Stop Bindery before you run the script, so that nothing else
writes to the database at the same time.

The story must be tracked in Bindery, with at least one stored chapter. The
script asks four questions:

1. Which stored chapter comes just before the gap. Type "start" if the gap
   is at the start of the story.
2. Which ePub chapter is the same as that stored chapter. After "start", the
   question is which ePub chapter to add first, so that you can leave out the
   cover and the title pages.
3. Which stored chapter comes just after the gap. Type "end" if the gap is at
   the end of the story.
4. Which ePub chapter is the same as that stored chapter. After "end", the
   question is which ePub chapter to add last.

The script suggests an answer where the chapter titles match, and puts it on
the input line. Press Enter to accept it, or change it. Type "list" to see all
the chapters again. The script then shows the ePub chapters that go into the
gap, in the order of the ePub, and adds them after you say yes. Their images
come with them. A second run with the same ePub does not add the same chapters
again.

The added chapters have no address on the site, so they show as "removed from
site" on the story page. A "full" download takes them out of the book, like
other chapters that the site does not list. Select them and use "Put back in
book" to put them back. To delete them, use "Delete…" on the story page, or run
the script with --remove. Rebuild the book after each change.
"""

import argparse
import hashlib
import posixpath
import re
import sqlite3
import sys
import warnings
import zipfile
from difflib import SequenceMatcher
from pathlib import Path

from bs4 import BeautifulSoup, XMLParsedAsHTMLWarning
from defusedxml import ElementTree

try:
    import readline
except ImportError:  # Windows
    readline = None

# Run as a file, Python puts scripts/ on the path, not the project folder that holds bindery/.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from fanficfare.story import imagetypes  # noqa: E402

from bindery import config, fff  # noqa: E402
from bindery.diffing import text_hash  # noqa: E402
from bindery.store import Store  # noqa: E402

OPF = "{http://www.idpf.org/2007/opf}"
NCX = "{http://www.daisy.org/z3986/2005/ncx/}"
DC = "{http://purl.org/dc/elements/1.1/}"
CONTAINER = "{urn:oasis:names:tc:opendocument:xmlns:container}"
URL_MARK = "#bindery-import-"
SOURCE = "epub-import"
# Images get a file: URL, so a build never asks the story site for them.
IMAGE_URL = "file:///bindery-epub-import/"
# A long list shows only this many chapters around the one that the question is about.
WINDOW = 12

warnings.filterwarnings("ignore", category=XMLParsedAsHTMLWarning)


class ScriptError(Exception):
    """A problem the user can fix."""


def _resolve(folder: str, href: str) -> str:
    """The path in the zip file of a link from a file in `folder`, without its #fragment."""
    return posixpath.normpath(posixpath.join(folder, href.split("#", 1)[0]))


def _toc_titles(z: zipfile.ZipFile, opf, folder: str) -> dict[str, str]:
    """The title of each file in the ePub's table of contents (EPUB 2 NCX or EPUB 3 nav), by path.

    The first entry of a file wins. Entries that are only a number, such as page numbers, are left out.
    """
    items = list(opf.find(f"{OPF}manifest"))
    titles: dict[str, str] = {}

    def add(path: str, label: str) -> None:
        label = " ".join(label.split())
        if label and not label.isdigit():
            titles.setdefault(path, label)

    ncx = next((i for i in items if i.get("media-type") == "application/x-dtbncx+xml"), None)
    if ncx is not None:
        ncx_path = _resolve(folder, ncx.get("href"))
        nav_map = ElementTree.fromstring(z.read(ncx_path)).find(f"{NCX}navMap")
        for point in nav_map.iter(f"{NCX}navPoint") if nav_map is not None else []:
            text, content = point.find(f"{NCX}navLabel/{NCX}text"), point.find(f"{NCX}content")
            if text is not None and content is not None and content.get("src"):
                add(_resolve(posixpath.dirname(ncx_path), content.get("src")), text.text or "")
    nav = next((i for i in items if "nav" in (i.get("properties") or "").split()), None)
    if nav is not None:
        nav_path = _resolve(folder, nav.get("href"))
        soup = BeautifulSoup(z.read(nav_path), "html5lib")
        toc = soup.find("nav", attrs={"epub:type": "toc"}) or soup.find("nav")
        for a in toc.find_all("a", href=True) if toc else []:
            add(_resolve(posixpath.dirname(nav_path), a["href"]), a.get_text(" "))
    return titles


def _only_links(soup: BeautifulSoup) -> bool:
    """Whether all the text of a file is in links, as in a table of contents."""
    texts = [t for t in soup.body.find_all(string=True) if t.strip()] if soup.body else []
    return bool(texts) and all(t.find_parent("a") for t in texts)


def _book_id(opf, path: Path) -> str:
    """A short ID of the ePub, from its identifier, or else its file name."""
    identifier = opf.find(f"{OPF}metadata/{DC}identifier")
    key = identifier.text.strip() if identifier is not None and identifier.text else path.name
    return hashlib.sha256(key.encode("utf-8")).hexdigest()[:8]


def _take_images(z: zipfile.ZipFile, soup: BeautifulSoup, folder: str, book: str) -> list[dict]:
    """Point each <img> of a chapter at a URL of its own, and return the images with those URLs.

    An image that the ePub does not have, or that FanFicFare cannot use, is left out.
    """
    images = []
    for img in soup.body.find_all("img", src=True):
        if img["src"].startswith(("data:", "http:", "https:")):
            continue
        path = _resolve(folder, img["src"])
        ext = path.rsplit(".", 1)[-1].lower()
        if ext in imagetypes and path in z.namelist():
            img["src"] = f"{IMAGE_URL}{book}/{path}"
            images.append({"url": img["src"], "ext": ext, "mime": imagetypes[ext], "data": z.read(path)})
        else:
            img.decompose()
    return images


def read_epub(path: Path) -> list[dict]:
    """Read the chapters of an ePub, in reading order, with the images that each one uses.

    The script reads the OPF manifest and spine directly, so it can read any
    ePub. A file that has only links, such as a table of contents, is not a
    chapter. Titles come from the table of contents of the ePub, or else from
    the first heading of the chapter. Each chapter has a stable `key`: the ID
    of the ePub and the path of the chapter in it.
    """
    try:
        with zipfile.ZipFile(path) as z:
            container = ElementTree.fromstring(z.read("META-INF/container.xml"))
            rootfile = container.find(f".//{CONTAINER}rootfile")
            if rootfile is None or not rootfile.get("full-path"):
                raise ScriptError(f"{path.name} has no OPF file")
            opf_path = rootfile.get("full-path")
            opf = ElementTree.fromstring(z.read(opf_path))
            folder = posixpath.dirname(opf_path)
            manifest = {item.get("id"): item for item in opf.find(f"{OPF}manifest")}
            toc = _toc_titles(z, opf, folder)
            book = _book_id(opf, path)
            chapters = []
            for itemref in opf.find(f"{OPF}spine").findall(f"{OPF}itemref"):
                item = manifest.get(itemref.get("idref"))
                if itemref.get("linear") == "no" or item is None:
                    continue
                if "nav" in (item.get("properties") or "").split():
                    continue
                href = item.get("href") or ""
                if not href.lower().endswith((".xhtml", ".html", ".htm")):
                    continue
                full = _resolve(folder, href)
                soup = BeautifulSoup(z.read(full), "html5lib")
                if not soup.body or _only_links(soup):
                    continue
                heading = soup.find(["h1", "h2", "h3"])
                if full in toc:
                    title = toc[full]
                elif heading and heading.get_text(strip=True):
                    title = heading.get_text(" ", strip=True)
                elif soup.title and soup.title.get_text(strip=True):
                    title = soup.title.get_text(strip=True)
                else:
                    title = href
                images = _take_images(z, soup, posixpath.dirname(full), book)
                html = "".join(str(c) for c in soup.body.contents).strip()
                if html:
                    chapters.append({"title": title, "html": html, "key": f"{book}-{full}", "images": images})
    except (zipfile.BadZipFile, KeyError, ElementTree.ParseError) as e:
        raise ScriptError(f"{path.name} is not a readable ePub: {e}") from e
    if not chapters:
        raise ScriptError(f"{path.name} has no chapters that the script can read")
    return chapters


def suggest_matches(title: str, candidates: list[str], limit: int = 3) -> list[int]:
    """Indexes of the candidates most like `title`, best first."""
    ratio = {i: SequenceMatcher(None, title.lower(), c.lower()).ratio() for i, c in enumerate(candidates)}
    return sorted((i for i in ratio if ratio[i] > 0.4), key=lambda i: ratio[i], reverse=True)[:limit]


def chapter_url(story_url: str, chapter: dict) -> str:
    """A made-up, stable URL for a chapter that has no link. The same ePub chapter always gets the same URL."""
    slug = re.sub(r"[^a-z0-9]+", "-", chapter["key"].lower()).strip("-")
    return f"{story_url.rstrip('/')}{URL_MARK}{slug}"


def splice_order(existing_urls: list[str], before_idx: int, new_urls: list[str]) -> list[str]:
    """`existing_urls` with `new_urls` put directly after position `before_idx` (-1 for the start)."""
    return existing_urls[: before_idx + 1] + new_urls + existing_urls[before_idx + 1 :]


def _show(items: list[str], marks: set[int] = frozenset(), around: int | None = None) -> None:
    """Print a numbered list. With `around`, a long list shows only the part near that index."""
    first, last = 0, len(items)
    if around is not None and len(items) > WINDOW * 2:
        first = max(0, min(around - WINDOW // 2, len(items) - WINDOW))
        last = first + WINDOW
    # The numbers go in brackets, padded so that the titles line up. A suggestion has a * in front.
    width = len(str(len(items))) + 2
    if first:
        print("  …")
    for i in range(first, last):
        print(f"{'*' if i in marks else ' '} {f'[{i + 1}]':<{width}} {items[i]}")
    if last < len(items):
        print(f"  … (type 'list' to see all {len(items)})")


def _ask(prompt: str, default: str) -> str:
    """Read a line. The default is already on the input line, where the terminal allows it."""
    prefill = bool(default) and readline is not None and sys.stdin.isatty()
    if prefill:
        readline.set_startup_hook(lambda: readline.insert_text(default))
    elif default:
        prompt = f"{prompt[:-2]} [{default}]: "
    try:
        answer = input(prompt).strip()
    finally:
        if prefill:
            readline.set_startup_hook()
    return answer or default


def choose(
    prompt: str,
    items: list[str],
    *,
    word: str | None = None,
    default: int | None = None,
    marks: set[int] = frozenset(),
    around: int | None = None,
    minimum: int = 0,
) -> int | str:
    """Show `items` and ask for one of them. Returns its index, or `word` if the user typed it.

    The prompt says what the question is, so the user does not have to look above a long list for it.
    """
    _show(items, marks, around)
    words = ["'list'"] + ([f"'{word}'"] if word else [])
    while True:
        answer = _ask(f"{prompt}: ", "" if default is None else str(default + 1)).lower()
        if answer == "list":
            _show(items, marks)
        elif word and answer == word:
            return word
        elif answer.isdigit() and minimum < int(answer) <= len(items):
            return int(answer) - 1
        else:
            print(f"Type a number from {minimum + 1} to {len(items)}, or {' or '.join(words)}.")


def _confirm(prompt: str) -> bool:
    return input(f"{prompt} [y/N] ").strip().lower() == "y"


def _added_by_script(url: str) -> bool:
    return URL_MARK in url


def remove_added_chapters(db_store: Store, story: dict, existing: list[dict]) -> None:
    added = [c for c in existing if _added_by_script(c["url"])]
    if not added:
        print(f"{story['title']} has no chapters that this script added.")
        return
    print(f"{story['title']} has {_count(len(added), 'chapter')} that this script added:")
    for c in added:
        print(f"  - {fff.plain_title(c['title'])}")
    if not _confirm("Delete them and all their revisions?"):
        return
    conn = sqlite3.connect(config.DB_PATH)
    try:
        # Without this, SQLite keeps the revisions of a deleted chapter.
        conn.execute("PRAGMA foreign_keys=ON")
        with conn:
            conn.executemany("DELETE FROM chapters WHERE id = ?", [(c["id"],) for c in added])
    finally:
        conn.close()
    print(f"Deleted {_count(len(added), 'chapter')}. Rebuild {story['title']} in Bindery.")


def _count(n: int, word: str) -> str:
    return f"{n} {word}{'' if n == 1 else 's'}"


def ask_for_gap(
    story_title: str, titles: list[str], labels: list[str], epub_name: str, epub_titles: list[str]
) -> tuple[int, int, int]:
    """Ask where the gap is. Returns the index of the stored chapter before the gap (-1 for the start),
    and the indexes of the first and last ePub chapters to add.

    `titles` are the stored chapter titles, and `labels` are the same titles as the lists show them.
    """
    print(f"\n{story_title} has {_count(len(titles), 'stored chapter')}:")
    before = choose("Stored chapter just BEFORE the gap (or 'start')", labels, word="start")
    print(f"\n{epub_name} has {_count(len(epub_titles), 'chapter')}:")
    if before == "start":
        before = -1
        first = choose("FIRST ePub chapter to add", epub_titles)
    else:
        guess = suggest_matches(titles[before], epub_titles)
        match = choose(
            f'ePub chapter that matches "{titles[before]}"',
            epub_titles,
            default=guess[0] if guess else None,
            marks=set(guess),
        )
        first = match + 1

    print(f"\n{story_title}, near the gap:")
    if before + 1 < len(titles):
        after = choose(
            "Stored chapter just AFTER the gap (or 'end')",
            labels,
            word="end",
            default=before + 1,
            around=before + 1,
            minimum=before + 1,
        )
    else:
        print(f'  None. The gap is after the last stored chapter, "{titles[before]}".')
        after = "end"
    print(f"\n{epub_name}, near the gap:")
    if after == "end":
        last = choose(
            "LAST ePub chapter to add", epub_titles, default=len(epub_titles) - 1, around=first, minimum=first
        )
    else:
        if after != before + 1:
            skipped = _count(after - before - 1, "stored chapter")
            print(f"Note: the {skipped} between those two stay where they are, after the added chapters.")
        later = epub_titles[first:]
        guess = [first + i for i in suggest_matches(titles[after], later)]
        match = choose(
            f'ePub chapter that matches "{titles[after]}"',
            epub_titles,
            default=guess[0] if guess else None,
            marks=set(guess),
            around=guess[0] if guess else first,
            minimum=first,
        )
        last = match - 1
    return before, first, last


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("story_url", help="The URL of the tracked story")
    parser.add_argument("epub", type=Path, nargs="?", help="The ePub that has the missing chapters")
    parser.add_argument("--remove", action="store_true", help="Delete the chapters that this script added")
    args = parser.parse_args(argv)

    normal = fff.normalize_url(args.story_url)
    if not normal:
        raise ScriptError(f"Not a supported story URL: {args.story_url}")
    if not args.remove and (args.epub is None or not args.epub.is_file()):
        raise ScriptError(f"No such file: {args.epub}" if args.epub else "Give the path to the ePub.")

    db_store = Store(config.DB_PATH)
    story = db_store.get_story_by_url(normal[0])
    if not story:
        raise ScriptError(f"No tracked story at {normal[0]}. Add it in Bindery first.")

    existing = db_store.list_chapters(story["id"])
    if args.remove:
        remove_added_chapters(db_store, story, existing)
        return
    if not existing:
        raise ScriptError(f"{story['title']} has no stored chapters yet. Download it in Bindery first.")

    epub_chapters = read_epub(args.epub)
    titles = [fff.plain_title(c["title"]) for c in existing]
    labels = [t + ("" if c["on_site"] else "  (not on site)") for t, c in zip(titles, existing, strict=True)]
    epub_titles = [c["title"] for c in epub_chapters]

    before, first, last = ask_for_gap(story["title"], titles, labels, args.epub.name, epub_titles)
    stored_urls = {c["url"] for c in existing}
    to_add = [c for c in epub_chapters[first : last + 1] if chapter_url(story["url"], c) not in stored_urls]
    already = last + 1 - first - len(to_add)
    if already:
        print(f"\n{_count(already, 'ePub chapter')} of the gap are in the story already, from an earlier run.")
    if not to_add:
        print("There are no chapters to add.")
        return

    where = "at the start" if before == -1 else f'after "{titles[before]}"'
    print(f"\nThe script will add {'this chapter' if len(to_add) == 1 else f'these {len(to_add)} chapters'} {where}:")
    for c in to_add:
        print(f"  - {c['title']}")
    if not _confirm("Add it?" if len(to_add) == 1 else "Add them?"):
        return

    new_chapters = [
        {
            "url": chapter_url(story["url"], c),
            "title": fff.escape_title(c["title"]),
            "html": c["html"],
            "text_hash": text_hash(c["html"]),
        }
        for c in to_add
    ]
    for img in {img["url"]: img for c in to_add for img in c["images"]}.values():
        db_store.put_image(story["id"], img["url"], img["ext"], img["mime"], img["data"])
    order = splice_order([c["url"] for c in existing], before, [c["url"] for c in new_chapters])
    for chapter_id in db_store.add_chapters(story["id"], new_chapters, SOURCE, order):
        # The site does not list these chapters. The next download would set this too.
        db_store.update_chapter(chapter_id, on_site=0)
    print(f"\nAdded {_count(len(new_chapters), 'chapter')} to {story['title']}. Rebuild the book in Bindery.")


if __name__ == "__main__":
    try:
        main()
    except ScriptError as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
    except (KeyboardInterrupt, EOFError):
        print("\nStopped. Nothing changed.")
        sys.exit(1)