src/pdf_loader.py

"""PDF upload validation and text extraction."""

from pathlib import Path

from pypdf import PdfReader

from src.config import DEFAULT_OCR_ENGINE, DEFAULT_OCR_MODE, OCR_ENGINE_NAMES, OCR_MODE_DPI
from src.ocr_service import extract_text_from_pdf_page_with_selected_ocr


def validate_pdf_file(file_name: str) -> None:
    """Raise a helpful error when the uploaded file is not a PDF."""
    if not file_name.lower().endswith(".pdf"):
        raise ValueError(f"{file_name} is not a PDF file. Please upload only PDF files.")


def save_pdf_bytes(file_name: str, file_bytes: bytes, upload_dir: Path) -> Path:
    """Save uploaded PDF bytes into the local uploads folder."""
    validate_pdf_file(file_name)

    if not file_bytes:
        raise ValueError(f"{file_name} is empty. Please upload a valid PDF file.")

    upload_dir.mkdir(parents=True, exist_ok=True)
    safe_file_name = Path(file_name).name
    saved_path = upload_dir / safe_file_name

    with saved_path.open("wb") as pdf_file:
        pdf_file.write(file_bytes)

    return saved_path


def count_pdf_pages(pdf_path: Path) -> int:
    """Return how many pages a PDF has."""
    try:
        reader = PdfReader(str(pdf_path))
    except Exception as error:
        raise ValueError(f"Could not read {pdf_path.name}. The PDF may be damaged.") from error

    return len(reader.pages)


def validate_page_range(start_page: int | None, end_page: int | None, total_pages: int) -> tuple[int, int]:
    """Convert a 1-based page range into safe 0-based bounds."""
    start = start_page or 1
    end = end_page or total_pages

    if start < 1:
        raise ValueError("Start page must be 1 or greater.")

    if end < start:
        raise ValueError("End page must be greater than or equal to start page.")

    if start > total_pages:
        raise ValueError(f"Start page {start} is greater than total pages {total_pages}.")

    end = min(end, total_pages)
    return start - 1, end


def validate_ocr_mode(ocr_mode: str) -> str:
    """Return a supported OCR mode or raise a beginner-friendly error."""
    normalized_mode = (ocr_mode or DEFAULT_OCR_MODE).strip().lower()

    if normalized_mode not in OCR_MODE_DPI:
        allowed_modes = ", ".join(OCR_MODE_DPI)
        raise ValueError(f"Unsupported OCR mode '{ocr_mode}'. Use one of: {allowed_modes}.")

    return normalized_mode


def validate_ocr_engine(ocr_engine: str) -> str:
    """Return a supported OCR engine or raise a helpful error."""
    normalized_engine = (ocr_engine or DEFAULT_OCR_ENGINE).strip().lower()

    if normalized_engine not in OCR_ENGINE_NAMES:
        allowed_engines = ", ".join(sorted(OCR_ENGINE_NAMES))
        raise ValueError(f"Unsupported OCR engine '{ocr_engine}'. Use one of: {allowed_engines}.")

    return normalized_engine


def extract_pages_from_pdf(
    pdf_path: Path,
    start_page: int | None = None,
    end_page: int | None = None,
    ocr_mode: str = DEFAULT_OCR_MODE,
    ocr_engine: str = DEFAULT_OCR_ENGINE,
) -> list[dict]:
    """Extract text page by page, using the selected OCR engine for scanned pages."""
    pages = []
    ocr_mode = validate_ocr_mode(ocr_mode)
    ocr_engine = validate_ocr_engine(ocr_engine)

    try:
        reader = PdfReader(str(pdf_path))
    except Exception as error:
        raise ValueError(f"Could not read {pdf_path.name}. The PDF may be damaged.") from error

    total_pages = len(reader.pages)
    start_index, stop_index = validate_page_range(start_page, end_page, total_pages)

    for page_index in range(start_index, stop_index):
        page = reader.pages[page_index]
        page_text = page.extract_text() or ""
        cleaned_text = " ".join(page_text.split())
        extraction_method = "embedded_text"

        if not cleaned_text:
            cleaned_text, extraction_method = extract_text_from_pdf_page_with_selected_ocr(
                pdf_path=pdf_path,
                page_index=page_index,
                ocr_mode=ocr_mode,
                ocr_engine=ocr_engine,
            )

        if cleaned_text:
            used_ocr = extraction_method != "embedded_text"
            pages.append(
                {
                    "text": cleaned_text,
                    "page_number": page_index + 1,
                    "source_file": pdf_path.name,
                    "extraction_method": extraction_method,
                    "ocr_mode": ocr_mode if used_ocr else "not_used",
                    "ocr_engine": ocr_engine if used_ocr else "not_used",
                    "total_pages": total_pages,
                }
            )

    if not pages:
        raise ValueError(
            f"No readable text was found in {pdf_path.name}. "
            "The selected OCR engine could not read enough text from this PDF. "
            "Try a clearer scan, better lighting, or typed notes."
        )

    return pages


def load_uploaded_pdfs(
    saved_pdf_paths: list[Path],
    start_page: int | None = None,
    end_page: int | None = None,
    ocr_mode: str = DEFAULT_OCR_MODE,
    ocr_engine: str = DEFAULT_OCR_ENGINE,
) -> list[dict]:
    """Extract readable text from all uploaded PDFs."""
    all_pages = []

    for pdf_path in saved_pdf_paths:
        all_pages.extend(
            extract_pages_from_pdf(
                pdf_path=pdf_path,
                start_page=start_page,
                end_page=end_page,
                ocr_mode=ocr_mode,
                ocr_engine=ocr_engine,
            )
        )

    if not all_pages:
        raise ValueError("No readable text was found in the uploaded PDFs.")

    return all_pages