doczyai-pipelines/src/pipelines/shared/preprocessing/preprocessing_funcs.py

import logging
import re
from typing import TYPE_CHECKING, Optional

import src.utils.string_utils as string_utils
from src.pipelines.saas.prompts import prompt_calls
from src import config

if TYPE_CHECKING:
    from src.pipelines.shared.extraction.page_funcs import Page
    from src.pipelines.shared.extraction.exhibit_funcs import Exhibit


def remove_page_indicators(contract_text: str) -> str:
    """Clean textract output by removing page number indicators in the form of "Page X of Y"

    This function processes input text to remove lines that indicate page numbers (e.g. 'Page 1 of 10')

    Args:
        contract_text (str): Raw text output from Textract to be cleaned

    Returns:
        str: cleaned text with newlines and page number indicators removed
    """

    if contract_text:
        cleaned_text = re.sub(r"Page [0-9]+ of [0-9]+\n\n", " ", contract_text)
    else:
        cleaned_text = contract_text
    return cleaned_text


def split_text(text: str) -> dict[str, str]:
    """Split text on pages by the string `Start of Page No. = '

    Args:
        text (str): Raw text output from Textract to be split

    Returns:
        dict[str, str]: A dictionary, keyed by the string page number and valued by the page text.
    """
    temp_list = text.split("Start of Page No. = ")
    text_list = re.split(r"Start of Page No. = [0-9]+\n", text)

    text_dict = {}
    for i in range(len(text_list)):
        if temp_list[i]:
            text_dict[temp_list[i].split()[0]] = text_list[
                i
            ]  # splits on whitespace characters, which includes spaces, tabs, and newline characters

    return {k: v for k, v in text_dict.items() if k != "Document"}


def clean_law_symbols(contract_text):
    if not contract_text:
        return contract_text
    # First correction: Replace '$$' with '$'
    contract_text = contract_text.replace("$$", "$")

    contract_text = re.sub(
        r"(U\.?S\.?C\.?) \$", r"\1§", contract_text
    )  # replaces $ with § when it follows abbreviations like "U.S.C." (United States Code).
    contract_text = re.sub(
        r"(C\.?F\.?R\.?) \$", r"\1§", contract_text
    )  # replaces $ with § when it follows abbreviations like "C.F.R." (Code of Federal Regulations).

    # Second correction: Replace '$' with '§' when followed by a number with three decimal places
    contract_text = re.sub(r"\$(?=\d+\.\d{3})", "§", contract_text)

    return contract_text


def clean_newlines(contract_text: str) -> str:
    """Clean textract output by removing newlines and page number indicators.

    This function processes input text to:
    1. Remove all newline characters.
    2. Remove lines that indicate page numbers (e.g. 'Page 1 of 10')

    Args:
        contract_text (str): Raw text output from Textract to be cleaned

    Returns:
        str: cleaned text with newlines and page number indicators removed
    """
    cleaned_text = re.sub(r"(?<!\n)\n(?!\n)", " ", contract_text)
    cleaned_text = re.sub(r"Page [0-9]+ of [0-9]+\n\n", " ", contract_text)
    return cleaned_text


def filter_quick_review(text_dict):
    """
    Cover Sheets (aka Top Sheets or Quick Review pages) are pages stapled to the front of the contract that contain manually written summaries of the contract's contents.
    They are NOT legally binding documents, and because they are often manually filled out and hand-written, are more prone to have erroneous information than the rest of the contract.
    For Doczy.AI, we should not pull any information from top sheets, except as a very last resort.

    filter_quick_review() splits the input dictionary into two dictionaries, one containing the contract pages, the other containing the top sheet pages.

    The function checks each page's text for the keywords "QUICK REVIEW", "TOP SHEET", and "COVER SHEET". It categorizes the text into two separate dictionaries: one for texts that do not contain any of these keywords, and another for texts that do.

    Parameters:
    text_dict (dict): A dictionary where the key is the page number and the value is the text of that page.

    Returns:
    tuple of two dicts:
        - The first dictionary contains pages that do not have the specified keywords.
        - The second dictionary includes pages that contain any of the specified keywords.
    """
    keywords = [
        "QUICK REVIEW",
        "TOP SHEET",
        "COVER SHEET",
        "PROVIDER SUBMISSION FORM",
        "CONTRACT REVIEW FORM",
    ]
    regular_pages = {}
    quick_review_pages = {}

    for page_num, page_text in text_dict.items():
        page_text_upper = page_text.upper()
        if any(keyword in page_text_upper for keyword in keywords):
            quick_review_pages[page_num] = page_text
        else:
            regular_pages[page_num] = page_text

    return regular_pages, quick_review_pages


