doczyai-pipelines/src/utils/io_utils.py

import json
import logging
import os
import re
import tempfile
from io import StringIO
import json
from threading import Lock

import pandas as pd
import src.utils.tracking_utils as tracking_utils
from botocore.exceptions import ClientError
from pyxlsb import open_workbook
from src import config
from pathlib import Path
from typing import Optional

# Thread-safe lock for JSON file operations
json_write_lock = Lock()

RUN_FOLDER_PATTERN = re.compile(r"^run_(\d{8})_(\d{2})-(\d{2})")


def read_local(file_path) -> str | pd.DataFrame | None:
    """Reads a local (.txt, .csv, or .xlsx) file and returns its contents.

    Args:
        file_path (str): The path to the local file.

    Returns:
        str or pd.DataFrame: The contents of the file, either as a string (for text files)
        or as a pandas DataFrame (for CSV/XLSX files).
    """
    # Check if the file is a text file
    if os.path.isfile(file_path) and file_path.endswith(".txt"):
        try:
            # First attempt to open the file with UTF-8 encoding
            with open(file_path, "r", encoding="utf-8") as file:
                file_contents = file.read()
                return file_contents
        except UnicodeDecodeError:
            # If UTF-8 fails, try reading the file with ANSI encoding
            try:
                with open(file_path, "r", encoding="cp1252") as file:
                    file_contents = file.read()
                    return file_contents
            except UnicodeDecodeError:
                # If ANSI also fails, log an error message or handle it accordingly
                logging.error(
                    f"Failed to decode {file_path} with UTF-8 and cp1252 encodings."
                )
                return None

    if os.path.isfile(file_path) and file_path.endswith(".csv"):
        try:
            df = pd.read_csv(file_path)
            return df
        except Exception:
            logging.error("Failed to read .csv file.")
            return None

    if os.path.isfile(file_path) and file_path.endswith(".xlsx"):
        try:
            df = pd.read_excel(file_path)
            return df
        except Exception:
            logging.error("Failed to read .xlsx file.")
            return None

    # Unsupported file type
    logging.error(f"Unsupported file type for {file_path}.")
    return None


def list_s3_files(prefix=config.S3_PREFIX, bucket=config.S3_BUCKET):  # io_utils.py
    """List all files in an S3 folder."""
    response = config.S3_CLIENT.list_objects_v2(Bucket=bucket, Prefix=prefix)
    return [item["Key"] for item in response.get("Contents", [])]


def list_s3_files_paginated(
    prefix=config.S3_PREFIX, bucket=config.S3_BUCKET
) -> list[str]:
    """List all S3 objects under a prefix using pagination."""
    all_keys = []
    continuation_token = None

    while True:
        kwargs = {"Bucket": bucket, "Prefix": prefix}
        if continuation_token:
            kwargs["ContinuationToken"] = continuation_token

        response = config.S3_CLIENT.list_objects_v2(**kwargs)
        all_keys.extend([item["Key"] for item in response.get("Contents", [])])

        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    return all_keys


def list_s3_child_prefixes(
    prefix=config.S3_PREFIX, bucket=config.S3_BUCKET
) -> list[str]:
    """List immediate child prefixes under an S3 prefix using Delimiter='/' pagination."""
    child_prefixes = []
    continuation_token = None

    while True:
        kwargs = {"Bucket": bucket, "Prefix": prefix, "Delimiter": "/"}
        if continuation_token:
            kwargs["ContinuationToken"] = continuation_token

        response = config.S3_CLIENT.list_objects_v2(**kwargs)
        child_prefixes.extend([p["Prefix"] for p in response.get("CommonPrefixes", [])])

        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    return child_prefixes


def extract_matching_top_level_s3_prefixes(
    match_substring: str, bucket=config.S3_BUCKET
) -> list[str]:
    """Return top-level S3 prefixes that contain the provided substring."""
    top_level_prefixes = list_s3_child_prefixes(prefix="", bucket=bucket)
    matching_prefixes = {
        prefix.rstrip("/")
        for prefix in top_level_prefixes
        if match_substring in prefix.rstrip("/")
    }
    return sorted(matching_prefixes)


