SXXXXXXX_CreateIconFromFilesPng/tool_utils/communicator.py

# tool_utils/communicator.py

"""
Utility module for Python tools designed to integrate with ProjectUtility.

Provides standardized functions for argument parsing based on a tool's
configuration and for sending structured messages (progress, status, results)
via standard output using JSON Lines, as well as handling standard error
and exit codes according to the ProjectUtility Tool Interface specification.
"""

import argparse
import json
import sys
import os
import logging
from typing import Dict, Any, List, Optional, Tuple, Callable

# Set up a basic logger for the utility module itself
# Note: Tools using this should ideally configure their own application-level logging.
util_logger = logging.getLogger(__name__)
# Add a basic handler if no other logging is configured by the tool
if not util_logger.hasHandlers():
    util_logger.addHandler(logging.StreamHandler(sys.stderr))
    util_logger.setLevel(logging.WARNING)  # Default to WARNING to avoid being too noisy

# --- Constants for message types ---
MSG_TYPE_PROGRESS = "progress"
MSG_TYPE_STATUS = "status"
MSG_TYPE_LOG = "log"
MSG_TYPE_RESULT = "result"

# --- Argument Parsing ---


def parse_arguments(tool_parameters: List[Dict[str, Any]]) -> argparse.Namespace:
    """
    Parses command-line arguments based on the tool parameter definitions.

    Follows the convention:
    - String, Int, Float, File, Folder: --name <value>
    - Boolean: --name (flag is present if True)

    Args:
        tool_parameters: A list of dictionaries, where each dictionary defines
                         a parameter (matching the structure in tool_config.json,
                         requiring at least 'name', 'type', 'required', 'description').

    Returns:
        An argparse.Namespace object containing the parsed arguments.
        The script will exit if '--help' is invoked.
    """
    parser = argparse.ArgumentParser(
        description="Tool script adhering to ProjectUtility interface.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,  # Show defaults in help
    )

    for param in tool_parameters:
        name = param.get("name")
        ptype = str(param.get("type", "string")).lower()  # Default to string if missing
        required = param.get("required", False)
        default = param.get("default")
        description = param.get("description", "")  # Help text

        if not name:
            util_logger.warning(
                f"Skipping parameter definition missing 'name': {param}"
            )
            continue

        arg_name = f"--{name}"

        try:
            if ptype == "boolean":
                # For booleans, store_true means the arg name implies True
                # If required=True, it doesn't make much sense for a boolean flag,
                # but argparse handles 'required' flags correctly if needed.
                # Default behavior: if flag present -> True, else False.
                # We don't use the 'default' from config directly here as store_true handles it.
                parser.add_argument(
                    arg_name,
                    action="store_true",  # Automatically sets default=False
                    help=f"(Flag) {description}",
                    required=required,  # Can be required, though unusual for flags
                )
                # Note: If a boolean needs to default to True and be toggled off,
                # use action="store_false" and adjust logic/help text.
            else:
                # For other types, specify the type for basic conversion by argparse
                # Note: 'file' and 'folder' are treated as strings here; validation
                # would happen within the tool's logic.
                arg_type = str
                if ptype == "integer":
                    arg_type = int
                elif ptype == "float":
                    arg_type = float

                parser.add_argument(
                    arg_name,
                    type=arg_type,
                    required=required,
                    default=default,
                    help=description,
                    metavar=f"<{ptype.upper()}>",  # Provides hint in help message, e.g. <STRING>
                )
        except Exception as e:
            util_logger.error(
                f"Failed to add argument for parameter '{name}': {e}", exc_info=True
            )

    # Parse known args to allow flexibility if gestore sends extra args
    args, unknown_args = parser.parse_known_args()

    if unknown_args:
        util_logger.warning(f"Received unrecognized arguments: {unknown_args}")

    # Post-parsing validation (optional but recommended)
    # Example: Check if required arguments that are *not* booleans actually have a value
    #          (argparse usually handles this with 'required=True', but defensive check)
    for param in tool_parameters:
        name = param.get("name")
        required = param.get("required", False)
        ptype = str(param.get("type", "string")).lower()
        if required and name and ptype != "boolean":
            if (
                getattr(args, name.replace("-", "_"), None) is None
            ):  # Check if value is None
                # Argparse should have already exited if required=True and missing,
                # but this handles edge cases or if required=False but tool logic needs it.
                print_error(
                    f"Required parameter '--{name}' is missing or has no value.",
                    exit_code=2,
                )
                # print_error exits the script

    util_logger.debug(f"Arguments parsed successfully: {args}")
    return args


