doczyai-pipelines/src/pipelines/saas/file_processing.py

import concurrent.futures
import logging
from typing import TYPE_CHECKING, Dict, Optional

import pandas as pd
import src.codes.code_funcs as code_funcs
from src.pipelines.shared.preprocessing import (
    preprocessing_funcs,
    preprocess,
    hybrid_smart_chunking_funcs,
)
from src.pipelines.shared.preprocessing.exhibit_smart_chunking_funcs import ESC_CONFIG
from src.pipelines.shared.postprocessing import (
    aarete_derived,
    postprocess,
    postprocessing_funcs,
)
from src.pipelines.shared.extraction import (
    dynamic_funcs,
    one_to_n_funcs,
    one_to_one_funcs,
    row_funcs,
    tin_npi_funcs,
    exhibit_funcs,
)
from src.pipelines.shared.extraction.exhibit_funcs import Exhibit, ExhibitChunk
from src.utils import io_utils, logging_utils, string_utils, timing_utils
from src.constants.constants import Constants
from src import config
from src.prompts.fieldset import FieldSet
from src.utils.string_utils import datetime_str

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


def process_file(file_object, constants: Constants, run_timestamp):
    filename, contract_text = file_object

    # Set per-file logging context:
    # With this, all logging calls in this thread will now route to logs/{filename}.log
    # This includes logging from all called functions (preprocess, one_to_n_funcs, etc.)
    logging_utils.set_current_file(filename)
    logging.debug(f"{datetime_str()} Processing {filename}...")

    # Set default values
    dynamic_one_to_one_fields = FieldSet()
    process_one_to_n = config.FIELDS in ["all", "one_to_n"]
    process_one_to_one = config.FIELDS in ["all", "one_to_one"]
    # Initialize default fallback for final results
    final_results = pd.DataFrame([{"FILE_NAME": filename}])

    ################## PREPROCESS ##################
    with timing_utils.timed_block("preprocess", context=filename):
        contract_text = preprocess.clean_text(contract_text)
        text_dict, top_sheet_dict = preprocess.split_text(contract_text)
        if not text_dict:
            logging.warning(
                f"{datetime_str()} All pages removed as cover sheets, skipping processing - {filename}"
            )
            return final_results, pd.DataFrame([{"FILE_NAME": filename}])
        text_dict, removal_metadata = preprocess.clean_header_footer(text_dict)
    # ONE TO N PROCESSING
    one_to_n_results = pd.DataFrame()  # Initialize empty DataFrame for one_to_n_results
    if process_one_to_n:
        # Use split_text_with_pages() to get Page objects with table splitting
        # Note: split_text_with_pages() internally calls split_text() which applies headers/footers
        # But we need to apply headers/footers first, so we'll create pages_dict from the cleaned text_dict
        from src.pipelines.shared.extraction.page_funcs import Page
        from src.pipelines.shared.preprocessing import preprocessing_funcs as prep_funcs

        # Create pages_dict from text_dict (headers/footers already applied)
        pages_dict = prep_funcs.split_large_tables(text_dict)

        with timing_utils.timed_block("one_to_n_exhibit_chunking", context=filename):
            # Use pages_dict for exhibit chunking (preferred) or fall back to text_dict
            exhibit_header_dict = preprocess.one_to_n_exhibit_chunking(
                pages_dict=pages_dict,
                text_dict=text_dict,
                EXHIBIT_HEADER_MARKERS=constants.EXHIBIT_HEADER_MARKERS,
                filename=filename,
            )

            logging.info(f"{datetime_str()} Preprocessing Complete - {filename}")
        one_to_n_results = pd.DataFrame([{"FILE_NAME": filename}])  # Initialize here

        if string_utils.contains_reimbursement(contract_text):
            logging.info(
                f"{datetime_str()} Starting One-to-N extraction for {sum(len(value) for value in exhibit_header_dict.values() if isinstance(value, list))} exhibits - {filename}"
            )
            with timing_utils.timed_block("one_to_n_extraction", context=filename):
                # Get specific fields to extract (if configured)
                specific_fields = config.get_specific_fields_list()
                (
                    one_to_n_results,
                    dynamic_one_to_one_fields,
                ) = run_one_to_n_prompts(
                    pages_dict=pages_dict,
                    text_dict=text_dict,
                    exhibit_header_dict=exhibit_header_dict,
                    constants=constants,
                    filename=filename,
                    specific_fields=specific_fields,
                )
            if not one_to_n_results.empty:
                one_to_n_results["FILE_NAME"] = filename
                with timing_utils.timed_block(
                    "one_to_n_generate_reimb_ids", context=filename
                ):
                    one_to_n_results = postprocessing_funcs.generate_reimb_ids(
                        one_to_n_results
                    )
            logging.debug(f"{datetime_str()} One to N Complete - {filename}")

        final_results = one_to_n_results  # Set as final results

    # ONE TO ONE PROCESSING
    if process_one_to_one:
        # Get specific fields to extract (if configured)
        specific_fields = config.get_specific_fields_list()
        with timing_utils.timed_block("one_to_one_extraction", context=filename):
            one_to_one_results = run_one_to_one_prompts(
                filename,
                contract_text,
                text_dict,
                dynamic_one_to_one_fields,
                constants,
                specific_fields=specific_fields,
            )
        one_to_one_results["FILE_NAME"] = filename
        logging.debug(f"{datetime_str()} One to One Complete - {filename}")

        # Decide how to handle one_to_one results
        with timing_utils.timed_block(
            "merge_one_to_one_into_one_to_n", context=filename
        ):
            if not one_to_n_results.empty:
                # BOTH processed - merge into one_to_n
                final_results = row_funcs.merge_one_to_one_into_one_to_n(
                    one_to_n_results, one_to_one_results, constants
                )
            else:
                # ONLY one_to_one processed - convert dict to DataFrame
                final_results = pd.DataFrame([one_to_one_results])
    else:
        logging.debug(
            f"{datetime_str()} Fields not configured for One to One, Skipping - {filename}"
        )

    # split the rows based service term and duplicate the other fields accordingly
    if not final_results.empty and "SERVICE_TERM" in final_results.columns:
        with timing_utils.timed_block("split_service_terms", context=filename):
            final_results = one_to_n_funcs.split_service_terms(final_results, filename)
        logging.debug(f"{datetime_str()} Split Service Terms Complete - {filename}")

    # APPLY CODES IF ONE_TO_N WAS PROCESSED
    if not one_to_n_results.empty:
        with timing_utils.timed_block("code_processing", context=filename):
            results_with_code = code_funcs.code_breakout(final_results, constants)
            final_results = code_funcs.grouper_breakout(results_with_code)
        logging.debug(f"{datetime_str()} Codes Complete - {filename}")

    # POSTPROCESS
    with timing_utils.timed_block("postprocess", context=filename):
        cc_df, dashboard_df = postprocess.postprocess(final_results, constants)
    logging.debug(f"{datetime_str()} Postprocessing Complete - {filename}")

    ################## WRITE INDIVIDUAL ##################
    with timing_utils.timed_block("write_individual", context=filename):
        if config.WRITE_TO_S3:
            io_utils.write_s3(cc_df, filename, run_timestamp, "individual_cc")
        else:
            io_utils.write_local(cc_df, filename, "", "individual_cc")

    logging.debug(f"{datetime_str()} Writing Complete - {filename}")

    return cc_df, dashboard_df


