path: root/app/converter/extractors.py

"""Text extraction from book files (TXT, PDF, EPUB) and text/HTML cleaning."""

import codecs
import logging
import posixpath
import re
import zipfile
from html import unescape
from pathlib import Path
from typing import Dict, List, NamedTuple, Optional
from urllib.parse import unquote

try:
    from bs4 import BeautifulSoup
    BS4_AVAILABLE = True
except ImportError:
    BS4_AVAILABLE = False

logger = logging.getLogger(__name__)


class Section(NamedTuple):
    """A titled chunk of a book (e.g. an EPUB chapter)."""

    title: str
    text: str


class Book(NamedTuple):
    """A book's metadata plus its titled sections."""

    title: str
    author: str
    sections: List[Section]


def extract_text(file_path: Path) -> str:
    """Extract text from a book file based on its extension."""
    extension = file_path.suffix.lower()
    if extension == ".txt":
        return _extract_txt(file_path)
    if extension == ".pdf":
        return _extract_pdf(file_path)
    raise ValueError(f"Unsupported file format: {extension}")


def extract_sections(file_path: Path) -> List[Section]:
    """Extract the book's text as titled sections (chapters).

    EPUB files are split on their spine documents so they can be converted
    one chapter at a time. TXT and PDF files have no chapter structure and
    always yield a single section.
    """
    if file_path.suffix.lower() == ".epub":
        chapters = _extract_epub_chapters(file_path)
        if not chapters:
            raise RuntimeError("All EPUB extraction methods failed")
        return chapters

    return [Section(file_path.stem, extract_text(file_path))]


def extract_book(file_path: Path) -> Book:
    """Extract sections plus book-level metadata (title, author).

    EPUB and PDF files carry embedded metadata; missing fields (and TXT
    files, which have none) fall back to the file stem for the title and
    an empty author.
    """
    title, author = "", ""
    extension = file_path.suffix.lower()
    if extension == ".epub":
        title, author = _epub_metadata(file_path)
    elif extension == ".pdf":
        title, author = _pdf_metadata(file_path)
    return Book(title or file_path.stem, author.strip(), extract_sections(file_path))


def _epub_metadata(file_path: Path) -> tuple:
    """Return (title, author) from an EPUB's Dublin Core metadata."""
    try:
        from ebooklib import epub  # raises ImportError when unavailable

        book = epub.read_epub(str(file_path))
        title = _first_dc_value(book.get_metadata("DC", "title"))
        author = _first_dc_value(book.get_metadata("DC", "creator"))
        return title, author
    except Exception as exc:
        logger.warning("Could not read EPUB metadata: %s", exc)
        return "", ""


def _pdf_metadata(file_path: Path) -> tuple:
    """Return (title, author) from a PDF's document info dictionary."""
    try:
        from pypdf import PdfReader

        reader = PdfReader(str(file_path))
        info = reader.metadata or {}
        title = str(info.get("/Title") or "")
        author = str(info.get("/Author") or "")
        return title, author
    except Exception as exc:
        logger.warning("Could not read PDF metadata: %s", exc)
        return "", ""


def _first_dc_value(entries) -> str:
    """First value of an ebooklib DC metadata list: [(value, ...), ...]."""
    if not entries:
        return ""
    value = entries[0][0]
    return str(value).strip() if value else ""


def _extract_epub_chapters(file_path: Path) -> List[Section]:
    """Return one Section per EPUB spine document (chapter), in reading order."""
    book = None
    for method in (_read_epub_ebooklib, _read_epub_zipfile, _read_epub_manual):
        try:
            book = method(file_path)
        except Exception as exc:
            logger.warning("EPUB chapter method %s failed: %s", method.__name__, exc)
            continue
        if book:
            break

    if book is None:
        return []

    chapters = []
    for title, text in book:
        cleaned = clean_html(text)
        if cleaned.strip():
            chapters.append(Section(title or file_path.stem, cleaned))
    return chapters


def _toc_titles(book) -> dict:
    """Flatten an ebooklib TOC into a ``{href: title}`` mapping."""
    titles = {}

    def walk(nodes) -> None:
        for node in nodes:
            if isinstance(node, (tuple, list)):
                walk(node[1] if len(node) > 1 else [])
                continue
            href = getattr(node, "href", None)
            title = getattr(node, "title", None)
            if href and title:
                titles[href.split("#")[0]] = title

    walk(book.toc)
    return titles


