src/flashcard_generator/extraction.py

from __future__ import annotations

from pathlib import Path

from .text_processing import clean_text


def extract_text_from_pdf(path: str | Path) -> str:
    try:
        import fitz
    except ImportError as exc:
        raise RuntimeError("PyMuPDF is required for PDF extraction. Install pymupdf.") from exc

    document = fitz.open(path)
    pages: list[str] = []
    for page in document:
        pages.append(page.get_text("text"))
    document.close()
    return clean_text("\n".join(pages))


def extract_pdf_pages(path: str | Path) -> list[str]:
    try:
        import fitz
    except ImportError as exc:
        raise RuntimeError("PyMuPDF is required for PDF extraction. Install pymupdf.") from exc

    document = fitz.open(path)
    pages = [clean_text(page.get_text("text")) for page in document]
    document.close()
    return pages


def extract_text_from_txt(path: str | Path) -> str:
    return clean_text(Path(path).read_text(encoding="utf-8", errors="ignore"))


def extract_text_from_image(path: str | Path) -> str:
    try:
        from PIL import Image
        import pytesseract
    except ImportError as exc:
        raise RuntimeError(
            "Image extraction requires pillow and pytesseract. Install both and make sure Tesseract OCR is available."
        ) from exc

    with Image.open(path) as image:
        return clean_text(pytesseract.image_to_string(image))


def extract_text(path: str | Path) -> str:
    suffix = Path(path).suffix.lower()
    if suffix == ".pdf":
        return extract_text_from_pdf(path)
    if suffix in {".txt", ".md"}:
        return extract_text_from_txt(path)
    if suffix in {".png", ".jpg", ".jpeg", ".webp", ".tiff", ".bmp"}:
        return extract_text_from_image(path)
    raise ValueError(f"Unsupported file type: {suffix or 'unknown'}")