pdf_editor.py

from __future__ import annotations

import io
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import pdfplumber
from PyPDF2 import PdfReader, PdfWriter
from PyPDF2.generic import BooleanObject, NameObject
from reportlab.lib.utils import ImageReader
from reportlab.pdfgen import canvas


class PDFEditor:
    """Utility class for common PDF editing operations."""

    def __init__(self, file_path: Optional[str] = None) -> None:
        self.file_path: Optional[str] = None
        self.reader: Optional[PdfReader] = None
        self.writer: PdfWriter = PdfWriter()

        if file_path:
            self.load(file_path)

    def load(self, file_path: str) -> None:
        """Load a PDF and initialize an in-memory writer copy."""
        self.file_path = str(file_path)
        self.reader = PdfReader(self.file_path)
        self.writer = PdfWriter()

        for page in self.reader.pages:
            self.writer.add_page(page)

        if self.reader.metadata:
            self.writer.add_metadata(dict(self.reader.metadata))

    @staticmethod
    def merge_pdfs(file_paths: Sequence[str], output_path: str) -> None:
        """Merge multiple PDFs into one output file."""
        if len(file_paths) < 2:
            raise ValueError("Need at least two files to merge")

        writer = PdfWriter()
        for path in file_paths:
            reader = PdfReader(path)
            for page in reader.pages:
                writer.add_page(page)

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as handle:
            writer.write(handle)

    def split_pdf(self, ranges: Sequence[Sequence[int]], output_prefix: str) -> List[str]:
        """Split a PDF into separate files by 1-based inclusive page ranges."""
        self._require_loaded()

        if not ranges:
            raise ValueError("At least one range is required")

        num_pages = len(self.reader.pages)
        output_files: List[str] = []

        for idx, rng in enumerate(ranges, start=1):
            if len(rng) != 2:
                raise ValueError(f"Invalid range format at item {idx}: {rng}")

            start, end = int(rng[0]), int(rng[1])
            if start < 1 or end < 1 or end < start:
                raise ValueError(f"Invalid range values: {start}-{end}")
            if end > num_pages:
                raise ValueError(f"Range {start}-{end} exceeds total pages ({num_pages})")

            writer = PdfWriter()
            for page_index in range(start - 1, end):
                writer.add_page(self.reader.pages[page_index])

            output_file = f"{output_prefix}_{idx}.pdf"
            Path(output_file).parent.mkdir(parents=True, exist_ok=True)
            with open(output_file, "wb") as handle:
                writer.write(handle)
            output_files.append(output_file)

        return output_files

    def extract_pages(self, pages: Sequence[int], output_path: str) -> None:
        """Extract specific 1-based page numbers into a new PDF."""
        self._require_loaded()

        if not pages:
            raise ValueError("At least one page is required")

        num_pages = len(self.writer.pages)
        writer = PdfWriter()

        for page_num in pages:
            page_idx = int(page_num) - 1
            if page_idx < 0 or page_idx >= num_pages:
                raise ValueError(f"Page number out of range: {page_num}")
            writer.add_page(self.writer.pages[page_idx])

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as handle:
            writer.write(handle)

    def rotate_page(self, page_num: int, angle: int) -> None:
        """Rotate a single page by angle (must be multiple of 90)."""
        self._require_loaded()

        if int(angle) % 90 != 0:
            raise ValueError("Angle must be a multiple of 90")

        page_idx = int(page_num) - 1
        self._validate_page_index(page_idx)
        page = self.writer.pages[page_idx]

        if hasattr(page, "rotate"):
            page.rotate(int(angle))
        elif int(angle) >= 0:
            page.rotate_clockwise(int(angle))
        else:
            page.rotate_counter_clockwise(abs(int(angle)))

    def add_text(
        self,
        page_num: int,
        text: str,
        x: float,
        y: float,
        font_size: int = 12,
    ) -> None:
        """Overlay text on a page at the given coordinates."""
        self._require_loaded()

        page_idx = int(page_num) - 1
        self._validate_page_index(page_idx)

        target_page = self.writer.pages[page_idx]
        page_width = float(target_page.mediabox.width)
        page_height = float(target_page.mediabox.height)

        packet = io.BytesIO()
        cv = canvas.Canvas(packet, pagesize=(page_width, page_height))
        cv.setFont("Helvetica", int(font_size))
        cv.drawString(float(x), float(y), text)
        cv.save()
        packet.seek(0)

        overlay_page = PdfReader(packet).pages[0]
        target_page.merge_page(overlay_page)

    def add_image(
        self,
        page_num: int,
        image_path: str,
        x: float,
        y: float,
        width: Optional[float] = None,
        height: Optional[float] = None,
    ) -> None:
        """Overlay an image on a page at the given coordinates."""
        self._require_loaded()

        page_idx = int(page_num) - 1
        self._validate_page_index(page_idx)

        target_page = self.writer.pages[page_idx]
        page_width = float(target_page.mediabox.width)
        page_height = float(target_page.mediabox.height)

        image_reader = ImageReader(image_path)
        native_width, native_height = image_reader.getSize()
        draw_width = float(width) if width else float(native_width)
        draw_height = float(height) if height else float(native_height)

        packet = io.BytesIO()
        cv = canvas.Canvas(packet, pagesize=(page_width, page_height))
        cv.drawImage(
            image_reader,
            float(x),
            float(y),
            width=draw_width,
            height=draw_height,
            preserveAspectRatio=False,
            mask="auto",
        )
        cv.save()
        packet.seek(0)

        overlay_page = PdfReader(packet).pages[0]
        target_page.merge_page(overlay_page)

    def extract_text(self, pages: Optional[Sequence[int]] = None) -> Dict[str, str]:
        """Extract text from selected pages (or all pages when omitted)."""
        if not self.file_path:
            raise ValueError("No file loaded")

        text_by_page: Dict[str, str] = {}
        with pdfplumber.open(self.file_path) as pdf:
            total_pages = len(pdf.pages)

            if pages:
                selected = [int(p) for p in pages]
            else:
                selected = list(range(1, total_pages + 1))

            for page_num in selected:
                page_idx = page_num - 1
                if page_idx < 0 or page_idx >= total_pages:
                    raise ValueError(f"Page number out of range: {page_num}")

                extracted = pdf.pages[page_idx].extract_text() or ""
                text_by_page[str(page_num)] = extracted

        return text_by_page

    def fill_form(self, form_data: Dict[str, str], output_path: str) -> None:
        """Fill form fields and write the result to output_path."""
        self._require_loaded()

        writer = PdfWriter()
        for page in self.reader.pages:
            writer.add_page(page)

        fields = self.reader.get_fields() or {}
        if fields:
            for page in writer.pages:
                try:
                    writer.update_page_form_field_values(page, form_data)
                except Exception:
                    # Some pages in mixed documents may not carry form widgets.
                    continue

        root = self.reader.trailer.get("/Root")
        if hasattr(root, "get_object"):
            root = root.get_object()

        if isinstance(root, dict) and "/AcroForm" in root:
            acro_form = root["/AcroForm"]
            writer._root_object.update({NameObject("/AcroForm"): acro_form})
            form_obj = writer._root_object["/AcroForm"]
            if hasattr(form_obj, "get_object"):
                form_obj = form_obj.get_object()
            if hasattr(form_obj, "update"):
                form_obj.update({NameObject("/NeedAppearances"): BooleanObject(True)})

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as handle:
            writer.write(handle)

    def get_form_fields(self) -> Dict[str, Dict[str, str]]:
        """Return form fields keyed by field name."""
        self._require_loaded()
        fields = self.reader.get_fields() or {}

        normalized: Dict[str, Dict[str, str]] = {}
        for name, data in fields.items():
            normalized[name] = {
                "value": str(data.get("/V", "")),
                "type": str(data.get("/FT", "")),
            }
        return normalized

    def encrypt(self, password: str, output_path: str) -> None:
        """Encrypt current PDF with a password."""
        self._require_loaded()

        writer = PdfWriter()
        for page in self.writer.pages:
            writer.add_page(page)

        if self.reader and self.reader.metadata:
            writer.add_metadata(dict(self.reader.metadata))

        writer.encrypt(password)

        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as handle:
            writer.write(handle)

    def get_metadata(self) -> Dict[str, str]:
        """Get metadata for the loaded document."""
        self._require_loaded()

        metadata = self.reader.metadata or {}
        clean_meta: Dict[str, str] = {}
        for key, value in dict(metadata).items():
            clean_meta[str(key)] = "" if value is None else str(value)
        return clean_meta

    def set_metadata(self, metadata: Dict[str, str]) -> None:
        """Set metadata on the in-memory writer."""
        self._require_loaded()

        if not metadata:
            raise ValueError("Metadata dictionary cannot be empty")

        clean_meta = {}
        for key, value in metadata.items():
            meta_key = key if str(key).startswith("/") else f"/{key}"
            clean_meta[str(meta_key)] = str(value)

        self.writer.add_metadata(clean_meta)

    def save(self, output_path: str) -> None:
        """Save current in-memory writer state."""
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as handle:
            self.writer.write(handle)

    def _require_loaded(self) -> None:
        if not self.reader:
            raise ValueError("No PDF loaded. Initialize with a file path or call load().")

    def _validate_page_index(self, page_idx: int) -> None:
        num_pages = len(self.writer.pages)
        if page_idx < 0 or page_idx >= num_pages:
            raise ValueError(f"Page number out of range: {page_idx + 1}")