def run_one_to_one_prompts(
    filename: str,
    contract_text: str,
    text_dict: dict[str, str],
    dynamic_one_to_one_fields: FieldSet,
    constants: Constants,
    specific_fields: Optional[list] = None,
):

    ################## INITIALIZE FIELDS ##################
    one_to_one_fields = FieldSet(
        relationship="one_to_one", file_path=config.FIELD_JSON_PATH
    ).combine(dynamic_one_to_one_fields)

    # Filter fields if specific_fields is provided
    if specific_fields is not None:
        one_to_one_fields = one_to_one_fields.filter_by_names(specific_fields)

    # Check if we need provider info fields
    provider_fields = set(config.FIELD_GROUPS.get("provider", []))
    run_provider_info = specific_fields is None or any(
        f in provider_fields for f in specific_fields
    )

    # Initialize results
    one_to_one_results = {}

    ################## RUN HYBRID SMART CHUNKED PROMPTS ##################
    # RAG function loads retrieval questions internally and matches with investment_prompts.json
    # Only run if there are fields to extract
    hybrid_smart_chunked_answers_dict = {}
    if one_to_one_fields.contains_fields():
        with timing_utils.timed_block(
            "one_to_one.hybrid_smart_chunking", context=filename
        ):
            hybrid_smart_chunked_answers_dict = (
                hybrid_smart_chunking_funcs.run_hybrid_smart_chunked_fields(
                    one_to_one_fields, constants, contract_text, filename, text_dict
                )
            )
    logging.debug(
        f"Hybrid Smart Chunked Answers for {filename}: {hybrid_smart_chunked_answers_dict}"
    )

    ################## COMBINE ANSWERS ##################
    for key, value in hybrid_smart_chunked_answers_dict.items():
        if not string_utils.is_empty(value):
            one_to_one_results[key] = value
        elif key not in one_to_one_results:
            # Add HSC's N/A if field doesn't exist yet
            one_to_one_results[key] = value

    ################## RUN PROVIDER INFO ##################
    if run_provider_info:
        with timing_utils.timed_block("one_to_one.provider_info", context=filename):
            prov_info_results = tin_npi_funcs.run_provider_info_fields(
                text_dict, filename, payer_name=""
            )
        one_to_one_results.update(prov_info_results)
        one_to_one_results = tin_npi_funcs.add_group_and_other(
            one_to_one_results, filename
        )

    ################## ADD AD FIELDS ################
    # Only run if date fields are requested
    date_fields = set(config.FIELD_GROUPS.get("dates", []))
    if specific_fields is None or any(f in date_fields for f in specific_fields):
        one_to_one_results = one_to_one_funcs.get_aarete_derived_dates(
            one_to_one_results, text_dict, filename
        )

    ################## ADD SIGNATURE COUNT ##################
    # Only run if signature count is requested
    if specific_fields is None or "NUM_SIGNED_SIGNATORY_LINE_COUNT" in specific_fields:
        one_to_one_results = one_to_one_funcs.add_signature_count(
            one_to_one_results, contract_text
        )

    ################## Crosswalk Fields ##################
    # All field format normalization is handled at prompt_calls level via field-aware parsers
    # Derived fields (AARETE_DERIVED_*, NUM_*_SIGNATORY_*) are created as strings directly
    # No additional normalization needed here
    one_to_one_results = aarete_derived.get_crosswalk_fields(
        [one_to_one_results], constants
    )

    ################## Fill NA Mapping ##################
    one_to_one_results = aarete_derived.fill_na_mapping(one_to_one_results)

    return one_to_one_results[0]


