ObsiGate/backend/pdf_reader.py

"""PDF text extraction for ObsiGate — pypdf backend, pymupdf optional fallback."""
from __future__ import annotations

import logging
from pathlib import Path

logger = logging.getLogger(__name__)

PDF_READER: str = "pypdf"
PdfReader = None  # type: ignore
try:
    import fitz  # pymupdf
    PDF_READER = "pymupdf"
    logger.info("PDF reader: pymupdf (high performance)")
except ImportError:
    try:
        from pypdf import PdfReader  # type: ignore  # noqa: F811
        logger.info("PDF reader: pypdf (pure Python)")
    except ImportError:
        logger.warning("No PDF reader available — install pypdf or pymupdf")


def extract_pdf_text(file_path: Path, max_chars: int = 100000) -> str:
    """Extract text from a PDF file. Returns empty string on failure."""
    if PdfReader is None and PDF_READER == "pypdf":
        return ""
    try:
        if PDF_READER == "pymupdf":
            return _extract_pymupdf(file_path, max_chars)
        else:
            return _extract_pypdf(file_path, max_chars)
    except Exception as e:
        logger.warning("Failed to extract PDF text from %s: %s", file_path, e)
        return ""


def extract_pdf_metadata(file_path: Path) -> dict:
    """Extract metadata from a PDF file."""
    info = {"pages": 0, "title": "", "author": ""}
    try:
        if PDF_READER == "pymupdf":
            doc = fitz.open(str(file_path))
            info["pages"] = doc.page_count
            meta = doc.metadata or {}
            info["title"] = meta.get("title", "")
            info["author"] = meta.get("author", "")
            doc.close()
        else:
            reader = PdfReader(str(file_path))
            info["pages"] = len(reader.pages)
            meta = reader.metadata or {}
            if meta:
                info["title"] = str(meta.get("/Title", ""))
                info["author"] = str(meta.get("/Author", ""))
    except Exception as e:
        logger.warning("Failed to extract PDF metadata from %s: %s", file_path, e)
    return info


def extract_pdf_toc(file_path: Path) -> list[dict]:
    """Extract table of contents (bookmarks/outline) from a PDF file.
    
    Returns a list of {title, page, level} dicts, or empty list on failure.
    """
    toc: list[dict] = []
    try:
        if PDF_READER == "pymupdf":
            doc = fitz.open(str(file_path))
            raw = doc.get_toc(simple=False)
            for item in raw:
                toc.append({
                    "title": str(item[1]),
                    "page": item[2],
                    "level": item[0],
                })
            doc.close()
        elif PdfReader is not None:
            reader = PdfReader(str(file_path))
            outline = reader.outline
            if outline:
                def _flatten(items, level=1):
                    for item in items:
                        if isinstance(item, list):
                            _flatten(item, level + 1)
                        elif hasattr(item, 'title') and hasattr(item, 'page'):
                            page_num = reader.get_page_number(item.page) + 1 if hasattr(item, 'page') and item.page else 1
                            toc.append({
                                "title": str(item.title),
                                "page": page_num,
                                "level": level,
                            })
                _flatten(outline)
    except Exception as e:
        logger.warning("Failed to extract PDF TOC from %s: %s", file_path, e)
    return toc


def _extract_pymupdf(file_path: Path, max_chars: int) -> str:
    doc = fitz.open(str(file_path))
    parts = []
    total = 0
    for page in doc:
        text = page.get_text()
        if text:
            if total + len(text) > max_chars:
                remaining = max_chars - total
                if remaining > 0:
                    parts.append(text[:remaining])
                break
            parts.append(text)
            total += len(text)
    doc.close()
    return "\f".join(parts)


def _extract_pypdf(file_path: Path, max_chars: int) -> str:
    reader = PdfReader(str(file_path))
    parts = []
    total = 0
    for page in reader.pages:
        text = page.extract_text() or ""
        if text:
            if total + len(text) > max_chars:
                remaining = max_chars - total
                if remaining > 0:
                    parts.append(text[:remaining])
                break
            parts.append(text)
            total += len(text)
    return "\f".join(parts)