def get_latest_run_folder(batch_folder: str, bucket=config.S3_BUCKET) -> str | None:
    """Find latest run_YYYYMMDD_HH-MM... folder under a batch folder."""
    prefix = f"{batch_folder}/"
    run_prefixes = list_s3_child_prefixes(prefix=prefix, bucket=bucket)
    run_candidates = {
        rp[len(prefix) :].rstrip("/")
        for rp in run_prefixes
        if rp[len(prefix) :].startswith("run_")
    }

    if not run_candidates:
        return None

    def sort_key(run_folder: str):
        match = RUN_FOLDER_PATTERN.match(run_folder)
        if not match:
            return ("", "", "", run_folder)
        date_part, hour_part, minute_part = match.groups()
        return (date_part, hour_part, minute_part, run_folder)

    return sorted(run_candidates, key=sort_key, reverse=True)[0]


def find_usage_key(
    batch_folder: str, run_folder: str, bucket=config.S3_BUCKET
) -> str | None:
    """Find the latest -USAGE.csv key in tracking/ for a batch run."""
    tracking_prefix = f"{batch_folder}/{run_folder}/tracking/"
    tracking_keys = list_s3_files_paginated(prefix=tracking_prefix, bucket=bucket)
    usage_keys = sorted([key for key in tracking_keys if key.endswith("-USAGE.csv")])
    if not usage_keys:
        return None
    return usage_keys[-1]


def count_s3_files_under_prefix(
    prefix: str, bucket=config.S3_BUCKET, include_folders: bool = False
) -> int:
    """Count S3 objects under a prefix, optionally excluding folder placeholders."""
    count = 0
    continuation_token = None

    while True:
        kwargs = {"Bucket": bucket, "Prefix": prefix}
        if continuation_token:
            kwargs["ContinuationToken"] = continuation_token

        response = config.S3_CLIENT.list_objects_v2(**kwargs)

        if include_folders:
            count += len(response.get("Contents", []))
        else:
            count += len(
                [
                    item["Key"]
                    for item in response.get("Contents", [])
                    if not item["Key"].endswith("/")
                ]
            )

        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    return count


def read_s3_pdf(pdf_path):
    """
    Downloads a PDF file from an S3 bucket and returns a temporary file path.

    Args:
        pdf_path (str): The S3 URL of the PDF file (e.g., s3://bucket-name/path/to/file.pdf).

    Returns:
        str: The local temporary file path of the downloaded PDF.
    """

    s3_client = config.S3_CLIENT
    s3_parts = pdf_path[5:].split("/", 1)
    bucket = s3_parts[0]
    prefix = s3_parts[1] if len(s3_parts) > 1 else ""

    try:
        response = s3_client.get_object(Bucket=bucket, Key=prefix)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
            temp_pdf.write(response["Body"].read())
            temp_pdf_path = temp_pdf.name
        return temp_pdf_path
    except ClientError as e:
        logging.error(f"Error fetching PDF from S3: {e}")
        return None


def read_s3() -> dict:
    """Reads all .txt files from a specific S3 bucket and prefix.

    Returns:
        dict: A dictionary containing the contents of the files, keyed by filename.

    Notes:
        - Doesn't take input arguments; instead, it uses configurations from the config
            module for S3 client, bucket, and prefix.
        - Only reads files with a .txt extension.
    """
    s3_client = config.S3_CLIENT
    bucket = config.S3_BUCKET
    prefix = config.S3_PREFIX

    file_list = []
    continuation_token = None

    while True:

        if continuation_token:
            response = s3_client.list_objects_v2(
                Bucket=bucket, Prefix=prefix, ContinuationToken=continuation_token
            )
        else:
            response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)

        file_list += [
            obj["Key"]
            for obj in response.get("Contents", [])
            if not obj["Key"].endswith("/")
        ]

        if response.get(
            "IsTruncated", False
        ):  # default to False to avoid infinite loop
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    contract_list = sorted(file_list)
    files = {}

    for contract in contract_list:
        logging.info(f"Reading {contract}")
        try:
            data = s3_client.get_object(Bucket=bucket, Key=contract)
            if contract.endswith(".txt"):
                contents = data["Body"].read()
                context = contents.decode("utf-8")
                path, filename = os.path.split(contract)
                files[filename] = context
        except Exception as e:
            logging.warning(f"Failed to read {contract}: {e}")
            pass

    return files


