dotfiles/python/ebook_search/ingest.py

"""EPUB ingestion into Richie DB."""

from __future__ import annotations

import asyncio
import hashlib
import logging
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING

import tiktoken
from sqlalchemy import or_, select

from python.ebook_search.epub_parse import parse_epub
from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book
from python.orm.richie import EbookChapter, EbookChunk, EbookSource

logger = logging.getLogger(__name__)
DEFAULT_CHUNK_TOKENS = 700
DEFAULT_CHUNK_OVERLAP = 100

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

    from python.ebook_search.config import EbookSearchConfig
    from python.ebook_search.epub_parse import ParsedChapter


@dataclass(frozen=True)
class TextChunk:
    """A token-bounded chunk of text."""

    text: str
    token_start: int
    token_count: int


def chunk_text(
    text: str,
    *,
    chunk_tokens: int = DEFAULT_CHUNK_TOKENS,
    overlap_tokens: int = DEFAULT_CHUNK_OVERLAP,
) -> list[TextChunk]:
    """Split text into overlapping token chunks."""
    if chunk_tokens <= 0:
        msg = "chunk_tokens must be positive"
        raise ValueError(msg)
    if overlap_tokens < 0 or overlap_tokens >= chunk_tokens:
        msg = "overlap_tokens must be non-negative and smaller than chunk_tokens"
        raise ValueError(msg)

    encoding = tiktoken.get_encoding("cl100k_base")
    tokens = encoding.encode(text)
    if not tokens:
        return []

    chunks: list[TextChunk] = []
    step = chunk_tokens - overlap_tokens
    for start in range(0, len(tokens), step):
        chunk = tokens[start : start + chunk_tokens]
        if not chunk:
            continue
        chunks.append(
            TextChunk(
                text=encoding.decode(chunk).strip(),
                token_start=start,
                token_count=len(chunk),
            )
        )
        if start + chunk_tokens >= len(tokens):
            break
    return [chunk for chunk in chunks if chunk.text]


def find_library_epubs(library_path: str) -> tuple[Path, list[Path] | None]:
    """Resolve one configured library path and collect its EPUB files (blocking filesystem walk).

    Returns:
        tuple[Path, list[Path] | None]: The expanded path and its EPUB files, or ``None`` when
        the path is neither an EPUB file nor a directory.
    """
    path = Path(library_path).expanduser()
    if path.is_file() and path.suffix.lower() == ".epub":
        return path, [path]
    if path.is_dir():
        return path, sorted(path.rglob("*.epub"))
    return path, None


async def ingest_configured_paths(session: AsyncSession, config: EbookSearchConfig) -> int:
    """Ingest every EPUB found under configured library paths."""
    count = 0
    for library_path in config.library_paths:
        path, epub_paths = await asyncio.to_thread(find_library_epubs, library_path)
        logger.info(f"ebook_ingest_path_start {path=}")
        if epub_paths is None:
            logger.warning(f"ebook_ingest_path_missing {path=}")
            continue
        for epub_path in epub_paths:
            count += int(await ingest_file(session, epub_path, config))
    logger.info(f"ebook_ingest_paths_complete {count=} configured_paths={len(config.library_paths)}")
    return count


def resolve_ingest_path(path: Path) -> Path:
    """Expand and resolve an ingest path (blocking filesystem call)."""
    return path.expanduser().resolve()


async def ingest_file(session: AsyncSession, path: Path, config: EbookSearchConfig) -> bool:
    """Ingest one EPUB file. Return True when the database changed."""
    try:
        resolved_path = await asyncio.to_thread(resolve_ingest_path, path)
        logger.info(f"ebook_ingest_file_start {resolved_path=}")
        file_hash = await asyncio.to_thread(sha256_file, resolved_path)
        existing = await find_existing_source(session, resolved_path, file_hash)
        if existing is not None and existing.file_sha256 == file_hash:
            stat = resolved_path.stat()
            existing.file_path = str(resolved_path)
            existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC)
            existing.file_size = stat.st_size
            await session.flush()
            logger.info(f"ebook_ingest_file_unchanged {existing.id=} {resolved_path=}")
            return False
        if existing is not None:
            logger.info(f"ebook_ingest_file_replacing {existing.id=} {resolved_path=}")
            await session.delete(existing)
            await session.flush()

        stat = resolved_path.stat()
        parsed = await asyncio.to_thread(parse_epub, resolved_path)
        source = EbookSource(
            title=parsed.title,
            author=parsed.author,
            language=parsed.language,
            publisher=parsed.publisher,
            identifier=parsed.identifier,
            file_path=str(resolved_path),
            file_sha256=file_hash,
            file_mtime=datetime.fromtimestamp(stat.st_mtime, tz=UTC),
            file_size=stat.st_size,
        )
        session.add(source)
        await session.flush()

        chunk_index = 0
        for spine_index, parsed_chapter in enumerate(parsed.chapters):
            chapter = EbookChapter(
                source_id=source.id,
                spine_index=spine_index,
                title=parsed_chapter.title,
                href=parsed_chapter.href,
            )
            session.add(chapter)
            await session.flush()
            chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index, config)

        await session.commit()
        mention_count = await index_chunk_phrase_mentions_for_book(session, source.id, config)
        logger.info(
            f"ebook_ingest_file_complete {source.id=} {resolved_path=} chapters={len(parsed.chapters)} {chunk_index=} "
            f"{mention_count=}"
        )
    except Exception:
        logger.exception(f"ebook_ingest_file_error {path=}")
        return False
    else:
        return True


async def find_existing_source(session: AsyncSession, path: Path, file_hash: str) -> EbookSource | None:
    """Find an existing source by canonical path or file hash."""
    return await session.scalar(
        select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash))
    )


def add_chapter_chunks(
    session: AsyncSession,
    source: EbookSource,
    chapter: EbookChapter,
    parsed_chapter: ParsedChapter,
    chunk_index: int,
    config: EbookSearchConfig,
) -> int:
    """Add chunk rows for one parsed chapter and return the next chunk index."""
    page_label = parsed_chapter.page_labels[0] if parsed_chapter.page_labels else None
    for text_chunk in chunk_text(
        parsed_chapter.text,
        chunk_tokens=config.chunk_tokens,
        overlap_tokens=config.chunk_overlap,
    ):
        session.add(
            EbookChunk(
                source_id=source.id,
                chapter_id=chapter.id,
                chunk_index=chunk_index,
                text=text_chunk.text,
                token_start=text_chunk.token_start,
                token_count=text_chunk.token_count,
                page_label=page_label,
                content_sha256=hashlib.sha256(text_chunk.text.encode()).hexdigest(),
                search_text=f"{source.title} {source.author or ''} {chapter.title or ''} {text_chunk.text}",
            )
        )
        chunk_index += 1
    return chunk_index


def sha256_file(path: Path) -> str:
    """Calculate the SHA-256 digest for a file."""
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for block in iter(lambda: file.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()