doczyai-pipelines/src/pipelines/shared/extraction/page_funcs.py

"""
Page class for managing page text and table splitting functionality.

This module provides the Page class which can represent either a regular page
or a collection of sub-pages (for split tables). It handles table detection,
parsing, and splitting into manageable chunks.
"""

import ast
import logging
import re
from typing import Optional

from src import config
from src.constants.delimiters import TABLE_START_MARKER, TABLE_END_MARKER


class Page:
    """
    Represents a single page or a collection of sub-pages (for split tables).

    A Page can be in one of two states:
    1. Regular page: Contains a single text string
    2. Split table page: Contains multiple sub-pages (chunks of a table)

    Attributes:
        page_num (str): Original page number from the document
        is_split (bool): Whether this page has been split into sub-pages
        text (str): Full text of the page (if not split) or combined text of all sub-pages
        sub_pages (dict[str, str]): Mapping of sub-page identifiers to sub-page text
                                   Format: {"page_num.1": "text", "page_num.2": "text", ...}
        original_text (str): Original unsplit text (preserved for reference)
    """

    def __init__(self, page_num: str, text: str):
        """
        Initialize a Page object.

        Args:
            page_num: The page number identifier
            text: The text content of the page
        """
        self.page_num = page_num
        self.original_text = text
        self.text = text
        self.is_split = False
        self.sub_pages: dict[str, str] = {}

    def has_table(self) -> bool:
        """
        Check if page contains a table.

        Returns:
            bool: True if page contains table markers, False otherwise
        """
        return TABLE_START_MARKER in self.text and TABLE_END_MARKER in self.text

    def split_table(
        self,
        max_cells: Optional[int] = None,
        min_rows_to_split: Optional[int] = None,
    ) -> None:
        """
        Split table into sub-pages based on cell count (rows × columns).

        This method will:
        1. Extract the table content
        2. Parse table rows and detect column count
        3. Split into chunks where rows × columns <= max_cells
        4. Create sub-pages with appropriate identifiers
        5. Preserve text before/after the table

        Args:
            max_cells: Maximum number of cells (rows × columns) per sub-page.
                      If None, uses config.MAX_TABLE_CELLS_PER_SUBPAGE
            min_rows_to_split: Minimum rows to trigger splitting.
                              If None, uses config.MIN_TABLE_ROWS_TO_SPLIT
        """
        if max_cells is None:
            max_cells = config.MAX_TABLE_CELLS_PER_SUBPAGE
        if min_rows_to_split is None:
            min_rows_to_split = config.MIN_TABLE_ROWS_TO_SPLIT
        if not self.has_table():
            return

        table_content = detect_table_in_text(self.text)
        if not table_content:
            return

        table_text, text_before, text_after = table_content

        # Parse table rows and detect column count
        table_rows, num_columns = parse_table_rows(table_text, self.page_num)

        if len(table_rows) < min_rows_to_split:
            # Table is too small to split
            return

        # Calculate if table needs splitting based on total cell count
        total_cells = len(table_rows) * num_columns
        if total_cells <= max_cells:
            # Table is small enough, don't split
            return

        # Split table into chunks based on cell count
        table_chunks = split_table_by_cell_count(table_rows, num_columns, max_cells)

        # Create sub-pages
        self.sub_pages = {}
        for i, chunk in enumerate(table_chunks, start=1):
            sub_page_id = f"{self.page_num}.{i}"
            sub_page_text = ""

            # Add text before table only to first sub-page
            if i == 1:
                sub_page_text += text_before

            # Add table chunk with markers
            sub_page_text += f"{TABLE_START_MARKER}\n{chunk}\n{TABLE_END_MARKER}"

            # Add text after table only to last sub-page
            if i == len(table_chunks):
                sub_page_text += text_after

            self.sub_pages[sub_page_id] = sub_page_text

        # Mark as split
        self.is_split = True

        # Update combined text
        self.text = "\n".join(self.sub_pages.values())

    def get_text(self, sub_page: Optional[str] = None) -> str:
        """
        Get text for a specific sub-page or full page.

        Args:
            sub_page: Optional sub-page identifier. Can be:
                     - None: Returns full page text (combined sub-pages if split)
                     - Full identifier (e.g., "1.1"): Returns that specific sub-page
                     - Partial identifier (e.g., "1"): Returns sub-page "page_num.1"

        Returns:
            str: The requested text content
        """
        if not self.is_split:
            return self.text

        if sub_page is None:
            # Return combined text of all sub-pages
            return self.text

        # Handle full sub-page identifier (e.g., "1.1")
        if sub_page in self.sub_pages:
            return self.sub_pages[sub_page]

        # Handle partial identifier (e.g., "1" -> "page_num.1")
        sub_page_id = f"{self.page_num}.{sub_page}"
        if sub_page_id in self.sub_pages:
            return self.sub_pages[sub_page_id]

        # Not found, return empty string
        return ""

    def get_all_sub_pages(self) -> dict[str, str]:
        """
        Get all sub-pages if split, else return single page.

        Returns:
            dict: Mapping of page/sub-page identifiers to text content
        """
        if not self.is_split:
            return {self.page_num: self.text}

        return self.sub_pages.copy()

    def get_all_page_identifiers(self) -> list[str]:
        """
        Get list of all page/sub-page identifiers.

        Returns:
            list[str]: List of page identifiers (e.g., ["27"] or ["27.1", "27.2", ...])
        """
        if not self.is_split:
            return [self.page_num]

        return sorted(self.sub_pages.keys(), key=lambda x: float(x.split(".")[-1]))

    def __repr__(self) -> str:
        """String representation for debugging."""
        if self.is_split:
            return (
                f"Page(page_num={self.page_num}, is_split=True, "
                f"sub_pages={len(self.sub_pages)})"
            )
        return f"Page(page_num={self.page_num}, is_split=False)"


