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.runtime_context import get_active_client
from src.pipelines.vendors.shared.extraction import (
    run_full_context_fields as vendor_run_full_context_fields,
)
from src.pipelines.shared.extraction.exhibit_funcs import Exhibit, ExhibitChunk
from src.utils import (
    instrumentation,
    instrumentation_context,
    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}...")
    instrumentation.log("file_start", filename=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 (
        instrumentation_context.instr_scope(segment="preprocessing"),
        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 (
            instrumentation_context.instr_scope(segment="preprocessing"),
            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,
                contract_text=contract_text,
            )

            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 (
            instrumentation_context.instr_scope(segment="one_to_one"),
            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 results for unified output
                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 (
            instrumentation_context.instr_scope(segment="code_breakout"),
            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}")

    # Derive AARETE_DERIVED_PROGRAM / AARETE_DERIVED_PRODUCT from PROGRAM / PRODUCT
    if not final_results.empty:
        with timing_utils.timed_block(
            "add_aarete_derived_program_product", context=filename
        ):
            final_results = one_to_n_funcs.add_aarete_derived_program_product(
                final_results, constants, filename
            )

    # Derive AARETE_DERIVED_LOB from AARETE_DERIVED_PROGRAM / AARETE_DERIVED_PRODUCT
    if not final_results.empty:
        with timing_utils.timed_block("fill_na_mapping", context=filename):
            final_results = aarete_derived.fill_na_mapping(final_results, constants)

    # POSTPROCESS
    with (
        instrumentation_context.instr_scope(segment="postprocess"),
        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}")

    instrumentation.log(
        "file_end",
        filename=filename,
        extra_json=f'{{"final_rows": {len(cc_df)}}}',
    )
    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)

    ################## LOAD CLIENT-SPECIFIC FIELDS ##################
    client_name = get_active_client()
    client_prompts_path = config.CLIENT_PROMPTS_MAP.get(client_name)
    client_full_context_fields = None

    if client_prompts_path:
        client_fields = FieldSet(
            relationship="one_to_one", file_path=client_prompts_path
        )
        if client_fields.contains_fields():
            # smart_chunked fields (e.g. BCBS OFFSET_TERM) are combined into HSC fields
            smart_chunked = client_fields.filter(field_type="smart_chunked")
            if smart_chunked.contains_fields():
                one_to_one_fields = one_to_one_fields.combine(smart_chunked)

            # full_context fields (e.g. Clover's 12 fields) are extracted separately after HSC
            full_context = client_fields.filter(field_type="full_context")
            if full_context.contains_fields():
                client_full_context_fields = full_context

            logging.info(
                f"Loaded {len(client_fields.fields)} client-specific fields for {client_name}"
            )

    # 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: dict = {}
    hybrid_smart_chunked_metadata_dict: 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_chunked_metadata_dict,
            ) = hybrid_smart_chunking_funcs.run_hybrid_smart_chunked_fields(
                one_to_one_fields,
                constants,
                contract_text,
                filename,
                text_dict,
                return_metadata=True,
            )
    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

    ################## SCORE PER-FIELD CONFIDENCE ##################
    # T2 (DAIP2-2692): rule-based confidence per HSC field. Combines
    # exact / token-overlap / fuzzy / number-match / regex / grounding /
    # na-check signals with field-type-aware weights. T3 will layer an
    # LLM verifier on top for gray-zone scores.
    from src.qc_qa.confidence import score_field

    for field_name, metadata in hybrid_smart_chunked_metadata_dict.items():
        meta = metadata if isinstance(metadata, dict) else {}
        # Pull the field's prompt for the na_check signal's keyword set.
        field_obj = (
            one_to_one_fields.get_field(field_name)
            if one_to_one_fields.contains_fields()
            else None
        )
        field_prompt = getattr(field_obj, "prompt", "") if field_obj else ""

        rule_score = score_field(
            field_name=field_name,
            value=one_to_one_results.get(field_name),
            metadata=meta,
            contract_text=contract_text,
            field_prompt=field_prompt,
        )
        one_to_one_results[f"{field_name}_CONF"] = rule_score

    ################## RUN CLIENT FULL-CONTEXT FIELDS ##################
    # Client-specific full_context fields (e.g. Clover's 12 audit/filing/CA fields)
    if client_full_context_fields is not None:
        with timing_utils.timed_block(
            f"one_to_one.{client_name}_full_context", context=filename
        ):
            fc_answers = vendor_run_full_context_fields(
                contract_text=contract_text,
                fields=client_full_context_fields,
                filename=filename,
                usage_prefix=client_name.upper(),
                constants=constants,
            )
        one_to_one_results.update(fc_answers)

    ################## DERIVE CLIENT-SPECIFIC FIELDS ##################
    # BCBS Promise: derive OFFSET_INDICATOR from OFFSET_TERM
    if "OFFSET_TERM" in one_to_one_results:
        offset_term = one_to_one_results.get("OFFSET_TERM", "N/A")
        if not string_utils.is_empty(offset_term):
            one_to_one_results["OFFSET_INDICATOR"] = "Y"
        else:
            one_to_one_results["OFFSET_INDICATOR"] = "N"

    ################## 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
        )

    ################## Split Dynamic Primary Entities ##################
    # If the DYNAMIC_PRIMARY_ENTITIES smart_chunked fallback was triggered (all four
    # dynamic primary fields were empty across every one_to_n row), the HSC step above
    # will have populated DYNAMIC_PRIMARY_ENTITIES in one_to_one_results.
    # Classify it into LOB/PROGRAM/PRODUCT/NETWORK before crosswalk runs so that
    # AARETE_DERIVED_LOB and AARETE_DERIVED_NETWORK are derived correctly.
    if not string_utils.is_empty(one_to_one_results.get("DYNAMIC_PRIMARY_ENTITIES")):
        split_results = one_to_n_funcs.split_dynamic_primary_entities(
            [one_to_one_results], filename
        )
        one_to_one_results = split_results[0]

    ################## Derive AARETE_DERIVED_LOB/NETWORK ##################
    one_to_n_funcs.map_lob_network_to_aarete_derived(
        [one_to_one_results], constants, filename
    )

    ################## 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
    )

    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 [], []

    instrumentation.log("page_start", filename=filename, page_num=str(page_num))
    with instrumentation_context.instr_scope(page_num=str(page_num)):
        ############################### 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:
            with instrumentation_context.instr_scope(chunk_id=chunk.chunk_id):
                before = len(reimbursement_level_answers)
                chunk_answers = one_to_n_funcs.reimbursement_level(
                    chunk.text, constants, filename
                )
                if chunk_answers:
                    reimbursement_level_answers.extend(chunk_answers)
                instrumentation.log(
                    "row_count",
                    filename=filename,
                    stage="reimbursement_level_per_chunk",
                    page_num=str(page_num),
                    chunk_id=chunk.chunk_id,
                    n_in=before,
                    n_out=len(reimbursement_level_answers),
                )

        if not reimbursement_level_answers:
            logging.debug(
                f"{datetime_str()} Page {page_num}: No reimbursement answers found, skipping - {filename}"
            )
            instrumentation.log("page_end", filename=filename, page_num=str(page_num))
            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 ################################
        _n_before_carveout = len(reimbursement_level_answers)
        reimbursement_level_answers, special_case_answers = (
            one_to_n_funcs.carveout_and_special_case(
                reimbursement_level_answers, constants, filename
            )
        )
        instrumentation.log(
            "row_count",
            filename=filename,
            stage="carveout",
            page_num=str(page_num),
            n_in=_n_before_carveout,
            n_out=len(reimbursement_level_answers),
            extra_json=f'{{"special_case_rows": {len(special_case_answers)}}}',
        )

        ################################ Dynamic Code Assignment ################################
        _n_before_dca = len(reimbursement_level_answers)
        reimbursement_level_answers = one_to_n_funcs.dynamic_code_assignment(
            reimbursement_level_answers, constants, filename
        )
        instrumentation.log(
            "row_count",
            filename=filename,
            stage="dynamic_code_assignment",
            page_num=str(page_num),
            n_in=_n_before_dca,
            n_out=len(reimbursement_level_answers),
        )

        ############################### Lesser of Distribution ###############################
        # Note: We run lesser_of without dynamic fields in Step 1
        _n_before_lesser = len(reimbursement_level_answers)
        reimbursement_level_answers = one_to_n_funcs.lesser_of_distribution(
            reimbursement_level_answers,
            exhibit_text_simplified,
            page_num,
            constants,
            filename,
            exhibit,
        )
        instrumentation.log(
            "row_count",
            filename=filename,
            stage="lesser_of_distribution",
            page_num=str(page_num),
            n_in=_n_before_lesser,
            n_out=len(reimbursement_level_answers),
        )

        ################################ Get Breakouts ###############################
        _n_before_breakout = len(reimbursement_level_answers)
        reimbursement_level_answers, special_case_answers = one_to_n_funcs.breakout(
            reimbursement_level_answers,
            special_case_answers,
            filename,
            constants,
        )
        instrumentation.log(
            "row_count",
            filename=filename,
            stage="breakout",
            page_num=str(page_num),
            n_in=_n_before_breakout,
            n_out=len(reimbursement_level_answers),
        )

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

        instrumentation.log("page_end", filename=filename, page_num=str(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 i, exhibit in enumerate(exhibits):
        instrumentation.log(
            "exhibit_start",
            filename=filename,
            exhibit_idx=i,
            exhibit_header=(
                exhibit.exhibit_header[:80] if exhibit.exhibit_header else ""
            ),
            exhibit_page=exhibit.exhibit_page,
            extra_json=f'{{"num_pages": {len(exhibit.exhibit_page_nums)}}}',
        )
        with instrumentation_context.instr_scope(
            exhibit_idx=i,
            exhibit_header=(
                exhibit.exhibit_header[:80] if exhibit.exhibit_header else ""
            ),
            exhibit_page=exhibit.exhibit_page,
        ):
            # 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 {i + 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 = {
                    instrumentation_context.submit_with_context(
                        executor,
                        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}"
                )
                instrumentation.log(
                    "exhibit_gate_skip",
                    filename=filename,
                    exhibit_idx=i,
                    exhibit_page=exhibit.exhibit_page,
                    extra_json='{"reason": "no_reimbursements_after_step1"}',
                )
                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
            instrumentation.log(
                "stage_transition",
                filename=filename,
                stage="step2_exhibit_level_start",
                exhibit_idx=i,
            )
            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
            )
            instrumentation.log(
                "stage_transition",
                filename=filename,
                stage="step2_exhibit_level_done",
                exhibit_idx=i,
            )
            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
                instrumentation.log(
                    "stage_transition",
                    filename=filename,
                    stage="step3_start",
                    exhibit_idx=i,
                    n_in=len(exhibit.reimbursement_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,
                )
                instrumentation.log(
                    "row_count",
                    filename=filename,
                    stage="combine_one_to_n",
                    exhibit_idx=i,
                    n_in=len(reimbursement_rows_with_dynamic),
                    n_out=len(combined_rows),
                )

                # Run cleaning
                combined_rows = one_to_n_funcs.one_to_n_cleaning(
                    combined_rows, exhibit.exhibit_text, constants, filename
                )
                instrumentation.log(
                    "row_count",
                    filename=filename,
                    stage="one_to_n_cleaning",
                    exhibit_idx=i,
                    n_in=len(combined_rows),
                    n_out=len(combined_rows),
                )

                exhibit.final_rows = combined_rows
                one_to_n_results.extend(combined_rows)
                instrumentation.log(
                    "stage_transition",
                    filename=filename,
                    stage="step3_end",
                    exhibit_idx=i,
                    n_out=len(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