def read_s3_csv(csv_path) -> pd.DataFrame | None:
    """Reads a CSV file from an S3 bucket and returns it as a DataFrame.

    Args:
        csv_path (str): The S3 URL of the CSV file (e.g., s3://bucket-name/path/to/file.csv).

    Returns:
        pd.DataFrame | None: The contents of the CSV file as a pandas DataFrame,
        or None if an error occurs.
    """
    s3_client = config.S3_CLIENT

    # Remove 's3://' prefix and split into bucket and key
    if not csv_path.startswith("s3://"):
        logging.error(f"Invalid S3 path: {csv_path}")
        return None

    s3_parts = csv_path[5:].split("/", 1)  # Remove 's3://' prefix with magic number 5
    bucket = s3_parts[0]
    prefix = s3_parts[1] if len(s3_parts) > 1 else ""
    try:
        response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)
        for obj in response.get(
            "Contents", []
        ):  # If "Contents" key is missing, default to empty list
            key = obj["Key"]
            if key.endswith(".csv"):
                csv_object = s3_client.get_object(Bucket=bucket, Key=key)
                try:
                    df = pd.read_csv(
                        StringIO(csv_object["Body"].read().decode("utf-8"))
                    )
                    return df
                except Exception as e:
                    logging.error(f"Error parsing CSV file {key}: {e}")
                    continue
    except ClientError as e:
        logging.error(f"Error listing S3 objects: {e}")
        return None

    # If no CSV files were found
    logging.warning(f"No CSV files found at {csv_path}")
    return None


def read_input(local_path=config.LOCAL_PATH) -> dict:
    """Reads input files from the specified local path or S3 bucket.

    Args:
        local_path (str, optional): The local path to read files from. Defaults to config.LOCAL_PATH.
            CAUTION: If config.READ_MODE is set to "s3", this parameter is ignored.

    Returns:
        dict: A dictionary containing the contents of the input files.

    Notes:
        - This function relies on the global configuration setting config.READ_MODE
        - If config.READ_MODE is "local", it reads all files from the specified local path.
        - If config.READ_MODE is "s3", it reads files from the configured S3 bucket and prefix.
    """
    if config.READ_MODE == "local":
        files = {}
        try:
            for file in os.listdir(local_path):
                full_path = os.path.join(local_path, file)
                if os.path.isfile(full_path):  # Only process files and not directories
                    file_text = read_local(full_path)
                    files[file] = file_text
            return files
        except (FileNotFoundError, PermissionError, NotADirectoryError) as e:
            logging.error(f"Error accessing local path {local_path}: {e}")
            return {}
    elif config.READ_MODE == "s3":
        return read_s3()
    else:
        logging.error(f"Invalid READ_MODE: {config.READ_MODE}")
        return {}


def read_input_csv():  # TODO: This function is unreferenced outside of unit tests.  Consider removal.
    if config.DF_READ_MODE == "local":
        ac_df = read_local(config.AC_DF) if config.AC_DF else None
        b_df = read_local(config.B_DF) if config.B_DF else None
    elif config.DF_READ_MODE == "s3":
        ac_df = read_s3_csv(config.AC_DF) if config.AC_DF else None
        b_df = read_s3_csv(config.B_DF) if config.B_DF else None

    return ac_df, b_df


def filter_already_processed(
    input_dict, input_dict_b
):  # TODO: this function is unreferenced outside of unit tests.  Consider removal.
    """Filter out already processed files based on either local files or S3 master tracking"""
    if config.WRITE_TO_S3:
        # Get processed files from S3 master tracking for this specific batch
        already_processed_ac, already_processed_b = tracking_utils.get_processed_files(
            config.BATCH_ID
        )

        # Keep files that still need either AC or B processing
        input_dict_ac = {
            k: v for k, v in input_dict.items() if k not in already_processed_ac
        }

        input_dict_b = {
            k: v for k, v in input_dict_b.items() if k not in already_processed_b
        }

        logging.info(
            f"Filtering based on S3 master tracking for batch {config.BATCH_ID}:"
        )
        logging.info(
            f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining"
        )
        logging.info(
            f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining"
        )

    else:
        # Original local directory checking logic
        already_processed_ac, already_processed_b = [], []
        for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
            folder_path = os.path.join(config.OUTPUT_DIRECTORY, folder_name)
            if os.path.isdir(folder_path):
                if config.AC_RESULTS_NAME in os.listdir(folder_path):
                    already_processed_ac.append(folder_name + ".txt")
                if config.B_RESULTS_NAME in os.listdir(folder_path):
                    already_processed_b.append(folder_name + ".txt")

        input_dict_ac = {
            k: v for k, v in input_dict.items() if k not in already_processed_ac
        }

        input_dict_b = {
            k: v for k, v in input_dict_b.items() if k not in already_processed_b
        }

        logging.info(f"Filtering based on local output directory:")
        logging.info(
            f"AC: Found {len(already_processed_ac)} already processed, {len(input_dict_ac)} remaining"
        )
        logging.info(
            f"B: Found {len(already_processed_b)} already processed, {len(input_dict_b)} remaining"
        )

    return input_dict_ac, input_dict_b


