replace_page.py

"""Replace a page in a PDF file with an image."""

from __future__ import annotations

import argparse
import io
import logging
import os
from typing import Optional

from PIL import Image  # pylint: disable=import-error
from PyPDF2 import (
    PdfReader,
    PdfWriter,
    PageObject,
    Transformation,
)  # pylint: disable=import-error

__version__ = "0.1.0"

logger = logging.getLogger(__name__)


class ImageToPdfError(RuntimeError):
    """Raised when image-to-PDF conversion fails."""


class PdfPageError(RuntimeError):
    """Raised when PDF page operations fail."""


def image_to_pdf_bytes(image_path: str, page_width: float, page_height: float) -> bytes:
    """Convert an image file to PDF bytes, scaled to fit the specified page dimensions."""
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"image not found: {image_path}")

    try:
        img = Image.open(image_path)
        if img.mode == "RGBA":
            img = img.convert("RGB")

        # Calculate scaling to fit the page while maintaining aspect ratio
        img_width, img_height = img.size
        aspect_ratio = img_width / img_height
        page_aspect = page_width / page_height

        if aspect_ratio > page_aspect:
            # Image is wider relative to page
            new_width = page_width
            new_height = page_width / aspect_ratio
        else:
            # Image is taller relative to page
            new_height = page_height
            new_width = page_height * aspect_ratio

        # Resize image
        img_resized = img.resize(
            (int(new_width), int(new_height)), Image.Resampling.LANCZOS
        )

        # Create a new image with the page dimensions and place the resized image centered
        bg = Image.new("RGB", (int(page_width), int(page_height)), "white")
        x_offset = int((page_width - new_width) / 2)
        y_offset = int((page_height - new_height) / 2)
        bg.paste(img_resized, (x_offset, y_offset))

        # Convert to PDF using PIL's built-in method
        output_io = io.BytesIO()
        bg.save(output_io, format="PDF")
        output_io.seek(0)

        return output_io.getvalue()

    except (OSError, ValueError) as exc:
        msg = f"failed to convert image to PDF: {exc}"
        raise ImageToPdfError(msg) from exc