def get_exhibit_pages(
    text_dict: Optional[dict[str, str]] = None,
    pages_dict: Optional[dict[str, "Page"]] = None,
    EXHIBIT_HEADER_MARKERS: Optional[list] = None,
    filename: Optional[str] = None,
) -> tuple[list[str], dict[str, str]]:
    """
    Identify exhibit pages from the text dictionary or pages dictionary.

    This function processes the dictionary to identify pages that are considered exhibits.
    It uses the first page as an exhibit and checks subsequent pages for exhibit headers.
    If a page is identified as an exhibit, it is added to the exhibit pages list and its
    header is stored in the exhibit header dictionary.
    If a page does not have a header or is not an exhibit, it is skipped.

    Args:
        text_dict: A dictionary where keys are page numbers (as strings) and values are
                  the text of those pages (for backward compatibility)
        pages_dict: A dictionary where keys are page numbers and values are Page objects
                   (preferred)
        EXHIBIT_HEADER_MARKERS: List of markers to identify exhibit headers
        filename: The name of the file being processed, used for logging and LLM invocation.

    Returns:
        tuple[list[str], dict[str, str]]: A tuple containing:
            - A list of page numbers (as strings) that are identified as exhibit pages.
            - A dictionary where keys are page numbers and values are the headers of those exhibits.

    Note:
        Either text_dict or pages_dict must be provided. If both are provided,
        pages_dict takes precedence.
    """
    if EXHIBIT_HEADER_MARKERS is None:
        EXHIBIT_HEADER_MARKERS = []
    if filename is None:
        filename = ""

    # Determine which dictionary to use
    if pages_dict is not None:
        page_dict = pages_dict
        use_pages = True
    elif text_dict is not None:
        page_dict = text_dict
        use_pages = False
    else:
        raise ValueError("Either text_dict or pages_dict must be provided")

    sorted_pages = sorted(page_dict.keys(), key=string_utils.page_key_sort)
    exhibit_pages = []
    exhibit_header_dict = {}

    for page_num in sorted_pages:
        # Get page content - handle both Page objects and strings
        if use_pages:
            page_obj = page_dict[page_num]
            page_content = page_obj.get_text()
        else:
            page_content = page_dict[page_num]

        if "." not in page_num or ".0" in page_num:
            # Note: Header markers are now included in EXHIBIT_HEADER_INSTRUCTION() for caching
            exhibit_header = prompt_calls.prompt_exhibit_header(page_content, filename)
            if "N/A" in exhibit_header:
                is_exhibit = False
            else:
                exhibit_pages.append(page_num)
                exhibit_header_dict[page_num] = exhibit_header
                is_exhibit = True
        elif is_exhibit == True:
            exhibit_pages.append(page_num)
    return exhibit_pages, exhibit_header_dict


