doczyai-pipelines/src/pipelines/clients/bcbs_promise/file_processing.py

import concurrent.futures
import logging
from typing import TYPE_CHECKING, 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.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
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.info(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)
        text_dict, header, footer = preprocess.find_headers_and_footers(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_chunk_mapping, all_exhibit_headers = (
                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 {len(exhibit_chunk_mapping)} exhibits - {filename}"
            )
            with timing_utils.timed_block("one_to_n_extraction", context=filename):
                (
                    one_to_n_results,
                    dynamic_one_to_one_fields,
                    first_reimbursement_page,
                ) = run_one_to_n_prompts(
                    pages_dict=pages_dict,
                    text_dict=text_dict,
                    exhibit_chunk_mapping=exhibit_chunk_mapping,
                    all_exhibit_headers=all_exhibit_headers,
                    constants=constants,
                    filename=filename,
                )
            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.info(f"{datetime_str()} One to N Complete - {filename}")
        else:
            first_reimbursement_page = "1"
            logging.info(
                f"{datetime_str()} No Reimbursement Found, Skipping - {filename}"
            )

        final_results = one_to_n_results  # Set as final results
    else:
        first_reimbursement_page = "1"
        logging.info(
            f"{datetime_str()} Fields not configured for One to N, Skipping - {filename}"
        )

    # ONE TO ONE PROCESSING
    if process_one_to_one:
        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,
                top_sheet_dict,
                dynamic_one_to_one_fields,
                first_reimbursement_page,
                constants,
            )
        one_to_one_results["FILE_NAME"] = filename
        logging.info(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.info(
            f"{datetime_str()} Fields not configured for One to One, Skipping - {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.info(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.info(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.info(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],
    top_sheet_dict,
    dynamic_one_to_one_fields: FieldSet,
    first_reimbursement_page: str,
    constants: Constants,
):

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

    ################## INITIALIZE BCBS PROMISE FIELDS ##################
    bcbs_promise_fields = FieldSet(
        relationship="one_to_one", file_path=config.BCBS_PROMISE_PROMPTS_PATH
    )
    # Combine bcbs_promise fields with one_to_one_fields for HSC processing
    hsc_fields = one_to_one_fields.combine(bcbs_promise_fields)

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

    ################## RUN HYBRID SMART CHUNKED PROMPTS ##################
    # RAG function loads retrieval questions internally
    # hsc_fields includes both investment_prompts.json AND bcbs_promise_prompts.json 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(
                hsc_fields, constants, contract_text, filename, text_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

    one_to_one_results = tin_npi_funcs.merge_provider_info_with_one_to_one(
        one_to_one_results, filename
    )

    ################## DERIVE BCBS PROMISE FIELDS ##################
    # Derive OFFSET_INDICATOR from OFFSET_TERM
    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"

    ################## ADD AD FIELDS ################
    one_to_one_results = one_to_one_funcs.get_aarete_derived_dates(
        one_to_one_results, text_dict, filename
    )

    ################## ADD SIGNATURE COUNT ##################
    one_to_one_results = one_to_one_funcs.add_signature_count(
        one_to_one_results, contract_text
    )

    ################## Crosswalk Fields ##################
    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,
    pages_dict: Optional[dict[str, "Page"]] = None,
    text_dict: Optional[dict[str, str]] = None,
    exhibit_page_nums: Optional[list[str]] = None,
    constants: Optional[Constants] = None,
    filename: Optional[str] = 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

    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
    """
    # Get page text - prefer using exhibit.get_page_text() if available
    if exhibit.pages is not None:
        page_text = exhibit.get_page_text(page_num)
        exhibit_page_nums = exhibit.exhibit_page_nums
    elif pages_dict is not None:
        # Handle sub-pages: "27.1" means get sub-page "1" from page "27"
        if "." in page_num and not page_num.startswith("."):
            base_page, sub_page = page_num.rsplit(".", 1)
            if base_page in pages_dict:
                page_text = pages_dict[base_page].get_text(sub_page)
            else:
                page_text = ""
        elif page_num in pages_dict:
            page_text = pages_dict[page_num].get_text()
        else:
            page_text = ""
        exhibit_page_nums = (
            exhibit.exhibit_page_nums
            if exhibit_page_nums is None
            else exhibit_page_nums
        )
    elif text_dict is not None:
        page_text = text_dict.get(page_num, "")
        exhibit_page_nums = (
            exhibit_page_nums
            if exhibit_page_nums is not None
            else exhibit.exhibit_page_nums
        )
    else:
        raise ValueError(
            "Either pages_dict, text_dict, or exhibit.pages must be provided"
        )

    # Simplify exhibit text - prefer using exhibit object
    exhibit_text_simplified = preprocessing_funcs.simplify_exhibit(
        exhibit=exhibit,
        current_page_num=page_num,
    )

    ############################### Reimbursement Primary ###############################
    reimbursement_level_answers = one_to_n_funcs.reimbursement_level(
        page_text, constants, filename
    )

    if not reimbursement_level_answers:
        return [], []

    ################################ Carveouts and Special Case ################################
    reimbursement_level_answers, special_case_answers = (
        one_to_n_funcs.carveout_and_special_case(
            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,  # ← Added this parameter
    )

    ################################ 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_chunk_mapping: Optional[dict[str, list[str]]] = None,
    all_exhibit_headers: Optional[dict[str, str]] = None,
    constants: Optional[Constants] = None,
    filename: Optional[str] = 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_chunk_mapping: Dictionary mapping exhibit pages to lists of page numbers
        all_exhibit_headers: Dictionary mapping exhibit pages to their headers
        constants: Constants object
        filename: Name of the file being processed
    """
    if exhibit_chunk_mapping is None:
        exhibit_chunk_mapping = {}
    if all_exhibit_headers is None:
        all_exhibit_headers = {}

    total_exhibits = len(exhibit_chunk_mapping)
    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)
    exhibits = exhibit_funcs.get_exhibit_list(
        pages_dict=pages_dict,
        text_dict=text_dict,
        exhibit_chunk_mapping=exhibit_chunk_mapping,
        all_exhibit_headers=all_exhibit_headers,
    )  # <- Returns list[Exhibit]

    # Process exhibits serially; within each exhibit, process pages in parallel
    total_pages = sum(len(exhibit.exhibit_page_nums) for exhibit in exhibits)
    completed_pages = 0
    one_to_n_results = []
    first_reimbursement_page = "1"

    for exhibit in exhibits:
        exhibit_pages = exhibit.exhibit_page_nums
        if exhibit_pages:
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=min(len(exhibit_pages), 20)
            ) as executor:
                page_futures = {
                    executor.submit(
                        process_page_reimbursements,
                        exhibit,
                        page_num,
                        pages_dict,
                        text_dict,
                        exhibit_pages,
                        constants,
                        filename,
                    ): page_num
                    for page_num in exhibit_pages
                }

                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
                        )
                        completed_pages += 1
                        if completed_pages % 20 == 0 or completed_pages == total_pages:
                            logging.debug(
                                f"{datetime_str()} Page progress: {completed_pages}/{total_pages} - {filename}"
                            )
                    except Exception as e:
                        logging.error(f"Error processing page: {str(e)}")

        if not exhibit.has_reimbursements:
            continue

        # 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,
            )
        )
        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}"
        )

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

        # Track first reimbursement page
        if first_reimbursement_page == "1" and exhibit.has_reimbursements:
            first_reimbursement_page = exhibit.exhibit_page

    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, first_reimbursement_page