doczyai-pipelines/src/pipelines/runner.py

"""
Common Pipeline Runner

Handles shared orchestration (DTC, QC/QA, cache warming) and routes
to client-specific file_processing implementations.

Usage:
    from src.pipelines.runner import main
    main(client="saas")  # Run default SaaS pipeline
    main(client="clover")  # Run Clover client pipeline
"""

import concurrent.futures
import logging
import os
import traceback
import warnings
from datetime import datetime

import pandas as pd

# Suppress Pydantic protected namespace warning from langchain-aws
warnings.filterwarnings("ignore", message=".*protected namespace.*")

import src.prompts.prompt_templates as prompt_templates
import src.utils.io_utils as io_utils
import src.utils.llm_utils as llm_utils
import src.utils.logging_utils as logging_utils
import src.utils.usage_tracking as usage_tracking
from src.pipelines.runtime_context import set_active_client
from src.constants.constants import Constants
from src import config
from src.core.registry import get_registry
from src.parent_child import main as parent_child_main
from src.document_classification import main as dtc_main
from src.qc_qa.pipeline import (
    generate_statistics,
    run_qc_qa_pipeline,
    save_qc_qa_outputs,
)


def get_file_processing_module(client: str):
    """Get the file_processing module (or processor instance) for a client or vendor.

    Vendors are always routed through the generic VendorProcessor, which is
    driven by vendor_fields.json + per-vendor config.yaml. No per-vendor
    file_processing.py is required or consulted for vendor pipelines.

    Args:
        client: Client/vendor name ("saas", "clover", "la_care", etc.)

    Returns:
        An object with a process_file(file_object, constants, run_timestamp) method.
    """
    registry = get_registry()
    if registry.is_vendor(client):
        from src.pipelines.vendors.shared.generic_processor import VendorProcessor

        return VendorProcessor(client)

    # Set active client so shared resolver can select function-level overrides.
    set_active_client(client)
    from src.pipelines.shared import file_processing

    return file_processing


def safe_process_file(item, constants, run_timestamp, file_processing):
    """Process a single file with comprehensive error handling.

    Args:
        item: Tuple of (filename, contract_text) representing the file to process
        constants: Constants object containing configuration and processing rules
        run_timestamp: Timestamp string for output file naming and tracking
        file_processing: The file_processing module to use

    Returns:
        tuple: Either:
            - (cc_df, dashboard_df) tuple with extracted field results (success case)
            - (error_df, error_df) tuple with error details (failure case)

    Note:
        Returns both CC and dashboard versions from file processing.
    """
    file_id = item[0] if isinstance(item, (list, tuple)) and len(item) > 0 else item
    try:
        cc_df, dashboard_df = file_processing.process_file(
            item, constants, run_timestamp
        )
        return cc_df, dashboard_df
    except Exception as e:
        error_type = type(e).__name__
        error_message = str(e)
        full_traceback = traceback.format_exc()

        logging.error(f"Error processing file {file_id}")
        logging.error(f"Error type: {error_type}")
        logging.error(f"Error Message: {error_message}")
        logging.error(f"Full traceback:\n{full_traceback}")

        error_df = pd.DataFrame(
            [
                {
                    "FILE_NAME": file_id,
                    "error": error_message.replace(",", ";").replace("\n", " "),
                    "error_type": error_type,
                    "traceback": full_traceback.replace(",", ";").replace("\n", " | "),
                }
            ]
        )
        return error_df, error_df  # Return tuple for consistency


