document.py

"""Extract plain text from uploaded documents."""

from __future__ import annotations

import io
from pathlib import Path


SUPPORTED_EXTENSIONS = {".txt", ".md", ".docx", ".pdf"}


def extract_text_from_path(path: str | Path) -> str:
    path = Path(path)
    suffix = path.suffix.lower()
    if suffix not in SUPPORTED_EXTENSIONS:
        raise ValueError(
            f"Unsupported file type '{suffix}'. Use .txt, .md, .docx, or .pdf."
        )
    data = path.read_bytes()
    return extract_text_from_bytes(data, suffix, path.name)


def extract_text_from_bytes(data: bytes, suffix: str, filename: str = "") -> str:
    suffix = suffix.lower()
    if not suffix.startswith("."):
        suffix = f".{suffix}"

    if suffix in {".txt", ".md"}:
        return _decode_text(data)

    if suffix == ".docx":
        return _extract_docx(data)

    if suffix == ".pdf":
        return _extract_pdf(data)

    raise ValueError(f"Unsupported file type for '{filename or suffix}'.")


def _decode_text(data: bytes) -> str:
    for encoding in ("utf-8", "utf-8-sig", "cp1252", "latin-1"):
        try:
            text = data.decode(encoding)
            break
        except UnicodeDecodeError:
            continue
    else:
        text = data.decode("utf-8", errors="replace")
    text = text.strip()
    if not text:
        raise ValueError("The file is empty.")
    return text


def _extract_docx(data: bytes) -> str:
    try:
        from docx import Document
    except ImportError as exc:
        raise RuntimeError("python-docx is required for .docx files.") from exc

    doc = Document(io.BytesIO(data))
    parts: list[str] = []
    for para in doc.paragraphs:
        line = para.text.strip()
        if line:
            parts.append(line)
    for table in doc.tables:
        for row in table.rows:
            cells = [c.text.strip() for c in row.cells if c.text.strip()]
            if cells:
                parts.append(" | ".join(cells))
    text = "\n\n".join(parts).strip()
    if not text:
        raise ValueError("No readable text found in the Word document.")
    return text


def _extract_pdf(data: bytes) -> str:
    try:
        from pypdf import PdfReader
    except ImportError as exc:
        raise RuntimeError("pypdf is required for .pdf files.") from exc

    reader = PdfReader(io.BytesIO(data))
    parts: list[str] = []
    for page in reader.pages:
        page_text = (page.extract_text() or "").strip()
        if page_text:
            parts.append(page_text)
    text = "\n\n".join(parts).strip()
    if not text:
        raise ValueError(
            "No readable text found in the PDF. Scanned/image PDFs are not supported."
        )
    return text


def word_count(text: str) -> int:
    return len(text.split())