def _read_epub_ebooklib(file_path: Path):
    """Read EPUB spine documents as (title, html) pairs via ebooklib."""
    import ebooklib
    from ebooklib import epub

    book = epub.read_epub(str(file_path))
    titles = _toc_titles(book)
    items = []
    for entry in book.spine:
        item_id = entry[0] if isinstance(entry, (tuple, list)) else entry
        try:
            item = book.get_item_with_id(item_id)
        except Exception as exc:
            logger.debug("Skipping EPUB spine item %r: %s", item_id, exc)
            continue
        if not item or item.get_type() != ebooklib.ITEM_DOCUMENT:
            continue
        if isinstance(item, epub.EpubNav):
            continue
        content = item.get_body_content()
        if content:
            if isinstance(content, bytes):
                content = content.decode("utf-8", errors="ignore")
            title = (titles.get(item.file_name)
                     or titles.get(item.get_name())
                     or getattr(item, "title", None)
                     or item.get_name())
            items.append((title, str(content)))
    return items


def _read_epub_zipfile(file_path: Path):
    """Read EPUB HTML members as (title, html) pairs, in spine order.

    Fallback for EPUBs ebooklib cannot read. The package's OPF describes
    the reading order (its ``<spine>`` itemrefs reference manifest items
    by id), so documents are emitted in that order; the manifest's
    ``properties="nav"`` item (the table of contents) and any document
    outside the spine are skipped so the TOC is never narrated as a
    chapter. EPUBs without a parsable OPF fall back to natural filename
    order over every HTML member.
    """
    items = []
    with zipfile.ZipFile(file_path, "r") as epub_zip:
        names = epub_zip.namelist()
        html_names = [name for name in names
                      if name.lower().endswith((".html", ".xhtml", ".htm"))]
        order = _epub_spine_order(epub_zip, html_names)
        if order is None:
            order = sorted(html_names, key=_natural_key)
        for file_name in order:
            try:
                content = epub_zip.read(file_name).decode("utf-8", errors="ignore")
                items.append((Path(file_name).stem, content))
            except Exception as exc:
                logger.debug("Skipping EPUB member %r: %s", file_name, exc)
    return items


def _epub_spine_order(epub_zip: zipfile.ZipFile, html_names: List[str]):
    """The EPUB's HTML members in spine order, or None when unparsable.

    Parses the package OPF (located via META-INF/container.xml, else the
    only *.opf member): manifest item id -> href, then the spine's
    idrefs. Hrefs are URL-encoded, may carry "./" or "../" segments and
    XML entities, and are relative to the OPF's folder — they are
    normalized before matching against the zip's member names, so valid
    chapters are not silently dropped when their written form differs
    from the archived path. Returns member paths limited to HTML_NAMES;
    nav documents (``properties`` containing "nav") and non-HTML items
    are excluded.
    """
    container = "META-INF/container.xml"
    opf_name = None
    try:
        rootfile = epub_zip.read(container).decode("utf-8", errors="ignore")
        match = re.search(r"full-path\s*=\s*[\"']([^\"']+)[\"']", rootfile)
        if match and match.group(1) in epub_zip.namelist():
            opf_name = match.group(1)
    except (KeyError, OSError):
        pass
    if opf_name is None:
        opf_candidates = [name for name in epub_zip.namelist()
                          if name.lower().endswith(".opf")]
        if len(opf_candidates) != 1:
            return None
        opf_name = opf_candidates[0]
    try:
        opf = epub_zip.read(opf_name).decode("utf-8", errors="ignore")
    except (KeyError, OSError):
        return None

    def attr(tag: str, name: str) -> Optional[str]:
        match = re.search(rf"\b{name}\s*=\s*[\"']([^\"']*)[\"']", tag)
        return match.group(1) if match else None

    base = "/".join(opf_name.split("/")[:-1])
    item_tags = re.findall(r"<item\b[^>]*>", opf)

    # Zip member names indexed under their normalized spellings, so an
    # exact archive path can be found from any equivalent href form.
    # (Percent-encoded and plain spellings both map to the same member.)
    members: Dict[str, str] = {}
    for name in html_names:
        members.setdefault(posixpath.normpath(name), name)
        members.setdefault(posixpath.normpath(unquote(name)), name)

    def member_for_href(href: str) -> Optional[str]:
        """The zip member an OPF manifest HREF points at (best effort)."""
        path = unescape(href).strip()
        path = path.split("#", 1)[0]
        path = unquote(path).strip()
        if not path:
            return None
        if base:
            path = f"{base}/{path}"
        return members.get(posixpath.normpath(path).lstrip("/"))

    def is_nav(item_tag: str) -> bool:
        properties = attr(item_tag, "properties") or ""
        return "nav" in properties.split()

    order: List[str] = []
    for ref_tag in re.findall(r"<itemref\b[^>]*>", opf):
        idref = attr(ref_tag, "idref")
        if not idref:
            continue
        match = next((item_tag for item_tag in item_tags
                      if attr(item_tag, "id") == idref), None)
        if match is None or is_nav(match):
            continue
        href = attr(match, "href")
        if not href:
            continue
        member = member_for_href(href)
        if member and member not in order:
            order.append(member)
    return order or None