def main(client: str = "saas", testing=False, test_params={}):
    """Main entry point for all clients and vendors.

    Args:
        client: Client or vendor name ("saas", "clover", "la_care", etc.)
        testing: Whether running in test mode
        test_params: Test parameters if testing=True
    """
    # Check if batch_id is specified
    if not config.BATCH_ID:
        raise ValueError(
            "batch_id is required. Please specify batch_id in your command"
        )

    # Check if max_workers is not negative
    if config.MAX_WORKERS < 0:
        raise ValueError(
            f"Invalid configuration: MAX_WORKERS must be non-negative (got {config.MAX_WORKERS})"
        )

    # Validate client/vendor
    registry = get_registry()
    if client != "saas" and not registry.has_client(client):
        raise ValueError(
            f"Unknown client/vendor: {client}. " f"Available: {registry.list_clients()}"
        )

    # Ensure dynamic fallback resolution targets the requested client.
    set_active_client(client)

    # Get client-specific file processing module
    file_processing = get_file_processing_module(client)
    logging.info(f"Using {client} pipeline")

    # Set up per-file logging
    logging_utils.setup_per_file_logging()

    # Windows paths cannot contain ':'; use '-' in time component
    run_timestamp = datetime.now().strftime(f"run_%Y%m%d_%H-%M_{config.BATCH_ID}")

    # Read Constants
    constants = Constants()

    # Read and process input
    if not testing:
        input_dict = io_utils.read_input()
        max_workers = config.MAX_WORKERS
    else:
        input_dict = io_utils.read_input(test_params["local_input_dir"])
        max_workers = test_params["max_workers"]
        if "input_files" in test_params:
            input_dict = {
                k: v for k, v in input_dict.items() if k in test_params["input_files"]
            }

    logging.debug(f"Total Input Files : {len(input_dict)}")

    # ========== COMMON: Document Type Classification ==========
    if config.PERFORM_DTC:
        logging.info("=" * 80)
        logging.info("DOCUMENT TYPE CLASSIFICATION MODE - Running DTC and exiting")
        logging.info("=" * 80)
        dtc_main.main(input_dict, run_timestamp)
        logging.info(
            "DTC classification complete. Exiting (investment pipeline not run)."
        )
        logging.info("=" * 80)
        return

    # ========== COMMON: Warm prompt caches ==========
    logging.info("Warming prompt caches before parallel processing...")
    try:
        cacheable_instructions = prompt_templates.get_cacheable_instructions()
        for cache_key, instruction in cacheable_instructions.items():
            llm_utils.warm_prompt_cache(
                instruction=instruction, cache_key=cache_key, model_id="sonnet_latest"
            )
        logging.debug(
            f"Successfully warmed {len(cacheable_instructions)} prompt caches"
        )
    except Exception as e:
        logging.warning(f"Error warming caches (will proceed anyway): {str(e)}")

    # ========== CLIENT-SPECIFIC: File Processing ==========
    successful_results_cc = []
    successful_results_dashboard = []
    error_results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(
                safe_process_file, item, constants, run_timestamp, file_processing
            )
            for item in input_dict.items()
        ]

        for future in concurrent.futures.as_completed(futures):
            try:
                cc_result, dashboard_result = future.result()
                if "error" in cc_result.columns:
                    error_results.append(cc_result)
                else:
                    successful_results_cc.append(cc_result)
                    if (
                        config.RUN_DASHBOARD_POSTPROCESSING
                        and dashboard_result is not None
                    ):
                        successful_results_dashboard.append(dashboard_result)
            except Exception as e:
                logging.error(f"Error processing future: {str(e)}")
                error_df = pd.DataFrame([{"error": str(e)}])
                error_results.append(error_df)

    # Combine results
    if len(successful_results_cc) > 0:
        FINAL_RESULT_DF_CC = pd.concat(successful_results_cc, ignore_index=True)
        if successful_results_dashboard:
            FINAL_RESULT_DF_DASHBOARD = pd.concat(
                successful_results_dashboard, ignore_index=True
            )
        else:
            FINAL_RESULT_DF_DASHBOARD = pd.DataFrame()
    else:
        FINAL_RESULT_DF_CC = pd.DataFrame()
        FINAL_RESULT_DF_DASHBOARD = pd.DataFrame()

    if len(error_results) > 0:
        ERROR_RESULT_DF = pd.concat(error_results, ignore_index=True)
    else:
        ERROR_RESULT_DF = pd.DataFrame()

    # ========== COMMON: QC/QA Validation ==========
    # Run QC/QA validation by default (preserves automatic behavior for single files and batches)
    # QC/QA runs on CC version only (dashboard version and error files do not need QC/QA)
    validated_cc_df = None
    qc_qa_stats_df = None

    # Run QC/QA on CC version if we have successful results
    if not FINAL_RESULT_DF_CC.empty:
        logging.info("=" * 80)
        logging.info("Running QC/QA Validation Pipeline on CC results...")
        logging.info("=" * 80)
        try:
            # Ensure unique index before QC/QA (safety check)
            if FINAL_RESULT_DF_CC.index.has_duplicates:
                logging.warning(
                    "Detected duplicate indices in FINAL_RESULT_DF_CC, resetting index..."
                )
                FINAL_RESULT_DF_CC = FINAL_RESULT_DF_CC.reset_index(drop=True)

            validated_cc_df = run_qc_qa_pipeline(
                FINAL_RESULT_DF_CC.copy(), stats_df=None
            )
            logging.debug("QC/QA validation on CC results completed successfully")

            # Generate validation statistics from CC version
            logging.debug("Generating QC/QA statistics...")
            qc_qa_stats_df = generate_statistics(validated_cc_df)
            logging.debug("QC/QA statistics generated successfully")
        except Exception as e:
            logging.error(f"QC/QA validation on CC results failed: {str(e)}")
            logging.error(f"Full traceback:\n{traceback.format_exc()}")
            logging.warning("Continuing with original unvalidated CC results...")
        logging.info("=" * 80)

    # ========== COMMON: Output & Parent-Child ==========
    if not testing:
        per_file_usage_df, batch_summary_df = usage_tracking.get_usage_dataframes(
            config.BATCH_ID, run_timestamp
        )

        if config.WRITE_TO_S3:
            # Write CC version
            doczy_output_for_pc = io_utils.write_s3(
                FINAL_RESULT_DF_CC, "", run_timestamp, "cc_results_full"
            )
            # Write dashboard version (only when dashboard postprocessing is enabled)
            if (
                config.RUN_DASHBOARD_POSTPROCESSING
                and not FINAL_RESULT_DF_DASHBOARD.empty
            ):
                io_utils.write_s3(
                    FINAL_RESULT_DF_DASHBOARD,
                    "",
                    run_timestamp,
                    "dashboard_results_full",
                )
            # Write error file
            if not ERROR_RESULT_DF.empty:
                io_utils.write_s3(ERROR_RESULT_DF, "", run_timestamp, "error")

            # Write QC/QA outputs (only for CC version, not dashboard)
            if validated_cc_df is not None:
                io_utils.write_s3(validated_cc_df, "", run_timestamp, "qc_qa_cc_full")
            if qc_qa_stats_df is not None:
                io_utils.write_s3(qc_qa_stats_df, "", run_timestamp, "qc_qa_stats")

            if not per_file_usage_df.empty and not batch_summary_df.empty:
                logging.debug("Exporting usage and cost tracking data to S3...")
                io_utils.write_s3(per_file_usage_df, "", run_timestamp, "usage")
                io_utils.write_s3(batch_summary_df, "", run_timestamp, "usage_summary")
        else:
            # Write CC version
            doczy_output_for_pc = io_utils.write_local(
                FINAL_RESULT_DF_CC, "", run_timestamp, "cc_results_full"
            )
            # Write dashboard version (only when dashboard postprocessing is enabled)
            if (
                config.RUN_DASHBOARD_POSTPROCESSING
                and not FINAL_RESULT_DF_DASHBOARD.empty
            ):
                io_utils.write_local(
                    FINAL_RESULT_DF_DASHBOARD,
                    "",
                    run_timestamp,
                    "dashboard_results_full",
                )
            # Write error file
            if not ERROR_RESULT_DF.empty:
                io_utils.write_local(ERROR_RESULT_DF, "", run_timestamp, "error")

            # Write QC/QA outputs (only for CC version, not dashboard)
            if validated_cc_df is not None:
                io_utils.write_local(
                    validated_cc_df, "", run_timestamp, "qc_qa_cc_full"
                )
            if qc_qa_stats_df is not None:
                io_utils.write_local(qc_qa_stats_df, "", run_timestamp, "qc_qa_stats")

        if config.PERFORM_PARENT_CHILD_MAPPING:
            if FINAL_RESULT_DF_CC.empty:
                logging.warning(
                    "Skipping Parent-Child mapping: No successful results to process (FINAL_RESULT_DF_CC is empty)"
                )
            else:
                logging.debug(
                    f"Starting Parent-Child mapping with {len(FINAL_RESULT_DF_CC)} records from: {doczy_output_for_pc}"
                )
                pc_output_dir = os.path.join(
                    config.CONSOLIDATED_OUTPUT_DIRECTORY,
                    run_timestamp,
                    "parent-child",
                )
                _, pc_local_path = parent_child_main.main(
                    doczy_output_for_pc,
                    client=client,
                    output_dir=pc_output_dir,
                )
                if config.WRITE_TO_S3 and pc_local_path:
                    io_utils.upload_local_file_to_s3(
                        pc_local_path, run_timestamp, "parent-child"
                    )

        if not per_file_usage_df.empty and not batch_summary_df.empty:
            logging.debug("Exporting usage and cost tracking data locally...")
            io_utils.write_local(per_file_usage_df, "", run_timestamp, "usage")
            io_utils.write_local(batch_summary_df, "", run_timestamp, "usage_summary")
    else:
        return FINAL_RESULT_DF_CC, FINAL_RESULT_DF_DASHBOARD, ERROR_RESULT_DF


if __name__ == "__main__":
    # Default to saas, can be overridden via config or CLI arg
    client_arg = getattr(config, "CLIENT", ["saas"])
    # CLIENT is a list from config, get first element or default to "saas"
    client = (
        client_arg[0]
        if client_arg and isinstance(client_arg, list)
        else client_arg or "saas"
    )
    main(client=client)