def remove_txt_extension(filename: str) -> str:
    """Remove the .txt extension from a filename.

    Args:
        filename (str): The filename from which to remove the extension.

    Returns:
        str: The filename, with .txt extension removed if it was present.
    """
    if isinstance(filename, str) and filename.lower().endswith(".txt"):
        return filename[:-4]
    return filename


def read_xlsb(
    path,
):  # TODO: this function is unreferenced outside of unit tests.  Consider removal.
    with open_workbook(path) as wb:
        with wb.get_sheet(1) as sheet:
            rows = []
            for row in sheet.rows():
                rows.append([item.v for item in row])  # Extracting cell values
            df = pd.DataFrame(
                rows[1:], columns=rows[0]
            )  # Converting to DataFrame, assuming first row is header
    return df


def split_dataframe_by_filename(
    df: pd.DataFrame, max_rows_per_split: int = 70000
) -> list[pd.DataFrame]:
    """
    Split a DataFrame into multiple DataFrames based on filename groups.

    This function ensures that all rows from the same filename stay together.
    It sorts filenames by row count (ascending) and creates splits that are as
    close to max_rows_per_split as possible without breaking up filename groups.

    Args:
        df (pd.DataFrame): The DataFrame to split. Must have a 'FILE_NAME' column.
        max_rows_per_split (int): Maximum number of rows per split file. Default: 70000.

    Returns:
        list[pd.DataFrame]: List of DataFrames, each representing a split file.

    Raises:
        ValueError: If DataFrame doesn't have 'FILE_NAME' column or is empty.
    """
    if df.empty:
        return [df]

    if "FILE_NAME" not in df.columns:
        raise ValueError("DataFrame must have a 'FILE_NAME' column for splitting")

    # Count rows per filename
    filename_counts = df["FILE_NAME"].value_counts().sort_values(ascending=True)
    filename_order = filename_counts.index.tolist()

    # Create list to store split DataFrames
    splits = []
    current_split_rows = []
    current_split_count = 0

    for filename in filename_order:
        filename_df = df[df["FILE_NAME"] == filename]
        filename_row_count = len(filename_df)

        # If adding this filename would exceed the limit, start a new split
        if (
            current_split_count + filename_row_count > max_rows_per_split
            and current_split_rows
        ):
            # Create and save current split
            current_split = pd.concat(current_split_rows, ignore_index=True)
            splits.append(current_split)
            # Start new split with current filename
            current_split_rows = [filename_df]
            current_split_count = filename_row_count
        else:
            # Add filename to current split
            current_split_rows.append(filename_df)
            current_split_count += filename_row_count

    # Add the last split if it has any rows
    if current_split_rows:
        final_split = pd.concat(current_split_rows, ignore_index=True)
        splits.append(final_split)

    return splits if splits else [df]


def _write_local_split_files(
    df: pd.DataFrame,
    output_dir: str,
    full_filename: str,
    split_filename_template: str,
    return_path_for_pc: bool = False,
) -> Optional[str]:
    """
    Helper function to write split files locally.

    Args:
        df: DataFrame to write (will be split if needed)
        output_dir: Full output directory path
        full_filename: Filename for single file (e.g., "{BATCH_ID}-RESULTS-FULL.csv")
        split_filename_template: Template for split files (e.g., "{BATCH_ID}-RESULTS-SPLIT{idx}.csv")
        return_path_for_pc: Whether to return path for parent-child mapping

    Returns:
        Full file path string if return_path_for_pc is True, None otherwise
    """
    os.makedirs(output_dir, exist_ok=True)
    splits = split_dataframe_by_filename(df, config.MAX_ROWS_PER_SPLIT)

    if len(splits) == 1:
        output_path = os.path.join(output_dir, full_filename)
        splits[0].to_csv(output_path, index=False, quoting=1)
        return output_path if return_path_for_pc else None
    else:
        # Multiple splits - write numbered files
        for idx, split_df in enumerate(splits, start=1):
            split_filename = split_filename_template.format(idx=idx)
            split_output_path = os.path.join(output_dir, split_filename)
            split_df.to_csv(split_output_path, index=False, quoting=1)
        # For parent-child mapping, use the first split file
        if return_path_for_pc:
            first_split_filename = split_filename_template.format(idx=1)
            return os.path.join(output_dir, first_split_filename)
        return None


