doczyai-pipelines/fieldExtraction/src/utils/io_utils.py

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

import pandas as pd
import src.tracking.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()


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 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 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 == "parent_child":
        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}-Parent_Child.csv"),
            index=False,
            quoting=1,
        )
        return None
    elif output_type == "error":
        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}-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 == "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
    else:
        logging.error(
            f"Unknown output_type: {output_type}. Valid types are: 'final', 'individual', 'error', 'usage', 'usage_summary', 'dtc', 'dtc_summary'"
        )
        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"
    elif output_type == "individual":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/individual/{filename}-RESULTS.csv"
        )
    elif output_type == "parent_child":
        output_path = (
            f"{config.BATCH_ID}/{run_timestamp}/parent_child/{filename}-RESULTS.csv"
        )
    elif output_type == "error":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-ERRORS.csv"
    elif output_type == "usage":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE.csv"
    elif output_type == "usage_summary":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE-SUMMARY.csv"
    elif output_type == "dtc":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC.csv"
    elif output_type == "dtc_summary":
        output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC-SUMMARY.csv"
    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()
    try:
        config.S3_CLIENT.put_object(
            Bucket=config.S3_OUTPUT_BUCKET, Key=output_path, Body=csv_buffer
        )
        if 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}")
        else:
            logging.info(f"Saved unknown output type file: {output_path}")
        
        if config.PERFORM_PARENT_CHILD_MAPPING:
            return f"s3://{config.S3_OUTPUT_BUCKET}/{output_path}"
        
        return None
    except ClientError as e:
        logging.error(f"Error uploading 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 {}