def _read_epub_manual(file_path: Path):
    """Last-resort read of any markup-looking EPUB member."""
    skipped_extensions = (".jpg", ".jpeg", ".png", ".gif", ".css", ".js")
    items = []
    with zipfile.ZipFile(file_path, "r") as epub_zip:
        for file_name in sorted(epub_zip.namelist(), key=_natural_key):
            if file_name.lower().endswith(skipped_extensions):
                continue
            try:
                content = epub_zip.read(file_name).decode("utf-8", errors="ignore")
                if "<" in content and len(content.strip()) > 100:
                    items.append((Path(file_name).stem, content))
            except Exception as exc:
                logger.debug("Skipping EPUB member %r: %s", file_name, exc)
    return items


def clean_text(text: str) -> str:
    """Normalize whitespace, strip standalone page numbers, and keep
    paragraph breaks.

    Page numbers are removed only when a short number (up to three digits)
    appears alone on its own line (before whitespace collapsing), so inline
    numbers like "42 years", "1,000" or "3.5" are preserved, as are
    four-digit standalone lines, which are usually years ("1984") or
    chapter numbers rather than page numbers.

    Runs of blank lines survive as a single paragraph break (``\\n\\n``) so
    the smart chunker can still end requests at paragraph boundaries and
    reset its quotation state there; single line breaks (soft wraps) and
    all other whitespace collapse to spaces.
    """
    if not text:
        return ""
    # Standalone page numbers (digits alone on a line) must go BEFORE the
    # whitespace-collapsing steps below.
    text = re.sub(r"(?m)^\s*\d{1,3}\s*$", " ", text)
    # Normalize line breaks so the steps below see plain "\n".
    text = re.sub(r"\r\n?", "\n", text)
    # Any run of blank lines becomes exactly one paragraph break.
    text = re.sub(r"\n[^\S\n]*\n(?:[^\S\n]*\n)*", "\n\n", text)
    # Remaining single line breaks (soft wraps) become spaces.
    text = re.sub(r"(?<!\n)\n(?!\n)", " ", text)
    # Collapse every other whitespace run (including nbsp) to one space.
    text = re.sub(r"[^\S\n]+", " ", text)
    # Trim whatever sits against a paragraph break.
    text = re.sub(r"[^\S\n]*\n\n[^\S\n]*", "\n\n", text)
    return text.strip()


# Block-level elements structure the document into paragraphs and lines:
# replacing them with a paragraph break keeps the chunker's paragraph
# boundaries (and headings separated from prose). Inline elements must NOT
# break words, so they are left for text extraction to unwrap without
# spaces ("un<em>believ</em>able" stays one word). The lookahead keeps
# lookalike tags (<parameter>, <preheat>) out of the match.
_BLOCK_BREAK_RE = re.compile(
    r"</?(?:address|article|aside|blockquote|body|br|caption|center|col"
    r"|colgroup|dd|div|dl|dt|fieldset|figcaption|figure|footer|form"
    r"|h[1-6]|head|header|hr|html|legend|li|main|nav|ol|option|p|pre"
    r"|section|summary|table|tbody|td|tfoot|th|thead|tr|ul)(?=[\s/>])"
    r"[^>]*>",
    re.IGNORECASE)