def _write_s3_split_files(
    df: pd.DataFrame,
    run_timestamp: str,
    base_path: str,
    full_filename: str,
    split_filename_template: str,
    return_path_for_pc: bool = False,
) -> Optional[str]:
    """
    Helper function to write split files to S3.

    Args:
        df: DataFrame to write (will be split if needed)
        run_timestamp: Timestamp string for path
        base_path: Base S3 path (e.g., "full_outputs/cc_results")
        full_filename: Filename for single file (e.g., "{BATCH_ID}-RESULTS-FULL.csv")
        split_filename_template: Template for split files (e.g., "{BATCH_ID}-RESULTS-SPLIT{idx}.csv")
        return_path_for_pc: Whether to return path for parent-child mapping

    Returns:
        S3 path string if return_path_for_pc is True, None otherwise
    """
    splits = split_dataframe_by_filename(df, config.MAX_ROWS_PER_SPLIT)

    if len(splits) == 1:
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{base_path}/{full_filename}"
        csv_buffer = splits[0].to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
        logging.info(f"Saved {base_path.split('/')[-1]} full file: {output_path}")
        return output_path if return_path_for_pc else None
    else:
        # Multiple splits - upload each one
        for idx, split_df in enumerate(splits, start=1):
            split_filename = split_filename_template.format(idx=idx)
            split_output_path = (
                f"{config.BATCH_ID}/{run_timestamp}/{base_path}/{split_filename}"
            )
            csv_buffer = split_df.to_csv(index=False).encode()
            config.S3_CLIENT.put_object(
                Bucket=config.S3_OUTPUT_BUCKET, Key=split_output_path, Body=csv_buffer
            )
        logging.info(
            f"Saved {base_path.split('/')[-1]} split files ({len(splits)} files)"
        )
        if return_path_for_pc:
            first_split_filename = split_filename_template.format(idx=1)
            return (
                f"{config.BATCH_ID}/{run_timestamp}/{base_path}/{first_split_filename}"
            )
        return None


