mcp-pdf-tools/src/mcp_pdf/mixins/table_extraction.py

"""
Table Extraction Mixin - PDF table detection and extraction capabilities
"""

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

# PDF processing libraries
import camelot
import tabula
import pdfplumber
import pandas as pd

from .base import MCPMixin, mcp_tool
from ..security import validate_pdf_path, parse_pages_parameter, sanitize_error_message

logger = logging.getLogger(__name__)


class TableExtractionMixin(MCPMixin):
    """
    Handles all PDF table extraction operations with intelligent fallbacks.

    Tools provided:
    - extract_tables: Multi-method table extraction with automatic fallbacks
    """

    def get_mixin_name(self) -> str:
        return "TableExtraction"

    def get_required_permissions(self) -> List[str]:
        return ["read_files", "table_processing"]

    def _setup(self):
        """Initialize table extraction specific configuration"""
        self.table_accuracy_threshold = 0.8
        self.max_tables_per_page = 10

    @mcp_tool(
        name="extract_tables",
        description="Extract tables from PDF with automatic method selection and intelligent fallbacks"
    )
    async def extract_tables(
        self,
        pdf_path: str,
        pages: Optional[str] = None,
        method: str = "auto",
        table_format: str = "json"
    ) -> Dict[str, Any]:
        """
        Extract tables from PDF using various methods with automatic fallbacks.

        Args:
            pdf_path: Path to PDF file or URL
            pages: Page specification (e.g., "1-5,10,15-20" or "all")
            method: Extraction method ("auto", "camelot", "tabula", "pdfplumber")
            table_format: Output format ("json", "csv", "markdown")

        Returns:
            Dictionary containing extracted tables and metadata
        """
        start_time = time.time()

        try:
            # Validate inputs using centralized security functions
            path = await validate_pdf_path(pdf_path)
            parsed_pages = parse_pages_parameter(pages)

            all_tables = []
            methods_tried = []

            # Auto method: try methods in order until we find tables
            if method == "auto":
                for try_method in ["camelot", "pdfplumber", "tabula"]:
                    methods_tried.append(try_method)

                    if try_method == "camelot":
                        tables = self._extract_tables_camelot(path, parsed_pages)
                    elif try_method == "pdfplumber":
                        tables = self._extract_tables_pdfplumber(path, parsed_pages)
                    elif try_method == "tabula":
                        tables = self._extract_tables_tabula(path, parsed_pages)

                    if tables:
                        method = try_method
                        all_tables = tables
                        break
            else:
                # Use specific method
                methods_tried.append(method)
                if method == "camelot":
                    all_tables = self._extract_tables_camelot(path, parsed_pages)
                elif method == "pdfplumber":
                    all_tables = self._extract_tables_pdfplumber(path, parsed_pages)
                elif method == "tabula":
                    all_tables = self._extract_tables_tabula(path, parsed_pages)
                else:
                    raise ValueError(f"Unknown table extraction method: {method}")

            # Format tables based on output format
            formatted_tables = []
            for i, df in enumerate(all_tables):
                if table_format == "json":
                    formatted_tables.append({
                        "table_index": i,
                        "data": df.to_dict(orient="records"),
                        "shape": {"rows": len(df), "columns": len(df.columns)}
                    })
                elif table_format == "csv":
                    formatted_tables.append({
                        "table_index": i,
                        "data": df.to_csv(index=False),
                        "shape": {"rows": len(df), "columns": len(df.columns)}
                    })
                elif table_format == "markdown":
                    formatted_tables.append({
                        "table_index": i,
                        "data": df.to_markdown(index=False),
                        "shape": {"rows": len(df), "columns": len(df.columns)}
                    })

            return {
                "success": True,
                "tables": formatted_tables,
                "total_tables": len(formatted_tables),
                "method_used": method,
                "methods_tried": methods_tried,
                "pages_searched": pages or "all",
                "processing_time": round(time.time() - start_time, 2)
            }

        except Exception as e:
            error_msg = sanitize_error_message(str(e))
            logger.error(f"Table extraction failed: {error_msg}")
            return {
                "success": False,
                "error": error_msg,
                "methods_tried": methods_tried,
                "processing_time": round(time.time() - start_time, 2)
            }

    # Private helper methods (all synchronous for proper async pattern)
    def _extract_tables_camelot(self, pdf_path: Path, pages: Optional[List[int]] = None) -> List[pd.DataFrame]:
        """Extract tables using Camelot"""
        page_str = ','.join(map(str, [p+1 for p in pages])) if pages else 'all'

        # Try lattice mode first (for bordered tables)
        try:
            tables = camelot.read_pdf(str(pdf_path), pages=page_str, flavor='lattice')
            if len(tables) > 0:
                return [table.df for table in tables]
        except Exception:
            pass

        # Fall back to stream mode (for borderless tables)
        try:
            tables = camelot.read_pdf(str(pdf_path), pages=page_str, flavor='stream')
            return [table.df for table in tables]
        except Exception:
            return []

    def _extract_tables_tabula(self, pdf_path: Path, pages: Optional[List[int]] = None) -> List[pd.DataFrame]:
        """Extract tables using Tabula"""
        page_list = [p+1 for p in pages] if pages else 'all'

        try:
            tables = tabula.read_pdf(str(pdf_path), pages=page_list, multiple_tables=True)
            return tables
        except Exception:
            return []

    def _extract_tables_pdfplumber(self, pdf_path: Path, pages: Optional[List[int]] = None) -> List[pd.DataFrame]:
        """Extract tables using pdfplumber"""
        tables = []

        with pdfplumber.open(str(pdf_path)) as pdf:
            page_range = pages if pages else range(len(pdf.pages))
            for page_num in page_range:
                page = pdf.pages[page_num]
                page_tables = page.extract_tables()
                for table in page_tables:
                    if table and len(table) > 1:  # Skip empty tables
                        df = pd.DataFrame(table[1:], columns=table[0])
                        tables.append(df)

        return tables