def detect_table_in_text(text: str) -> Optional[tuple[str, str, str]]:
    """
    Detect if text contains a table and return table content with context.

    Args:
        text: The text to analyze

    Returns:
        Optional[tuple[str, str, str]]: If table found, returns (table_text, text_before, text_after).
                                       Otherwise returns None.
    """
    if TABLE_START_MARKER not in text or TABLE_END_MARKER not in text:
        return None

    start_idx = text.find(TABLE_START_MARKER)
    end_idx = text.find(TABLE_END_MARKER, start_idx + len(TABLE_START_MARKER))

    if end_idx == -1:
        # Table start found but no end marker - return None
        return None

    # Extract components
    text_before = text[:start_idx].strip()
    table_text = text[start_idx + len(TABLE_START_MARKER) : end_idx].strip()
    text_after = text[end_idx + len(TABLE_END_MARKER) :].strip()

    return (table_text, text_before, text_after)


def parse_table_rows(table_text: str, page_num: str = "") -> tuple[list[str], int]:
    """
    Parse table text into individual rows and detect column count.

    Expects Textract output format: [['Service', 'Payment'], ['Item 1', '$100'], ...]

    The table text may contain descriptive text before/after the list structure.
    This function will find and extract the list structure from within the table text.

    Args:
        table_text: The table content between markers (may include text before/after list)
        page_num: Page number for logging purposes (optional)

    Returns:
        tuple[list[str], int]: List of table rows as strings and column count.
                              Returns empty list and 0 columns if format is unexpected.
    """
    if not table_text:
        return [], 0

    rows = []
    num_columns = 0

    # First, try to find the list structure within the table text
    # Look for the first '[' that starts a list structure
    list_start = table_text.find("[")

    if list_start == -1:
        # No list structure found
        logging.warning(
            f"Table on page {page_num} does not contain a list structure. "
            "Table will not be split."
        )
        return [], 0

    # Extract from the first '[' onwards
    # Use bracket matching to find the closing ']' for the entire list
    bracket_count = 0
    list_end = -1

    for i in range(list_start, len(table_text)):
        if table_text[i] == "[":
            bracket_count += 1
        elif table_text[i] == "]":
            bracket_count -= 1
            if bracket_count == 0:
                list_end = i + 1
                break

    if list_end == -1:
        # Couldn't find matching closing bracket
        logging.warning(
            f"Table on page {page_num} contains incomplete list structure. "
            "Table will not be split."
        )
        return [], 0

    # Extract the list text
    list_text = table_text[list_start:list_end]

    # Parse the list structure
    try:
        parsed = ast.literal_eval(list_text)
        if isinstance(parsed, list) and len(parsed) > 0:
            # Convert each row to a string representation
            for row in parsed:
                if isinstance(row, list):
                    # Update column count based on first row
                    if num_columns == 0:
                        num_columns = len(row)
                    # Join list items with a separator
                    row_str = [str(item) for item in row]
                    rows.append(row_str)
                else:
                    # Non-list row - log warning but include it
                    logging.warning(
                        f"Table on page {page_num} contains non-list row: {row}. "
                        "Expected all rows to be lists."
                    )
                    rows.append(str(row))
        else:
            logging.warning(
                f"Table on page {page_num} is not in expected format. "
                f"Expected list of lists, got: {type(parsed).__name__}. "
                "Table will not be split."
            )
            return [], 0
    except (ValueError, SyntaxError) as e:
        logging.warning(
            f"Table on page {page_num} could not be parsed as expected list format. "
            f"Error: {str(e)}. Table will not be split."
        )
        return [], 0

    if num_columns == 0:
        logging.warning(
            f"Table on page {page_num} has no columns detected. "
            "Table will not be split."
        )
        return [], 0

    return rows, num_columns


def split_table_by_cell_count(
    table_rows: list[str], num_columns: int, max_cells: int = 40
) -> list[str]:
    """
    Split table rows into chunks based on cell count (rows × columns).

    The product of rows × columns in each chunk will be <= max_cells.
    This means tables with more columns will have fewer rows per chunk.

    Args:
        table_rows: List of table row strings
        num_columns: Number of columns in the table
        max_cells: Maximum number of cells (rows × columns) per chunk (default: 40)

    Returns:
        list[str]: List of table chunks, each as a string
    """
    if not table_rows or num_columns == 0:
        return []

    # Calculate maximum rows per chunk based on column count
    # max_rows = floor(max_cells / num_columns)
    # Minimum of 1 row per chunk
    max_rows = max(1, max_cells // num_columns)

    chunks = []
    for i in range(0, len(table_rows), max_rows):
        chunk = table_rows[i : i + max_rows]
        chunk_strs = []
        for row in chunk:
            if isinstance(row, list):
                chunk_strs.append(" ".join(str(cell) for cell in row))
            else:
                chunk_strs.append(str(row))
        chunks.append("\n".join(chunk_strs))

    return chunks