def write_local(
    df: pd.DataFrame, filename: str, run_timestamp: str, output_type: str
) -> Optional[str]:
    """
    Saves a DataFrame as a CSV file in a local directory.

    Args:
        df (pd.DataFrame): The DataFrame to save.
        filename (str): The base filename to use when saving the file.
        run_timestamp (str): The timestamp used to organize output files.
        output_type (str): Type of output, either "final", "parent_child" or "individual".

    Returns:
        - Output path if parent-child mapping is performed.
        - None, when this function is called for its side effect of uploading a file to S3.

    Notes:
        - Creates output directories if they do not exist
        - Uses different naming conventions based on the output_type
            - "final" output_type saves the file as "<BATCH_ID>-RESULTS.csv"
            in a timestamped directory under CONSOLIDATED_OUTPUT_DIRECTORY
            - "individual" output_type saves the file as "<filename>_-RESULTS.csv"
            in the OUTPUT_DIRECTORY
    """
    if output_type == "final":
        output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp)
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-RESULTS.csv"),
            index=False,
            quoting=1,
        )
        if config.PERFORM_PARENT_CHILD_MAPPING:
            return os.path.join(output_dir, f"{config.BATCH_ID}-RESULTS.csv")
        return None
    elif output_type == "individual":
        base_filename = os.path.splitext(filename)[0].strip()
        os.makedirs(config.OUTPUT_DIRECTORY, exist_ok=True)
        df.to_csv(
            os.path.join(config.OUTPUT_DIRECTORY, f"{base_filename}_-RESULTS.csv"),
            index=False,
            quoting=1,
        )
        if config.PERFORM_PARENT_CHILD_MAPPING:
            return os.path.join(
                config.OUTPUT_DIRECTORY, f"{base_filename}_-RESULTS.csv"
            )

        return None
    elif output_type == "error":
        # Error file goes in full_outputs/ folder
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "full_outputs"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-ERRORS.csv"),
            index=False,
            quoting=1,
        )
        if config.PERFORM_PARENT_CHILD_MAPPING:
            return os.path.join(output_dir, f"{config.BATCH_ID}-RESULTS.csv")
        return None
    elif output_type == "cc_results_full":
        # Full CC results (may be split into multiple files)
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY,
            run_timestamp,
            "full_outputs",
            "cc_results",
        )
        return _write_local_split_files(
            df=df,
            output_dir=output_dir,
            full_filename=f"{config.BATCH_ID}-RESULTS-FULL.csv",
            split_filename_template=f"{config.BATCH_ID}-RESULTS-SPLIT{{idx}}.csv",
            return_path_for_pc=config.PERFORM_PARENT_CHILD_MAPPING,
        )
    elif output_type == "dashboard_results_full":
        # Full dashboard results (may be split into multiple files)
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY,
            run_timestamp,
            "full_outputs",
            "dashboard_results",
        )
        _write_local_split_files(
            df=df,
            output_dir=output_dir,
            full_filename=f"{config.BATCH_ID}-RESULTS-dashboard.csv",
            split_filename_template=f"{config.BATCH_ID}-RESULTS-dashboard-SPLIT{{idx}}.csv",
            return_path_for_pc=False,
        )
        return None
    elif output_type == "qc_qa_cc_full":
        # QC/QA CC full results (may be split into multiple files)
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "automation_qa-qc"
        )
        _write_local_split_files(
            df=df,
            output_dir=output_dir,
            full_filename=f"{config.BATCH_ID}-RESULTS-QC-QA-FULL.csv",
            split_filename_template=f"{config.BATCH_ID}-RESULTS-QC-QA-SPLIT{{idx}}.csv",
            return_path_for_pc=False,
        )
        return None
    elif output_type == "qc_qa_error":
        # QC/QA error results
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "automation_qa-qc"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-ERRORS-QC-QA.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "parent_child":
        # Parent-child output
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "parent-child"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-PC.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "usage":
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "tracking"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-USAGE.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "usage_summary":
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "tracking"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-USAGE-SUMMARY.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "dtc":
        output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp)
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-DTC.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "dtc_summary":
        output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp)
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-DTC-SUMMARY.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "qc_qa_validated":
        # QC/QA validated output goes to outputs/qc_qa folder
        output_dir = os.path.join(config.QC_QA_OUTPUT_DIRECTORY, run_timestamp)
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-QC-QA-VALIDATED.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "qc_qa_stats":
        # QC/QA statistics output goes to automation_qa-qc folder
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "automation_qa-qc"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-QC-QA-STATS.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "qc_qa_tin_stats":
        output_dir = os.path.join(config.QC_QA_OUTPUT_DIRECTORY, run_timestamp)
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{config.BATCH_ID}-QC-QA-TIN-STATS.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "individual_cc":
        # Individual CC results
        base_filename = os.path.splitext(filename)[0].strip()
        output_dir = os.path.join(
            config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "individual"
        )
        os.makedirs(output_dir, exist_ok=True)
        df.to_csv(
            os.path.join(output_dir, f"{base_filename}-RESULTS-CC.csv"),
            index=False,
            quoting=1,
        )
        return None
    else:
        logging.error(
            f"Unknown output_type: {output_type}. Valid types are: 'final', 'individual', 'error', 'usage', 'usage_summary', 'dtc', 'dtc_summary', 'qc_qa_validated', 'qc_qa_stats', 'cc_results_full', 'dashboard_results_full', 'qc_qa_cc_full', 'qc_qa_error', 'parent_child', 'individual_cc'"
        )
        return None