def clean_html(html_content: str) -> str:
    """Strip markup, scripts and styles from HTML content.

    Block-level tags become paragraph breaks (so the chunker sees the
    document's paragraphs); inline tags vanish without adding spaces, so
    no artificial word boundaries are injected.
    """
    if not html_content:
        return ""

    html_content = _BLOCK_BREAK_RE.sub("\n\n", html_content)

    if BS4_AVAILABLE:
        try:
            soup = BeautifulSoup(html_content, "html.parser")
            for tag in soup(["script", "style"]):
                tag.decompose()
            return clean_text(soup.get_text(separator=""))
        except Exception as exc:
            logger.debug("BeautifulSoup cleaning failed, falling back to regex: %s", exc)

    # Fallback regex cleaning. Block tags are already paragraph breaks;
    # every other tag (inline) disappears without adding a space, so no
    # artificial word boundary is injected.
    html_content = re.sub(r"<style[^>]*>.*?</style>", "", html_content, flags=re.DOTALL | re.IGNORECASE)
    html_content = re.sub(r"<script[^>]*>.*?</script>", "", html_content, flags=re.DOTALL | re.IGNORECASE)
    html_content = re.sub(r"<[^>]+>", "", html_content)
    html_content = unescape(html_content)
    return clean_text(html_content)


def _natural_key(name: str):
    """Sort key that orders numeric runs numerically (chapter2 before chapter10)."""
    return [int(part) if part.isdigit() else part.lower()
            for part in re.split(r"(\d+)", name)]


def _extract_txt(file_path: Path) -> str:
    """Extract from TXT, handling BOMs and common encodings (latin-1 is the catch-all).

    UTF-16 files without a BOM are detected via NUL bytes; otherwise they would
    silently decode as NUL-interleaved UTF-8 or cp1252/latin-1 garbage.
    """
    data = file_path.read_bytes()

    if data.startswith((codecs.BOM_UTF32_LE, codecs.BOM_UTF32_BE)):
        return clean_text(data.decode("utf-32"))
    if data.startswith((codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE)):
        return clean_text(data.decode("utf-16"))
    if data.startswith(codecs.BOM_UTF8):
        return clean_text(data.decode("utf-8-sig"))

    # No BOM: UTF-16 without BOM is common on Windows; detect via NUL bytes.
    # A real UTF-16 file of ASCII-range text has a NUL at every other byte
    # position, so require a substantial NUL share before committing to
    # UTF-16: a lone stray NUL in a UTF-8/cp1252 file must not flip the
    # whole book into mojibake (the decode ladder below handles that).
    sample = data[:4096]
    even_nuls = sum(1 for i, b in enumerate(sample) if b == 0 and i % 2 == 0)
    odd_nuls = sum(1 for i, b in enumerate(sample) if b == 0 and i % 2 == 1)
    threshold = max(len(sample) // 4, 1)
    if even_nuls >= threshold or odd_nuls >= threshold:
        encoding = "utf-16-be" if even_nuls > odd_nuls else "utf-16-le"
        return clean_text(data.decode(encoding))

    for encoding in ("utf-8", "cp1252", "latin-1"):
        try:
            return clean_text(data.decode(encoding))
        except UnicodeError:
            continue
    raise ValueError(f"Could not decode text file: {file_path}")


def _extract_pdf(file_path: Path) -> str:
    """Extract from PDF."""
    from pypdf import PdfReader

    text = ""
    with open(file_path, "rb") as file:
        pdf_reader = PdfReader(file)
        total_pages = len(pdf_reader.pages)
        logger.info("PDF has %d pages", total_pages)

        for page_num, page in enumerate(pdf_reader.pages, 1):
            try:
                page_text = page.extract_text() or ""
                if page_text.strip():
                    text += f"\n\n{page_text}"
                if page_num % 10 == 0:
                    logger.debug("Extracted %d/%d pages", page_num, total_pages)
            except Exception as exc:
                logger.warning("Failed to extract page %d: %s", page_num, exc)

        logger.info("Extracted text from %d pages, %d characters total", total_pages, len(text))
    return clean_text(text)