cloud-reader/backend/apps/books/services/process_ebook.py

"""Extract TOC, metadata, and page count from uploaded EPUB/PDF files."""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any

from apps.books.models import BookChapter, EBook

logger = logging.getLogger(__name__)


def process_ebook(file_path: str, original_filename: str | None = None) -> dict[str, Any]:
    """Parse an e-book file and return format, metadata, TOC, and page count."""
    ext = Path(original_filename or file_path).suffix.lower()
    if ext == ".epub" or file_path.lower().endswith(".epub"):
        return _process_epub(file_path)
    if ext == ".pdf" or file_path.lower().endswith(".pdf"):
        return _process_pdf(file_path)
    return {"format": ext.lstrip(".") or "unknown", "page_count": 0, "metadata": {}, "toc": []}


def store_chapters(ebook: EBook, toc: list[dict[str, Any]], parent_index: int = 0) -> None:
    """Recursively store TOC entries as BookChapter records."""
    for idx, entry in enumerate(toc):
        BookChapter.objects.create(
            ebook=ebook,
            title=entry.get("title", "Untitled"),
            index=parent_index + idx,
            href=entry.get("href", ""),
            children=entry.get("children", []),
        )
        children = entry.get("children", [])
        if children:
            store_chapters(ebook, children, parent_index + idx + 1)


def apply_processing_to_ebook(ebook: EBook) -> dict[str, Any]:
    """Run processing on an EBook instance and persist chapters + metadata."""
    if not ebook.file:
        raise ValueError("No file found for this e-book.")

    file_path = ebook.file.path
    result = process_ebook(file_path, original_filename=ebook.filename())

    ebook.format = result.get("format", ebook.format)
    ebook.page_count = result.get("page_count", 0)
    file_metadata = result.get("metadata") or {}
    if file_metadata:
        merged = dict(ebook.metadata_json or {})
        merged["file"] = file_metadata
        ebook.metadata_json = merged
    ebook.save(update_fields=["format", "page_count", "metadata_json", "updated_at"])

    raw_toc: list[dict[str, Any]] = result.get("toc", [])
    BookChapter.objects.filter(ebook=ebook).delete()
    store_chapters(ebook, raw_toc)

    return {
        "format": ebook.format,
        "page_count": ebook.page_count,
        "metadata": ebook.metadata_json,
        "toc_count": len(raw_toc),
        "status": "processed",
    }


def _process_epub(file_path: str) -> dict[str, Any]:
    from ebooklib import epub, ITEM_DOCUMENT

    book = epub.read_epub(file_path)
    metadata = _extract_epub_metadata(book)
    toc = _extract_epub_toc(book)

    if not toc:
        toc = _fallback_toc_from_spine(book, ITEM_DOCUMENT)

    flat_count = _count_toc_entries(toc)
    page_count = flat_count or len(book.spine)

    return {
        "format": "epub",
        "page_count": page_count,
        "metadata": metadata,
        "toc": toc,
    }


def _process_pdf(file_path: str) -> dict[str, Any]:
    return {"format": "pdf", "page_count": 0, "metadata": {}, "toc": []}


def _extract_epub_metadata(book) -> dict[str, Any]:
    metadata: dict[str, Any] = {}

    def first(namespace: str, name: str) -> str:
        values = book.get_metadata(namespace, name)
        if values:
            return str(values[0][0])
        return ""

    title = first("DC", "title")
    if title:
        metadata["title"] = title
    creator = first("DC", "creator")
    if creator:
        metadata["author"] = creator
    language = first("DC", "language")
    if language:
        metadata["language"] = language
    identifier = first("DC", "identifier")
    if identifier:
        metadata["identifier"] = identifier
    return metadata


def _parse_toc_item(item) -> dict[str, Any]:
    from ebooklib import epub

    if isinstance(item, epub.Link):
        return {
            "title": item.title or "Untitled",
            "href": item.href or "",
            "children": [],
        }
    if isinstance(item, tuple):
        section, children = item
        entry = {
            "title": getattr(section, "title", None) or "Untitled",
            "href": getattr(section, "href", None) or "",
            "children": [],
        }
        for child in children:
            entry["children"].append(_parse_toc_item(child))
        return entry
    if hasattr(item, "title"):
        return {
            "title": item.title or "Untitled",
            "href": getattr(item, "href", "") or "",
            "children": [],
        }
    return {"title": "Untitled", "href": "", "children": []}


def _extract_epub_toc(book) -> list[dict[str, Any]]:
    return [_parse_toc_item(item) for item in book.toc]


def _fallback_toc_from_spine(book, item_document_type) -> list[dict[str, Any]]:
    toc: list[dict[str, Any]] = []
    seen: set[str] = set()
    chapter_num = 0

    for spine_entry in book.spine:
        item_id = spine_entry[0] if isinstance(spine_entry, tuple) else spine_entry
        item = book.get_item_with_id(item_id)
        if not item or item.get_type() != item_document_type:
            continue
        href = item.get_name() or ""
        if not href or href in seen:
            continue
        seen.add(href)
        chapter_num += 1
        title = os.path.splitext(os.path.basename(href))[0] or f"Chapter {chapter_num}"
        toc.append({"title": title.replace("_", " ").replace("-", " "), "href": href, "children": []})

    return toc


def _count_toc_entries(toc: list[dict[str, Any]]) -> int:
    count = 0
    for entry in toc:
        count += 1
        count += _count_toc_entries(entry.get("children", []))
    return count