def process_page_reimbursements(
    exhibit: Exhibit,
    page_num: str,
    constants: Optional[Constants] = None,
    filename: Optional[str] = None,
    relevant_chunks: Optional[list[ExhibitChunk]] = None,
):
    """
    STEP 1 HELPER: Extract reimbursements from a single page (without exhibit-level fields).
    Runs: reimbursement_level → carveout_and_special_case → lesser_of_distribution → breakout

    Each relevant chunk on this page is processed individually through reimbursement_level,
    while downstream steps (lesser_of, etc.) receive page-based simplified exhibit context.
    Pages without relevant chunks are skipped (returns [], []).

    Args:
        exhibit: Exhibit object containing the page
        page_num: Page identifier (can be "27" or "27.1" for sub-pages)
        pages_dict: Dictionary mapping page numbers to Page objects (preferred)
        text_dict: Dictionary mapping page numbers to text strings (for backward compatibility)
        exhibit_page_nums: List of page numbers in the exhibit (for backward compatibility)
        constants: Constants object
        filename: Name of the file being processed
        specific_fields: Optional list of specific field names to extract.
                        If provided with exhibit-level-only fields (like claim_type),
                        skip detailed reimbursement extraction.
        relevant_chunks: List of relevant chunks on this page. Each chunk is processed
                        individually through reimbursement_level. If empty/None, page is skipped.
    """

    # No relevant chunks on this page - skip reimbursement extraction
    if not relevant_chunks:
        logging.debug(
            f"{datetime_str()} Page {page_num}: No relevant chunks, skipping - {filename}"
        )
        return [], []

    ############################### Reimbursement Primary ###############################
    # Run reimbursement_level per-chunk: each relevant chunk is processed individually
    # so the LLM sees focused, relevant text one chunk at a time
    reimbursement_level_answers = []
    for chunk in relevant_chunks:
        chunk_answers = one_to_n_funcs.reimbursement_level(
            chunk.text, constants, filename
        )
        if chunk_answers:
            reimbursement_level_answers.extend(chunk_answers)

    if not reimbursement_level_answers:
        logging.debug(
            f"{datetime_str()} Page {page_num}: No reimbursement answers found, skipping - {filename}"
        )
        return [], []

    # Page-level simplification of exhibit for downstream steps
    exhibit_text_simplified = preprocessing_funcs.simplify_exhibit(
        exhibit=exhibit,
        current_page_num=page_num,
    )

    ################################ Carveouts and Special Case ################################
    reimbursement_level_answers, special_case_answers = (
        one_to_n_funcs.carveout_and_special_case(
            reimbursement_level_answers, constants, filename
        )
    )

    ################################ Dynamic Code Assignment ################################
    reimbursement_level_answers = one_to_n_funcs.dynamic_code_assignment(
        reimbursement_level_answers, constants, filename
    )

    ############################### Lesser of Distribution ###############################
    # Note: We run lesser_of without dynamic fields in Step 1
    reimbursement_level_answers = one_to_n_funcs.lesser_of_distribution(
        reimbursement_level_answers,
        exhibit_text_simplified,
        page_num,
        constants,
        filename,
        exhibit,
    )

    ################################ Get Breakouts ###############################
    reimbursement_level_answers, special_case_answers = one_to_n_funcs.breakout(
        reimbursement_level_answers,
        special_case_answers,
        filename,
        constants,
    )

    ################################ Assign REIMB_PAGE ###############################
    for answer_dict in reimbursement_level_answers:
        answer_dict["REIMB_PAGE"] = page_num

    return reimbursement_level_answers, special_case_answers