def write_s3(
    df: pd.DataFrame, filename: str, run_timestamp: str, output_type: str
) -> Optional[str]:
    """
    Uploads a DataFrame to an S3 bucket as a CSV file.
    Returns Output path if parent-child mapping is performed.

    Args:
        df (pd.DataFrame): The DataFrame to be written to S3.
        filename (str): The base filename to use when saving the file.
        run_timestamp (str): A timestamp representing the current run.
        output_type (str): Type of output, either "final", "individual", "parent_child" or "error".

    The function generates an S3 key based on the configured BATCH_ID and
    uploads the DataFrame as a CSV file to the configured S3 bucket.
    The CSV file is saved in the "consolidated" folder with the naming
    convention "<BATCH_ID>-RESULTS.csv".

    Returns:
        - Output path if parent-child mapping is performed.
        - None, when this function is called for its side effect of uploading a file to S3.


    Configurations used:
        - config.BATCH_ID: Identifier for the batch process.
        - config.S3_CLIENT: The AWS S3 client for performing operations.
        - config.S3_OUTPUT_BUCKET: The target S3 bucket for storing the file.

    Prints:
        Confirmation message with the S3 key of the uploaded file.
    """
    if output_type == "final":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-RESULTS.csv"
        csv_buffer = df.to_csv(index=False).encode()
        try:
            config.S3_CLIENT.put_object(
                Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
            )
        except ClientError as e:
            logging.error(f"Error uploading to S3: {e}")
            return None
    elif output_type == "individual":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/individual/{filename}-RESULTS.csv"
        )
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "individual_cc":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/individual/{filename}-RESULTS-CC.csv"
        )
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "parent_child":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/parent-child/{config.BATCH_ID}-PC.csv"
        )
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "error":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/full_outputs/{config.BATCH_ID}-ERRORS.csv"
        csv_buffer = df.to_csv(index=False).encode()
        try:
            config.S3_CLIENT.put_object(
                Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
            )
        except ClientError as e:
            logging.error(f"Error uploading to S3: {e}")
            return None
    elif output_type in ("cc_results_full", "dashboard_results_full", "qc_qa_cc_full"):
        # Will be handled in try block using helper function
        output_path = None
        parent_child_path = None if output_type == "cc_results_full" else None
    elif output_type == "qc_qa_error":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/automation_qa-qc/{config.BATCH_ID}-ERRORS-QC-QA.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "usage":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE.csv"
        )
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "usage_summary":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE-SUMMARY.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "dtc":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "dtc_summary":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC-SUMMARY.csv"
        )
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "qc_qa_validated":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/automation_qa-qc/{config.BATCH_ID}-QC-QA-VALIDATED.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "qc_qa_stats":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/automation_qa-qc/{config.BATCH_ID}-QC-QA-STATS.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "qc_qa_tin_stats":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/automation_qa-qc/{config.BATCH_ID}-QC-QA-TIN-STATS.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    elif output_type == "individual_cc":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/individual/{filename}-RESULTS-CC.csv"
        )
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
    else:
        # Fallback path for unknown output types
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-unknown-{output_type}.csv"
        csv_buffer = df.to_csv(index=False).encode()
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )

    try:
        # Handle split files for cc_results_full, dashboard_results_full, and qc_qa_cc_full
        if output_type == "cc_results_full":
            parent_child_path = _write_s3_split_files(
                df=df,
                run_timestamp=run_timestamp,
                base_path="full_outputs/cc_results",
                full_filename=f"{config.BATCH_ID}-RESULTS-FULL.csv",
                split_filename_template=f"{config.BATCH_ID}-RESULTS-SPLIT{{idx}}.csv",
                return_path_for_pc=config.PERFORM_PARENT_CHILD_MAPPING,
            )
        elif output_type == "dashboard_results_full":
            _write_s3_split_files(
                df=df,
                run_timestamp=run_timestamp,
                base_path="full_outputs/dashboard_results",
                full_filename=f"{config.BATCH_ID}-RESULTS-dashboard.csv",
                split_filename_template=f"{config.BATCH_ID}-RESULTS-dashboard-SPLIT{{idx}}.csv",
                return_path_for_pc=False,
            )
        elif output_type == "qc_qa_cc_full":
            _write_s3_split_files(
                df=df,
                run_timestamp=run_timestamp,
                base_path="automation_qa-qc",
                full_filename=f"{config.BATCH_ID}-RESULTS-QC-QA-FULL.csv",
                split_filename_template=f"{config.BATCH_ID}-RESULTS-QC-QA-SPLIT{{idx}}.csv",
                return_path_for_pc=False,
            )
        elif output_type == "final":
            logging.info(f"Saved batch output file: {output_path}")
        elif output_type == "individual":
            logging.info(f"Saved individual output file: {output_path}")
        elif output_type == "parent_child":
            logging.info(f"Saved parent_child output file: {output_path}")
        elif output_type == "error":
            logging.info(f"Saved error output file: {output_path}")
        elif output_type == "usage":
            logging.info(f"Saved usage tracking file: {output_path}")
        elif output_type == "usage_summary":
            logging.info(f"Saved usage summary file: {output_path}")
        elif output_type == "dtc":
            logging.info(f"Saved DTC output file: {output_path}")
        elif output_type == "dtc_summary":
            logging.info(f"Saved DTC summary file: {output_path}")
        elif output_type == "qc_qa_validated":
            logging.info(f"Saved QC/QA validated data file: {output_path}")
        elif output_type == "qc_qa_stats":
            logging.info(f"Saved QC/QA statistics file: {output_path}")
        elif output_type == "qc_qa_tin_stats":
            logging.info(f"Saved QC/QA TIN statistics file: {output_path}")
        elif output_type == "qc_qa_error":
            logging.info(f"Saved QC/QA error file: {output_path}")
        elif output_type == "individual_cc":
            logging.info(f"Saved individual CC file: {output_path}")
        else:
            logging.info(f"Saved unknown output type file: {output_path}")

        if config.PERFORM_PARENT_CHILD_MAPPING and output_type == "cc_results_full":
            if "parent_child_path" in locals() and parent_child_path:
                return f"s3://{config.S3_OUTPUT_BUCKET}/{parent_child_path}"

        return None
    except ClientError as e:
        logging.error(f"Error uploading to S3: {e}")
        return None