def link_exhibit_pages(
    all_exhibit_headers: dict[str, str],
    text_dict: Optional[dict[str, str]] = None,
    pages_dict: Optional[dict[str, "Page"]] = None,
    filename: Optional[str] = None,
):
    """
    Filters exhibit pages to include only those with meaningfully different headers.

    This function processes exhibit headers sequentially and uses literal comparison and
    LLM-based semantic comparison to determine when exhibits represent continuation pages
    vs. distinct exhibits.  Continuation pages are filtered out of the headers list and
    mapping, and when used downstream are included in their parent exhibit's text chunk.
    Links exhibit pages based on their headers. If the header of a page is identical
    (either literally or semantically) to the previous page's header, it is not added
    to the final list of exhibit pages.
    This function will link contiguous header pages together, so that the final list of
    exhibit pages will only contain pages that are meaningfully different from the
    previous page.  E.g. "1" with header {"Exhibit 1"} and "2" with header {"Continued
    exhibit 1"} will be linked together, but "3" with header {"Exhibit 2"} will not be
    linked to either of them.

    Args:
        all_exhibit_headers: A dictionary mapping page numbers (as strings) to their
                            exhibit headers. Expected to be in chronological order
        text_dict: A dictionary mapping page numbers to their text content
                  (for backward compatibility)
        pages_dict: A dictionary mapping page numbers to Page objects (preferred)
        filename: The name of the file being processed, used for logging and LLM invocation.

    Returns:
        tuple[list[str], dict[str, str]]: A tuple containing:
            - A list of page numbers (as strings) that are identified as exhibit pages.
            - A dictionary where keys are page numbers and values are the headers of those exhibits.

    Example:
        Input: {"1": "Exhibit A", "2": "Exhibit A (continued)", "3": "Exhibit B"}
        Output (["1", "3"], {"1": "Exhibit A", "3": "Exhibit B"})

    Notes:
        - First exhibit page is always included.
        - Literal header comparison is performed first to optimize LLM usage and latency.
        - Uses `EXHIBIT_LINKAGE` prompt for semantic similarity assessment
        - Expects LLM to return "Y" or "N" enclosed in pipe delimiters
        - Pages with non-distinct headers are excluded from the final results.
        - Either text_dict or pages_dict must be provided (pages_dict takes precedence).
    """
    if filename is None:
        filename = ""
    first_exhibit = True
    final_exhibit_pages = []
    final_exhibit_headers = {}
    for page_num, page_header in all_exhibit_headers.items():
        if first_exhibit:
            exhibit_is_different = "N"
            previous_header = page_header
            final_exhibit_pages.append(page_num)
            final_exhibit_headers[page_num] = page_header
            first_exhibit = False
        else:
            # If literally identical, skip prompt (unless it has a table)
            if page_header.upper().strip() == previous_header.upper().strip():
                exhibit_is_different = "N"
            elif string_utils.is_empty(previous_header):
                exhibit_is_different = "Y"
                previous_header = page_header
            else:
                exhibit_is_different = prompt_calls.prompt_exhibit_linkage(
                    page_header, previous_header, filename
                )
                previous_header = page_header

            # If the prompt has identified that the current exhibit is meaningfully different than the previous
            if "Y" in exhibit_is_different:
                final_exhibit_pages.append(page_num)
                final_exhibit_headers[page_num] = page_header

    return final_exhibit_pages, final_exhibit_headers


def chunk_by_exhibit(
    text_dict: Optional[dict[str, str]] = None,
    pages_dict: Optional[dict[str, "Page"]] = None,
    exhibit_pages: Optional[list] = None,
    *args,  # For backward compatibility with positional arguments
) -> dict:
    """
    Organizes pages into groups based on their association with specific exhibits.

    This function assigns each page number from the dictionary to an exhibit
    based on the `exhibit_pages` list. Pages are grouped under the nearest preceding
    page number in `exhibit_pages`. If a page number is itself in `exhibit_pages`,
    it starts a new exhibit group.

    Parameters:
        text_dict: A dictionary where keys are page numbers and values are page text
                  (for backward compatibility)
        pages_dict: A dictionary where keys are page numbers and values are Page objects
                   (preferred)
        exhibit_pages: A list of page numbers that mark the beginning of a new exhibit.

    Returns:
        dict: A dictionary where keys are exhibit identifiers (page numbers from exhibit_pages)
              and values are lists of page numbers that belong to that exhibit.
              Example: {'1': ['1', '2'], '3': ['3', '4', '5']}

    Note:
        Either text_dict or pages_dict must be provided. If both are provided,
        pages_dict takes precedence.
        For backward compatibility: can be called as chunk_by_exhibit(text_dict, exhibit_pages)
    """
    # Backward compatibility: detect when called with old positional signature
    # Old signature: chunk_by_exhibit(text_dict, exhibit_pages)
    # When called positionally, Python binds: text_dict=dict, pages_dict=list, exhibit_pages=None
    if (
        pages_dict is not None
        and isinstance(pages_dict, list)
        and exhibit_pages is None
    ):
        # This is the old signature - pages_dict is actually exhibit_pages
        exhibit_pages = pages_dict
        pages_dict = None

    if exhibit_pages is None:
        exhibit_pages = []

    # Determine which dictionary to use
    if pages_dict is not None:
        if not isinstance(pages_dict, dict):
            raise TypeError(f"pages_dict must be a dict, got {type(pages_dict)}")
        page_dict = pages_dict
    elif text_dict is not None:
        if not isinstance(text_dict, dict):
            raise TypeError(f"text_dict must be a dict, got {type(text_dict)}")
        page_dict = text_dict
    else:
        raise ValueError("Either text_dict or pages_dict must be provided")

    if len(exhibit_pages) == 0:
        return {key: [key] for key in page_dict.keys()}

    exhibit_chunk_mapping = {}
    current_exhibit = list(page_dict.keys())[0]

    for page_num in page_dict.keys():
        if page_num in exhibit_pages:
            current_exhibit = page_num
            # Initialize the list for this exhibit if not already present
            if current_exhibit not in exhibit_chunk_mapping:
                exhibit_chunk_mapping[current_exhibit] = []

        # Add the page to the current exhibit's list
        if current_exhibit not in exhibit_chunk_mapping:
            exhibit_chunk_mapping[current_exhibit] = []
        exhibit_chunk_mapping[current_exhibit].append(page_num)

    return exhibit_chunk_mapping