# --- Standard Output Communication (JSON Lines) ---


def _send_json_message(msg_type: str, data: Dict[str, Any]) -> None:
    """Formats and sends a JSON message to stdout."""
    message = {"type": msg_type, **data}  # Combine type and data payload
    try:
        json_str = json.dumps(message, ensure_ascii=False)
        print(json_str, file=sys.stdout, flush=True)  # Print and FLUSH!
        util_logger.debug(f"Sent message: {json_str}")
    except (TypeError, ValueError) as e:
        util_logger.error(
            f"Failed to serialize message to JSON: {message}. Error: {e}", exc_info=True
        )
        # Also report the serialization error via stderr
        print(
            f"TOOL_UTILS_ERROR: Failed to send JSON message. Type: {msg_type}. Error: {e}",
            file=sys.stderr,
            flush=True,
        )


def send_progress(value: float, message: Optional[str] = None) -> None:
    """Sends a progress update message to stdout.

    Args:
        value: Progress value (ideally between 0.0 and 1.0).
        message: Optional descriptive message about the progress.
    """
    payload = {"value": max(0.0, min(1.0, value))}  # Clamp value between 0 and 1
    if message is not None:
        payload["message"] = str(message)
    _send_json_message(MSG_TYPE_PROGRESS, payload)


def send_status(message: str) -> None:
    """Sends a status message (general information) to stdout.

    Args:
        message: The status message string.
    """
    if not message:
        util_logger.warning("Attempted to send an empty status message.")
        return
    _send_json_message(MSG_TYPE_STATUS, {"message": str(message)})


def send_log(level: str, message: str) -> None:
    """Sends a log message to stdout (for potential capture by the manager).
       Note: This is distinct from the tool's internal logging to stderr/file.

    Args:
        level: Log level (e.g., "info", "warning", "debug", "error").
        message: The log message string.
    """
    if not message:
        util_logger.warning("Attempted to send an empty log message.")
        return
    _send_json_message(
        MSG_TYPE_LOG, {"level": str(level).lower(), "message": str(message)}
    )


def send_result(result_data: Dict[str, Any]) -> None:
    """Sends a result data message to stdout.

    Args:
        result_data: A dictionary containing the results to be sent.
                     Must be JSON serializable.
    """
    if not isinstance(result_data, dict):
        util_logger.error(
            f"Result data must be a dictionary, got {type(result_data)}. Cannot send result."
        )
        print(
            f"TOOL_UTILS_ERROR: Result data must be a dictionary. Data: {result_data}",
            file=sys.stderr,
            flush=True,
        )
        return
    _send_json_message(MSG_TYPE_RESULT, {"data": result_data})


# --- Standard Error and Exit Codes ---


def print_error(message: str, exit_code: Optional[int] = 1) -> None:
    """
    Prints an error message to stderr and optionally exits the script.

    Args:
        message: The error message string.
        exit_code: The exit code to use (e.g., 1 for general error).
                   If None, the script will not exit after printing.
                   Defaults to 1.
    """
    print(f"ERROR: {message}", file=sys.stderr, flush=True)  # Print and FLUSH!
    util_logger.error(message)  # Also log it via logging if configured

    if exit_code is not None:
        util_logger.info(f"Exiting with code {exit_code} due to error.")
        sys.exit(exit_code)


