potato/format_handlers/pdf_handler.py

"""
PDF Format Handler

Extracts text and layout information from PDF files using pdfplumber.
Supports text extraction with character-level position mapping.

Usage:
    from potato.format_handlers.pdf_handler import PDFHandler

    handler = PDFHandler()
    output = handler.extract("document.pdf", {
        "extraction_mode": "text",  # or "layout"
        "max_pages": 10,
    })

    # Access extracted content
    text = output.text
    html = output.rendered_html
    coords = output.coordinate_map
"""

from typing import Dict, List, Any, Optional
from pathlib import Path
import html
import logging

from .base import BaseFormatHandler, FormatOutput
from .coordinate_mapping import CoordinateMapper, PDFCoordinate

logger = logging.getLogger(__name__)

# Check if pdfplumber is available
try:
    import pdfplumber
    PDFPLUMBER_AVAILABLE = True
except ImportError:
    PDFPLUMBER_AVAILABLE = False
    pdfplumber = None


class PDFHandler(BaseFormatHandler):
    """
    Handler for PDF documents.

    Uses pdfplumber for text extraction with position information.
    Generates HTML representation suitable for span annotation.
    """

    format_name = "pdf"
    supported_extensions = [".pdf"]
    description = "PDF document text extraction with page/position mapping"
    requires_dependencies = ["pdfplumber"]

    def get_default_options(self) -> Dict[str, Any]:
        """Get default extraction options."""
        return {
            "extraction_mode": "text",  # "text" or "layout"
            "preserve_layout": False,
            "max_pages": None,
            "include_page_breaks": True,
            "page_separator": "\n\n--- Page {page} ---\n\n",
            "extract_tables": False,
            "x_tolerance": 3,  # Horizontal tolerance for word grouping
            "y_tolerance": 3,  # Vertical tolerance for line grouping
        }

    def extract(
        self,
        file_path: str,
        options: Optional[Dict[str, Any]] = None
    ) -> FormatOutput:
        """
        Extract text and layout from a PDF file.

        Args:
            file_path: Path to the PDF file
            options: Extraction options:
                - extraction_mode: "text" (plain) or "layout" (preserve layout)
                - max_pages: Maximum pages to process (None for all)
                - include_page_breaks: Include page separators in text
                - page_separator: Format string for page breaks ({page} replaced)
                - extract_tables: Also extract table structures

        Returns:
            FormatOutput with extracted text, HTML, and coordinate mappings
        """
        if not PDFPLUMBER_AVAILABLE:
            raise ImportError(
                "pdfplumber is required for PDF extraction. "
                "Install with: pip install pdfplumber"
            )

        opts = self.merge_options(options)
        mapper = CoordinateMapper()

        text_parts = []
        html_parts = []
        current_offset = 0

        metadata = {
            "format": "pdf",
            "pages": [],
            "total_pages": 0,
            "source_file": str(file_path),
        }

        html_parts.append('<div class="pdf-content">')

        with pdfplumber.open(file_path) as pdf:
            metadata["total_pages"] = len(pdf.pages)
            max_pages = opts.get("max_pages") or len(pdf.pages)

            for page_num, page in enumerate(pdf.pages[:max_pages], start=1):
                page_text, page_html, page_coords = self._extract_page(
                    page, page_num, opts, current_offset
                )

                # Add page coordinates to mapper
                for coord_info in page_coords:
                    mapper.add_mapping(
                        coord_info["start"],
                        coord_info["end"],
                        PDFCoordinate(
                            page=page_num,
                            bbox=coord_info.get("bbox", []),
                            line=coord_info.get("line"),
                        )
                    )

                # Add page separator
                if page_num > 1 and opts.get("include_page_breaks"):
                    separator = opts["page_separator"].format(page=page_num)
                    text_parts.append(separator)
                    current_offset += len(separator)

                text_parts.append(page_text)
                html_parts.append(page_html)
                current_offset += len(page_text)

                # Page metadata
                page_meta = {
                    "page_number": page_num,
                    "width": float(page.width),
                    "height": float(page.height),
                    "char_count": len(page_text),
                }
                metadata["pages"].append(page_meta)

        html_parts.append('</div>')

        full_text = "".join(text_parts)
        full_html = "\n".join(html_parts)

        # Create output with coordinate lookup function
        coord_dict = mapper.to_dict()
        coord_dict["get_coords_for_range"] = mapper.get_coords_for_range

        return FormatOutput(
            text=full_text,
            rendered_html=full_html,
            coordinate_map=coord_dict,
            metadata=metadata,
            format_name=self.format_name,
            source_path=str(file_path),
        )

    def _extract_page(
        self,
        page,
        page_num: int,
        opts: Dict[str, Any],
        base_offset: int
    ) -> tuple:
        """
        Extract text and HTML from a single page.

        Returns:
            Tuple of (text, html, coordinate_mappings)
        """
        extraction_mode = opts.get("extraction_mode", "text")

        if extraction_mode == "layout":
            return self._extract_page_layout(page, page_num, opts, base_offset)
        else:
            return self._extract_page_text(page, page_num, opts, base_offset)

    def _extract_page_text(
        self,
        page,
        page_num: int,
        opts: Dict[str, Any],
        base_offset: int
    ) -> tuple:
        """
        Extract text with word-level coordinate mapping.
        """
        text_parts = []
        html_parts = []
        coords = []
        current_offset = base_offset

        # Extract words with their positions
        words = page.extract_words(
            x_tolerance=opts.get("x_tolerance", 3),
            y_tolerance=opts.get("y_tolerance", 3),
        )

        html_parts.append(f'<div class="pdf-page" data-page="{page_num}">')

        if not words:
            # Fall back to full text extraction if no words found
            text = page.extract_text() or ""
            text_parts.append(text)
            html_parts.append(f'<span class="pdf-text">{html.escape(text)}</span>')

            if text:
                coords.append({
                    "start": current_offset,
                    "end": current_offset + len(text),
                    "bbox": [0, 0, float(page.width), float(page.height)],
                })
        else:
            # Process words with positions
            current_line_top = None
            line_words = []

            for word in words:
                word_top = word["top"]

                # Check if this is a new line
                if current_line_top is None:
                    current_line_top = word_top
                elif abs(word_top - current_line_top) > opts.get("y_tolerance", 3):
                    # Flush current line
                    if line_words:
                        line_text, line_html, line_coords = self._process_line(
                            line_words, current_offset
                        )
                        text_parts.append(line_text)
                        text_parts.append("\n")
                        html_parts.append(line_html)
                        html_parts.append("<br>")
                        coords.extend(line_coords)
                        current_offset += len(line_text) + 1  # +1 for newline

                    line_words = []
                    current_line_top = word_top

                line_words.append(word)

            # Process final line
            if line_words:
                line_text, line_html, line_coords = self._process_line(
                    line_words, current_offset
                )
                text_parts.append(line_text)
                html_parts.append(line_html)
                coords.extend(line_coords)

        html_parts.append('</div>')

        return "".join(text_parts), "\n".join(html_parts), coords

    def _process_line(
        self,
        words: List[Dict],
        base_offset: int
    ) -> tuple:
        """
        Process a line of words into text, HTML, and coordinates.
        """
        text_parts = []
        html_parts = []
        coords = []
        current_offset = base_offset

        for i, word in enumerate(words):
            word_text = word["text"]

            # Add space between words
            if i > 0:
                text_parts.append(" ")
                current_offset += 1

            start = current_offset
            end = start + len(word_text)

            text_parts.append(word_text)
            html_parts.append(
                f'<span class="pdf-word" '
                f'data-start="{start}" '
                f'data-end="{end}">'
                f'{html.escape(word_text)}</span>'
            )

            # Store coordinate mapping
            coords.append({
                "start": start,
                "end": end,
                "bbox": [
                    float(word["x0"]),
                    float(word["top"]),
                    float(word["x1"]),
                    float(word["bottom"]),
                ],
            })

            current_offset = end

        return "".join(text_parts), " ".join(html_parts), coords

    def _extract_page_layout(
        self,
        page,
        page_num: int,
        opts: Dict[str, Any],
        base_offset: int
    ) -> tuple:
        """
        Extract text preserving visual layout.
        """
        # Use extract_text with layout preservation
        text = page.extract_text(layout=True) or ""

        html_parts = []
        html_parts.append(f'<div class="pdf-page pdf-page-layout" data-page="{page_num}">')
        html_parts.append(f'<pre class="pdf-layout-text">{html.escape(text)}</pre>')
        html_parts.append('</div>')

        # For layout mode, we map the entire page
        coords = [{
            "start": base_offset,
            "end": base_offset + len(text),
            "bbox": [0, 0, float(page.width), float(page.height)],
        }]

        return text, "\n".join(html_parts), coords

    def get_page_count(self, file_path: str) -> int:
        """
        Get the number of pages in a PDF.

        Args:
            file_path: Path to the PDF file

        Returns:
            Number of pages
        """
        if not PDFPLUMBER_AVAILABLE:
            raise ImportError("pdfplumber is required")

        with pdfplumber.open(file_path) as pdf:
            return len(pdf.pages)

    def extract_page(
        self,
        file_path: str,
        page_number: int,
        options: Optional[Dict[str, Any]] = None
    ) -> FormatOutput:
        """
        Extract a single page from a PDF.

        Args:
            file_path: Path to the PDF file
            page_number: Page number (1-indexed)
            options: Extraction options

        Returns:
            FormatOutput for the single page
        """
        if not PDFPLUMBER_AVAILABLE:
            raise ImportError("pdfplumber is required")

        opts = self.merge_options(options)
        opts["max_pages"] = page_number  # Process up to this page
        opts["include_page_breaks"] = False

        # Extract only the requested page
        mapper = CoordinateMapper()

        with pdfplumber.open(file_path) as pdf:
            if page_number < 1 or page_number > len(pdf.pages):
                raise ValueError(
                    f"Page {page_number} out of range (1-{len(pdf.pages)})"
                )

            page = pdf.pages[page_number - 1]
            page_text, page_html, page_coords = self._extract_page(
                page, page_number, opts, 0
            )

            for coord_info in page_coords:
                mapper.add_mapping(
                    coord_info["start"],
                    coord_info["end"],
                    PDFCoordinate(
                        page=page_number,
                        bbox=coord_info.get("bbox", []),
                    )
                )

        coord_dict = mapper.to_dict()
        coord_dict["get_coords_for_range"] = mapper.get_coords_for_range

        return FormatOutput(
            text=page_text,
            rendered_html=page_html,
            coordinate_map=coord_dict,
            metadata={
                "format": "pdf",
                "page_number": page_number,
                "total_pages": len(pdf.pages),
            },
            format_name=self.format_name,
            source_path=str(file_path),
        )