def simplify_exhibit(
    text_dict: Optional[dict[str, str]] = None,
    pages_dict: Optional[dict[str, "Page"]] = None,
    exhibit: Optional["Exhibit"] = None,
    exhibit_page_nums: Optional[list[str]] = None,
    current_page_num: Optional[str] = None,
) -> str:
    """
    Simplify exhibit text by replacing table content with 'Table X' placeholders on all pages
    except the current page being processed.

    Args:
        text_dict: Dictionary mapping page numbers to page text (for backward compatibility)
        pages_dict: Dictionary mapping page numbers to Page objects (preferred)
        exhibit: Exhibit object (preferred - uses exhibit.pages and exhibit.exhibit_page_nums)
        exhibit_page_nums: List of page numbers that belong to this exhibit
                         (required if exhibit is not provided)
        current_page_num: The current page being processed (will NOT be simplified)
                        (required if exhibit is not provided)

    Returns:
        str: The exhibit text with tables simplified on all pages except current_page_num

    Note:
        Either (exhibit) or (text_dict/pages_dict + exhibit_page_nums + current_page_num) must be provided.
        If exhibit is provided, it takes precedence.
    """
    # Check if any parameter that should be a dict is actually an Exhibit object
    if text_dict is not None and hasattr(text_dict, "pages"):
        # text_dict is actually an Exhibit object
        exhibit = text_dict
        text_dict = None
    elif pages_dict is not None and hasattr(pages_dict, "pages"):
        # pages_dict is actually an Exhibit object
        exhibit = pages_dict
        pages_dict = None

    # Handle Exhibit object (preferred approach) - but only if it's actually an Exhibit
    if exhibit is not None and hasattr(exhibit, "pages"):
        if exhibit.pages is not None:
            pages_dict = exhibit.pages
        else:
            # Fall back to text_dict if pages not available
            text_dict = {
                page_id: exhibit.get_page_text(page_id)
                for page_id in exhibit.exhibit_page_nums
            }
        exhibit_page_nums = exhibit.exhibit_page_nums
        # For current_page_num, we need to determine which base page the current_page_num belongs to
        # If current_page_num is a sub-page like "27.1", the base page is "27"
        if current_page_num is None:
            # Try to infer from exhibit_page_nums - use first page as default
            current_page_num = exhibit_page_nums[0] if exhibit_page_nums else ""
        else:
            # Extract base page if it's a sub-page
            if "." in current_page_num and not current_page_num.startswith("."):
                current_page_num = current_page_num.rsplit(".", 1)[0]

    # Validate required parameters
    if exhibit_page_nums is None:
        raise ValueError(
            "exhibit_page_nums must be provided if exhibit is not provided"
        )
    if current_page_num is None:
        raise ValueError("current_page_num must be provided if exhibit is not provided")

    # Determine which dictionary to use
    use_pages = False
    if pages_dict is not None:
        page_dict = pages_dict
        use_pages = True
    elif text_dict is not None:
        page_dict = text_dict
        use_pages = False
    else:
        raise ValueError("Either text_dict, pages_dict, or exhibit must be provided")

    def replace_tables_with_placeholder(
        page_text: str, starting_table_num: int
    ) -> tuple[str, int]:
        """Replace table blocks with 'Table X' labels and return next table number."""
        if not page_text:
            return page_text, starting_table_num

        start_marker = "-------Table Start--------"
        end_marker = "-------Table End--------"
        result = []
        pos = 0
        table_num = starting_table_num

        while True:
            start_idx = page_text.find(start_marker, pos)
            if start_idx == -1:
                # No more tables, append remaining text
                result.append(page_text[pos:])
                break

            # Append text before table
            result.append(page_text[pos:start_idx])

            # Find end of table
            end_idx = page_text.find(end_marker, start_idx + len(start_marker))

            # Replace table with placeholder
            result.append(f"Table {table_num}")
            table_num += 1

            # Move position past the table end marker
            if end_idx != -1:
                pos = end_idx + len(end_marker)
            else:
                # No end marker found, stop processing
                break

        return "".join(result), table_num

    # Process each page/sub-page in the exhibit
    simplified_exhibit_pages = []
    table_counter = 1

    # Extract base page from current_page_num if it's a sub-page
    current_base_page = current_page_num
    if "." in current_page_num and not current_page_num.startswith("."):
        current_base_page = current_page_num.rsplit(".", 1)[0]

    for page_id in exhibit_page_nums:
        # Extract base page for comparison
        base_page = page_id
        if "." in page_id and not page_id.startswith("."):
            base_page = page_id.rsplit(".", 1)[0]

        # Get page text
        if use_pages:
            # Handle sub-pages: "27.1" means get sub-page "1" from page "27"
            if "." in page_id and not page_id.startswith("."):
                base_page_num, sub_page = page_id.rsplit(".", 1)
                if base_page_num in page_dict:
                    page_text = page_dict[base_page_num].get_text(sub_page)
                else:
                    page_text = ""
            elif page_id in page_dict:
                page_text = page_dict[page_id].get_text()
            else:
                page_text = ""
        else:
            page_text = page_dict.get(page_id, "")

        # Check if this is the current page (match base page)
        if base_page == current_base_page:
            # Don't simplify the current page - keep it as-is
            simplified_exhibit_pages.append(page_text)
            # But still need to count tables to keep numbering consistent
            if "-------Table Start--------" in page_text:
                _, table_counter = replace_tables_with_placeholder(
                    page_text, table_counter
                )
        else:
            # Simplify other pages by replacing tables with placeholders
            if "-------Table Start--------" in page_text:
                simplified_page, table_counter = replace_tables_with_placeholder(
                    page_text, table_counter
                )
                simplified_exhibit_pages.append(simplified_page)
            else:
                simplified_exhibit_pages.append(page_text)

    return "\n".join(simplified_exhibit_pages)


def split_large_tables(
    text_dict: dict[str, str],
    max_cells: Optional[int] = None,
    min_rows_to_split: Optional[int] = None,
) -> dict[str, "Page"]:
    """
    Identify and split large tables into sub-pages.

    This function converts a text dictionary to Page objects and splits
    large tables based on cell count (rows × columns).

    Args:
        text_dict: Original page number to text mapping
        max_cells: Maximum cells (rows × columns) per table 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

    Returns:
        dict mapping page numbers to Page objects
    """
    from src.pipelines.shared.extraction.page_funcs import Page

    if not config.TABLE_SPLITTING_ENABLED:
        # If table splitting is disabled, create Page objects without splitting
        return {
            page_num: Page(page_num, page_text)
            for page_num, page_text in text_dict.items()
        }

    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

    pages_dict = {}

    for page_num, page_text in text_dict.items():
        page = Page(page_num, page_text)

        # Check if page has a table and if it needs splitting
        if page.has_table():
            # Split table if it exceeds thresholds
            page.split_table(max_cells=max_cells, min_rows_to_split=min_rows_to_split)

        pages_dict[page_num] = page

    return pages_dict