def exit_success(message: Optional[str] = None) -> None:
    """
    Exits the script with exit code 0 (success).
    Optionally prints a final success message to stdout (as status).

    Args:
        message: Optional final success message to send as status.
    """
    if message:
        send_status(f"Success: {message}")  # Send final message via standard channel
    util_logger.info("Exiting with code 0 (Success).")
    sys.exit(0)


# --- Helper to Load Tool Config (Optional but useful for tools) ---


def load_tool_config(
    config_filename: str = "tool_config.json",
) -> Optional[Dict[str, Any]]:
    """
    Loads the tool's configuration file (typically 'tool_config.json')
    from the directory containing the currently running script.

    Args:
        config_filename: The name of the configuration file.

    Returns:
        A dictionary with the loaded configuration data, or None if the
        file is not found or cannot be parsed.
    """
    script_dir = os.path.dirname(
        os.path.abspath(sys.argv[0])
    )  # Directory of the main script
    config_path = os.path.join(script_dir, config_filename)

    if not os.path.isfile(config_path):
        util_logger.error(f"Tool configuration file not found: {config_path}")
        print(
            f"TOOL_UTILS_ERROR: Configuration file '{config_filename}' not found in script directory.",
            file=sys.stderr,
            flush=True,
        )
        return None

    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config_data = json.load(f)
        if not isinstance(config_data, dict):
            util_logger.error(
                f"Invalid format in configuration file: {config_path}. Expected a JSON object."
            )
            print(
                f"TOOL_UTILS_ERROR: Configuration file '{config_filename}' does not contain a valid JSON object.",
                file=sys.stderr,
                flush=True,
            )
            return None
        util_logger.debug(f"Loaded tool configuration from: {config_path}")
        return config_data
    except json.JSONDecodeError as e:
        util_logger.error(
            f"Invalid JSON in configuration file: {config_path}. Error: {e}"
        )
        print(
            f"TOOL_UTILS_ERROR: Invalid JSON in '{config_filename}'. Error: {e}",
            file=sys.stderr,
            flush=True,
        )
        return None
    except Exception as e:
        util_logger.exception(
            f"Failed to load or parse tool configuration file: {config_path}"
        )
        print(
            f"TOOL_UTILS_ERROR: Could not load '{config_filename}'. Error: {e}",
            file=sys.stderr,
            flush=True,
        )
        return None


# --- Example Usage (within a tool script) ---
#
# if __name__ == "__main__":
#     # 1. Load configuration to get parameter definitions
#     config = load_tool_config()
#     if config is None:
#         print_error("Failed to load tool configuration.", exit_code=3)
#
#     tool_params_def = config.get("parameters", [])
#
#     # 2. Parse arguments based on the definition
#     try:
#         args = parse_arguments(tool_params_def)
#     except Exception as e:
#         print_error(f"Error parsing arguments: {e}", exit_code=2)
#
#     # 3. Use the communicator functions during execution
#     send_status("Tool execution started.")
#     try:
#         # --- Your Tool's Logic ---
#         user_input = args.user_text # Access parsed args
#         print(f"Tool received text: {user_input}") # Normal diagnostic print to stdout (optional)
#         sys.stdout.flush()
#
#         total_steps = 5
#         for i in range(total_steps):
#             # Simulate work
#             import time
#             time.sleep(0.5)
#             progress_val = (i + 1) / total_steps
#             send_progress(progress_val, f"Step {i+1} completed.")
#             if i == 2:
#                 send_log("warning", "Potential issue detected during step 3.")
#
#         # Simulate an error condition
#         # if args.user_text == "error":
#         #     print_error("User triggered a simulated error!", exit_code=5)
#
#         # --- End of Tool's Logic ---
#
#         # 4. Send results (if any)
#         final_result = {"processed_text": user_input.upper(), "steps_done": total_steps}
#         send_result(final_result)
#
#         # 5. Exit successfully
#         exit_success("Tool finished processing.")
#
#     except Exception as e:
#         # Catch unexpected errors in the tool's logic
#         print_error(f"An unexpected error occurred in the tool: {e}", exit_code=1)
#