def upload_local_file_to_s3(
    local_path: str,
    run_timestamp: str,
    s3_subpath: str,
    content_type: Optional[str] = None,
) -> Optional[str]:
    """
    Upload a local file to S3 using the same path structure as other pipeline outputs.

    Args:
        local_path: Full path to the local file.
        run_timestamp: Run timestamp (e.g. run_20260212_13-53_test_batch).
        s3_subpath: Subpath under run_timestamp (e.g. parent-child).
        content_type: Optional MIME type. Inferred from extension if not provided.

    Returns:
        S3 URI on success, None on failure.
    """
    if not os.path.isfile(local_path):
        logging.error(f"Cannot upload: file not found: {local_path}")
        return None
    filename = os.path.basename(local_path)
    s3_key = f"{config.BATCH_ID}/{run_timestamp}/{s3_subpath}/{filename}"
    if content_type is None and filename.lower().endswith(".xlsx"):
        content_type = (
            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        )
    try:
        with open(local_path, "rb") as f:
            put_kwargs = {
                "Bucket": config.S3_OUTPUT_BUCKET,
                "Key": s3_key,
                "Body": f.read(),
            }
            if content_type:
                put_kwargs["ContentType"] = content_type
            config.S3_CLIENT.put_object(**put_kwargs)
        s3_uri = f"s3://{config.S3_OUTPUT_BUCKET}/{s3_key}"
        logging.info(f"Uploaded file to S3: {s3_uri}")
        return s3_uri
    except ClientError as e:
        logging.error(f"Failed to upload {local_path} to S3: {e}")
        return None


def save_result_to_json(filename, result, json_folder):
    """Save a result dictionary to an individual JSON file in the specified folder.

    This function is thread-safe and can be called concurrently from multiple threads.
    It's useful for saving intermediate results or individual processing outputs during
    batch operations.

    Args:
        filename: Name of the source file that was processed (used to generate JSON filename)
        result: Dictionary containing the processing results to be saved
        json_folder: Path to the folder where the JSON file should be saved

    Returns:
        None; this function is called for its side effect of writing a JSON file to disk.

    Notes:
        - Creates the output folder if it doesn't exist
        - Automatically sanitizes the filename to remove invalid characters (/, \\, :)
        - Adds .json extension if not already present
        - Thread-safe: uses a lock to prevent concurrent write conflicts
        - Saves both the original filename and result in the JSON structure

    Example:
        >>> save_result_to_json(
        ...     "contract.pdf",
        ...     {"status": "processed", "result": "approved"},
        ...     "/path/to/output"
        ... )
        # Creates: /path/to/output/contract.pdf.json with content:
        # {"filename": "contract.pdf", "result": {"status": "processed", "result": "approved"}}
    """
    with json_write_lock:
        # Ensure the folder exists
        os.makedirs(json_folder, exist_ok=True)

        # Create a safe filename for JSON (replace invalid characters)
        safe_filename = filename.replace("/", "_").replace("\\", "_").replace(":", "_")
        if not safe_filename.endswith(".json"):
            safe_filename = f"{safe_filename}.json"

        json_file_path = os.path.join(json_folder, safe_filename)

        # Save result to individual JSON file
        with open(json_file_path, "w") as f:
            json.dump({"filename": filename, "result": result}, f, indent=2)


def read_DataFrame(file_path: str) -> pd.DataFrame:
    """
    Takes in a file path and returns a pandas DataFrame.
    Supports reading from local file system and only CSV from S3.
    """

    is_s3 = file_path.startswith("s3://")

    if is_s3:
        file = read_s3_csv(file_path)
        return file
    file = read_local(file_path)
    return file


def convert_json_to_dict(json_file_path: str) -> dict:
    """
    Reads a JSON file and converts it to a dictionary.
    Converts string keys 'True' and 'False' to boolean keys True and False.
    """
    json_str = Path(json_file_path).read_text()
    try:
        data = json.loads(json_str)
        # Convert "True"/"False" string keys to boolean keys
        converted = {
            (True if k == "True" else False if k == "False" else k): v
            for k, v in data.items()
        }
        return converted
    except json.JSONDecodeError:
        return {}