def run_one_to_n_prompts(
    pages_dict: Optional[dict[str, "Page"]] = None,
    text_dict: Optional[dict[str, str]] = None,
    exhibit_header_dict: Optional[dict[str, list[str]]] = None,
    constants: Optional[Constants] = None,
    filename: Optional[str] = None,
    specific_fields: Optional[list] = None,
):
    """
    3-STEP APPROACH (per exhibit):
    Step 1: Extract reimbursements for exhibit pages (parallel within exhibit)
    Step 2: Run exhibit_level when exhibit has reimbursements
    Step 3: Run dynamic_assignment and combine results

    Args:
        pages_dict: Dictionary mapping page numbers to Page objects (preferred)
        text_dict: Dictionary mapping page numbers to text strings (for backward compatibility)
        exhibit_header_dict: Dictionary mapping page numbers to lists of headers starting on that page
                            Example: {'2': ['ARTICLE ONE', 'ARTICLE TWO'], '4': ['ARTICLE THREE']}
        constants: Constants object
        filename: Name of the file being processed
        specific_fields: Optional list of specific field names to extract.
                        If provided, only these fields will be extracted at exhibit level.
    """
    if exhibit_header_dict is None:
        exhibit_header_dict = {}

    total_exhibits = sum(len(headers) for headers in exhibit_header_dict.values())
    logging.debug(
        f"{datetime_str()} Processing {total_exhibits} exhibits with NEW 3-step approach - {filename}"
    )

    # Create Exhibit objects in order (maintains exhibit chain via prev_exhibit)
    # Uses header-based text splitting for precise exhibit boundaries
    exhibits = exhibit_funcs.create_exhibits_from_header_dict(
        exhibit_header_dict=exhibit_header_dict,
        pages_dict=pages_dict,
        text_dict=text_dict,
    )
    logging.info(
        f"{datetime_str()} Created {len(exhibits)} exhibit(s) from header dict - {filename}"
    )

    # Process exhibits serially; within each exhibit, process pages in parallel
    one_to_n_results = []

    for exhibit in exhibits:
        # Search for reimbursement-related chunks, then group by page
        relevant_chunks = exhibit.get_relevant_chunks(
            threshold=ESC_CONFIG.CHUNK_RELEVANCE_THRESHOLD
        )
        page_chunks_map = (
            exhibit.group_chunks_by_page(relevant_chunks) if relevant_chunks else {}
        )

        logging.debug(
            f"{datetime_str()} Processing exhibit {exhibits.index(exhibit) + 1}/{len(exhibits)}: "
            f"page={exhibit.exhibit_page}, header='{exhibit.exhibit_header}' - {filename}"
        )

        pages_to_process = exhibit.exhibit_page_nums

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=min(len(pages_to_process), 20)
        ) as executor:
            page_futures = {
                executor.submit(
                    process_page_reimbursements,
                    exhibit,
                    page_num,
                    constants,
                    filename,
                    relevant_chunks=page_chunks_map.get(
                        page_num
                    ),  # None for pages without relevant chunks
                ): page_num
                for page_num in pages_to_process
            }

            for future in concurrent.futures.as_completed(page_futures):
                try:
                    page_num = page_futures[future]
                    reimbursement_rows, special_case_rows = future.result()
                    exhibit.add_reimbursement_rows(
                        reimbursement_rows, special_case_rows
                    )
                except Exception as e:
                    logging.error(f"Error processing page {page_num}: {str(e)}")

        if not exhibit.has_reimbursements:
            logging.debug(
                f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: No reimbursements found, skipping steps 2-3 - {filename}"
            )
            continue

        logging.debug(
            f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: "
            f"{len(exhibit.reimbursement_rows)} reimbursement row(s), "
            f"{len(exhibit.special_case_rows)} special case row(s) - {filename}"
        )

        # STEP 2: exhibit_level for this exhibit
        exhibit_level_answers, dynamic_reimbursement_fields = (
            one_to_n_funcs.exhibit_level(
                exhibit.exhibit_text,
                exhibit.exhibit_header,
                exhibit.exhibit_page,
                constants,
                filename,
                specific_fields=specific_fields,
            )
        )
        exhibit.set_exhibit_level_data(
            exhibit_level_answers, dynamic_reimbursement_fields
        )
        logging.debug(
            f"Exhibit Level Answers for {filename}, Page {exhibit.exhibit_page}: {exhibit_level_answers}"
        )

        # Check if we need reimbursement-level processing
        # If specific_fields is set and doesn't include reimbursement fields, skip Step 3
        reimbursement_fields = {"SERVICE_TERM", "REIMB_TERM", "REIMB_PAGE"}
        skip_reimbursement_processing = specific_fields is not None and not any(
            f in reimbursement_fields for f in specific_fields
        )

        if skip_reimbursement_processing:
            # For exhibit-level-only extraction, just create minimal rows
            # with exhibit_level_answers to preserve the data
            if exhibit.exhibit_level_answers:
                minimal_row = exhibit.exhibit_level_answers.copy()
                minimal_row["REIMB_PAGE"] = exhibit.exhibit_page
                # Apply crosswalk mapping for derived fields (e.g., CLAIM_TYPE_CD -> AARETE_DERIVED_CLAIM_TYPE_CD)
                minimal_rows_with_crosswalk = aarete_derived.get_crosswalk_fields(
                    [minimal_row], constants
                )
                exhibit.final_rows = minimal_rows_with_crosswalk
                one_to_n_results.extend(minimal_rows_with_crosswalk)
        else:
            # STEP 3: dynamic assignment & combine for this exhibit
            # Get dynamic fields from previous exhibit if needed (for future use)
            if exhibit.dynamic_reimbursement_fields:
                dynamic_fields = exhibit.dynamic_reimbursement_fields
            elif (
                exhibit.prev_exhibit
                and not exhibit.prev_exhibit.has_reimbursements
                and exhibit.get_previous_exhibit_dynamic_fields().fields
            ):
                dynamic_fields = exhibit.get_previous_exhibit_dynamic_fields()
            else:
                dynamic_fields = None

            # Run dynamic assignment for ALL reimbursement rows in this exhibit
            # Note: dynamic_assignment expects a list of rows and returns a list of rows
            reimbursement_rows_with_dynamic = exhibit.reimbursement_rows
            if exhibit.reimbursement_rows and dynamic_fields is not None:
                reimbursement_rows_with_dynamic = dynamic_funcs.dynamic_assignment(
                    exhibit.reimbursement_rows,
                    dynamic_fields,
                    pages_dict,
                    text_dict,
                    exhibit,
                    constants,
                    filename,
                )

            # Combine reimbursement rows with exhibit-level answers
            combined_rows = row_funcs.combine_one_to_n(
                exhibit.exhibit_text,
                reimbursement_rows_with_dynamic,
                exhibit.special_case_rows,
                exhibit.exhibit_level_answers,
                filename,
            )

            # Run cleaning
            combined_rows = one_to_n_funcs.one_to_n_cleaning(
                combined_rows, exhibit.exhibit_text, constants, filename
            )

            exhibit.final_rows = combined_rows
            one_to_n_results.extend(combined_rows)
            logging.debug(
                f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: "
                f"{len(combined_rows)} final row(s) after combine & clean - {filename}"
            )

    logging.debug(
        f"{datetime_str()} STEP 3 COMPLETE: Combined {len(one_to_n_results)} total rows - {filename}"
    )

    ################## Add N/A Dynamic or Exhibit to One-to-One ##################
    dynamic_one_to_one_fields = dynamic_funcs.get_dynamic_one_to_one_fields(
        one_to_n_results, constants
    )

    ################## CONVERT TO DF ##################
    one_to_n_df = pd.DataFrame(one_to_n_results)

    return one_to_n_df, dynamic_one_to_one_fields