def replace_page(
    pdf_path: str,
    replacement_path: str,
    page_number: int,
    output_path: str,
    overwrite: bool = False,
) -> None:
    """Replace a specific page in a PDF with an image or another PDF page.

    Args:
        pdf_path: Path to the input PDF file.
        replacement_path: Path to the image file or single-page PDF to use as replacement.
        page_number: 1-indexed page number to replace.
        output_path: Path to save the modified PDF.
        overwrite: If True, overwrite the output file if it exists.

    Raises:
        FileNotFoundError: If input PDF or replacement file not found.
        PdfPageError: If page number is invalid or operation fails.
        ValueError: If output file exists and overwrite is False, or if replacement PDF has multiple pages.
    """
    if not os.path.isfile(pdf_path):
        raise FileNotFoundError(f"PDF file not found: {pdf_path}")

    if not os.path.isfile(replacement_path):
        raise FileNotFoundError(f"replacement file not found: {replacement_path}")

    if os.path.exists(output_path) and not overwrite:
        raise ValueError(
            f"output file {output_path} already exists; use --overwrite to replace"
        )

    if page_number < 1:
        raise PdfPageError("page number must be >= 1 (1-indexed)")

    # Determine if replacement is an image or PDF based on file extension
    _, ext = os.path.splitext(replacement_path.lower())
    is_image = ext in (".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".gif")
    is_pdf = ext == ".pdf"

    if not is_image and not is_pdf:
        raise ValueError(
            f"replacement file must be an image (.png, .jpg, .jpeg, .bmp, .tiff, .tif, .gif) "
            f"or PDF (.pdf), got: {ext}"
        )

    try:
        # Read the PDF into memory to avoid file handle issues
        with open(pdf_path, "rb") as f:
            pdf_data = f.read()

        reader = PdfReader(io.BytesIO(pdf_data))
        total_pages = len(reader.pages)

        if page_number > total_pages:
            raise PdfPageError(
                f"page number {page_number} exceeds total pages ({total_pages})"
            )

        # Get the dimensions of the page being replaced
        page_to_replace = reader.pages[page_number - 1]
        mediabox = page_to_replace.mediabox  # pylint: disable=no-member
        page_width = float(mediabox.width)
        page_height = float(mediabox.height)

        logger.debug(
            "page %d dimensions: %.2f x %.2f", page_number, page_width, page_height
        )

        # Get the replacement page
        if is_image:
            # Convert image to PDF bytes and read as PDF
            image_pdf_bytes = image_to_pdf_bytes(
                replacement_path, page_width, page_height
            )
            replacement_reader = PdfReader(io.BytesIO(image_pdf_bytes))
            replacement_page = replacement_reader.pages[0]
        else:  # is_pdf
            # Read the replacement PDF directly
            with open(replacement_path, "rb") as f:
                replacement_data = f.read()
            replacement_reader = PdfReader(io.BytesIO(replacement_data))

            if len(replacement_reader.pages) != 1:
                msg = (
                    "replacement PDF must contain exactly 1 page, got "
                    f"{len(replacement_reader.pages)}"
                )
                raise ValueError(msg)

            replacement_page = replacement_reader.pages[0]

            # Create a new page with the target dimensions
            new_page = PageObject.create_blank_page(
                width=page_width, height=page_height
            )

            # Get the mediabox dimensions of the replacement page
            replacement_mediabox = (
                replacement_page.mediabox
            )  # pylint: disable=no-member
            replacement_width = float(replacement_mediabox.width)
            replacement_height = float(replacement_mediabox.height)

            # Calculate scale factors
            scale_x = page_width / replacement_width
            scale_y = page_height / replacement_height

            # Use uniform scaling to maintain aspect ratio (use the smaller scale factor)
            scale = min(scale_x, scale_y)

            logger.debug(
                "scaling replacement PDF page from %.2f x %.2f to %.2f x %.2f (scale: %.2f)",
                replacement_width,
                replacement_height,
                replacement_width * scale,
                replacement_height * scale,
                scale,
            )

            # Calculate position to center the scaled page
            scaled_width = replacement_width * scale
            scaled_height = replacement_height * scale
            offset_x = (page_width - scaled_width) / 2
            offset_y = (page_height - scaled_height) / 2

            # Apply the transformation to the replacement page BEFORE merging
            # Create a transformation that scales and then translates
            transformation = Transformation()
            transformation = transformation.scale(sx=scale, sy=scale)
            transformation = transformation.translate(tx=offset_x, ty=offset_y)

            replacement_page.add_transformation(
                transformation
            )  # pylint: disable=no-member

            # Now merge the transformed page onto the new blank page
            new_page.merge_page(replacement_page)

            replacement_page = new_page

        # Create a new writer and add all pages, replacing the specified page
        writer = PdfWriter()
        for idx, page in enumerate(reader.pages, start=1):
            if idx == page_number:
                writer.add_page(replacement_page)
                logger.debug(
                    "replaced page %d with %s",
                    page_number,
                    "image" if is_image else "PDF page",
                )
            else:
                writer.add_page(page)

        # Write the modified PDF
        with open(output_path, "wb") as f:
            writer.write(f)

        logger.info("successfully replaced page %d in %s", page_number, output_path)

    except (FileNotFoundError, PdfPageError, ValueError):
        raise
    except Exception as exc:
        raise PdfPageError(f"failed to replace page: {exc}") from exc


def _configure_logging(args: argparse.Namespace) -> None:
    """Initialize logging with the desired verbosity."""
    level = (
        logging.WARNING
        if args.quiet
        else getattr(logging, args.log_level.upper(), logging.INFO)
    )
    logging.basicConfig(level=level, format="%(asctime)s %(levelname)s %(message)s")


def main(argv: Optional[list[str]] = None) -> int:
    """CLI entry point for replacing PDF pages with images or other PDFs."""
    parser = argparse.ArgumentParser(
        description="Replace a page in a PDF file with an image or another PDF page."
    )
    parser.add_argument("pdf", help="Path to the PDF file")
    parser.add_argument(
        "replacement", help="Path to the image or PDF file to use as replacement"
    )
    parser.add_argument(
        "-p",
        "--page",
        type=int,
        default=1,
        help="Page number to replace (1-indexed, default: 1)",
    )
    parser.add_argument(
        "-o",
        "--output",
        help="Path to the output PDF file (defaults to <input>_replaced.pdf)",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite the output file if it exists",
    )
    parser.add_argument(
        "--quiet",
        action="store_true",
        help="Suppress informational logging",
    )
    parser.add_argument(
        "--log-level",
        default="INFO",
        help="Logging level (DEBUG, INFO, WARNING, ERROR)",
    )
    parser.add_argument(
        "--version",
        action="version",
        version=f"replace-page {__version__}",
    )

    args = parser.parse_args(argv)
    _configure_logging(args)

    pdf_path = os.path.abspath(args.pdf)
    replacement_path = os.path.abspath(args.replacement)

    if args.output:
        output_path = os.path.abspath(args.output)
    else:
        base, ext = os.path.splitext(pdf_path)
        output_path = f"{base}_replaced{ext}"

    try:
        replace_page(
            pdf_path,
            replacement_path,
            args.page,
            output_path,
            overwrite=args.overwrite,
        )
        logger.info("PDF page replacement complete: %s", output_path)
        return 0
    except (FileNotFoundError, ValueError, PdfPageError, ImageToPdfError) as exc:
        logger.error("%s", exc)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())