"""
Prettify the execution information of the graph.
"""
from typing import Union

def prettify_exec_info(complete_result: list[dict], as_string: bool = True) -> Union[str, list[dict]]:
    """
    Formats the execution information of a graph showing node statistics.

    Args:
        complete_result (list[dict]): The execution information containing node statistics.
        as_string (bool, optional): If True, returns a formatted string table. 
                                  If False, returns the original list. Defaults to True.

    Returns:
        Union[str, list[dict]]: A formatted string table if as_string=True, 
        otherwise the original list of dictionaries.
    """
    if not as_string:
        return complete_result

    if not complete_result:
        return "Empty result"

    # Format the table
    lines = []
    lines.append("Node Statistics:")
    lines.append("-" * 100)
    lines.append(f"{'Node':<20} {'Tokens':<10} {'Prompt':<10} {'Compl.':<10} {'Requests':<10} {'Cost ($)':<10} {'Time (s)':<10}")
    lines.append("-" * 100)

    for item in complete_result:
        node = item['node_name']
        tokens = item['total_tokens']
        prompt = item['prompt_tokens']
        completion = item['completion_tokens']
        requests = item['successful_requests']
        cost = f"{item['total_cost_USD']:.4f}"
        time = f"{item['exec_time']:.2f}"

        lines.append(
            f"{node:<20} {tokens:<10} {prompt:<10} {completion:<10} {requests:<10} {cost:<10} {time:<10}"
        )

    return "\n".join(lines)
