Hunt-AI/Modules/submenu.py

import sys  # For exiting the program
from Modules.wrappers import *
from Modules.Submenus.analyst_notebook import analyst_notebook_submenu
from Modules.Submenus.about import about_submenu


# Global action status to persist across all submenus and their nested submenus
GLOBAL_ACTION_STATUS = {}
GLOBAL_TOTAL_PROGRESS = {"completed": 0, "total": 0}  # Tracks progress across all sections
GLOBAL_HIGHLIGHTS = {}  # Tracks highlighted actions


def build_submenu(menu_title, actions=None, module=None):
    """
    Builds a submenu dynamically by extracting functions from the given module,
    or using a provided `actions` dictionary for pre-defined submenus.
    Includes an `[ALL]` option to execute all functions or submenus on the page.
    Allows actions to be unmarked and highlighted.
    """
    if not actions and not module:
        raise ValueError("No module or actions provided to extract submenu items.")

    # If actions are not provided, dynamically extract from the module
    if not actions:
        actions = {
            str(idx + 1): {
                "description": func.__name__.replace("_", " ").title(),
                "function": func,
            }
            for idx, (name, func) in enumerate(module.items())
            if (
                callable(func)
                and func.__module__ == module["__name__"]  # Ensure the function is from this module
                and "submenu" not in func.__name__  # Exclude functions with 'submenu' in their name
            )
        }

    # Reassign sequential numbering starting at 1
    numbered_actions = {str(idx + 1): action for idx, action in enumerate(actions.values())}

    # Add the ALL option with key "ALL"
    numbered_actions["ALL"] = {
        "description": f"Run All {menu_title}",
        "function": lambda: run_all_actions(numbered_actions, module, GLOBAL_ACTION_STATUS[menu_title]),
    }

    # Initialize global action status and highlights for this submenu
    if menu_title not in GLOBAL_ACTION_STATUS:
        GLOBAL_ACTION_STATUS[menu_title] = {key: False for key in numbered_actions}
        GLOBAL_TOTAL_PROGRESS["total"] += len(numbered_actions) - 1  # Exclude the ALL option from the progress count

    if menu_title not in GLOBAL_HIGHLIGHTS:
        GLOBAL_HIGHLIGHTS[menu_title] = {key: False for key in numbered_actions}

    action_status = GLOBAL_ACTION_STATUS[menu_title]
    highlight_status = GLOBAL_HIGHLIGHTS[menu_title]

    while True:
        # Check if all individual actions are complete; if so, mark "ALL" as complete
        action_status["ALL"] = all(status for key, status in action_status.items() if key != "ALL")

        clear_screen()
        header()  # Display global app header

        print("=" * 70)
        print(f"         {menu_title} Menu        ")
        print("=" * 70, "\n")

        for key, action in numbered_actions.items():
            # Determine the status symbol with color
            status = f"{Fore.GREEN}✔️ " if action_status[key] else f"❌"
            # Highlight the key and description based on highlight status
            if key == "ALL":
                print(f"{Fore.BLUE}[{key}]{Style.RESET_ALL} {status} {action['description']}{Style.RESET_ALL}")
            elif highlight_status[key]:
                print(f"{Style.RESET_ALL}[{key}] {Fore.YELLOW}💡 {action['description']}{Style.RESET_ALL}")
            else:
                print(f"{Style.RESET_ALL}[{key}] {status} {action['description']}")

        print("[0] Return to Previous Menu")
        print("[q] EXIT\n")


        # Get user choice
        choice = input("Enter your choice: ").strip().upper()

         # Handle global commands
        if choice in ["NOTE", "ABOUT"]:
            global_command_handler(choice)
            continue

        # Handle highlighting (toggle with "!")
        if choice.endswith("!") and choice[:-1] in numbered_actions:
            key_to_toggle = choice[:-1]
            highlight_status[key_to_toggle] = not highlight_status[key_to_toggle]
            if highlight_status[key_to_toggle]:
                print(f"{Fore.YELLOW}Action [{key_to_toggle}] has been highlighted.")
            else:
                print(f"{Fore.YELLOW}Highlight removed from action [{key_to_toggle}].")
            input("Press Enter to continue...")
            continue

        # Handle unchecking of actions (with "-")
        if choice.startswith("-") and choice[1:] in numbered_actions:
            key_to_uncheck = choice[1:]
            if action_status[key_to_uncheck]:
                action_status[key_to_uncheck] = False
                GLOBAL_TOTAL_PROGRESS["completed"] -= 1
                print(f"{Fore.RED}Unchecked action [{key_to_uncheck}].")
            else:
                print(f"{Fore.RED}Action [{key_to_uncheck}] was already unchecked.")
            input("Press Enter to continue...")
            continue

        if choice in numbered_actions:
            if "function" in numbered_actions[choice]:
                # Call the action function
                numbered_actions[choice]["function"]()
                if choice != "ALL" and not action_status[choice]:  # Exclude ALL from progress tracking here
                    action_status[choice] = True
                    GLOBAL_TOTAL_PROGRESS["completed"] += 1
            elif "submenu" in numbered_actions[choice]:
                # Call the nested submenu
                numbered_actions[choice]["submenu"]()
                if is_submenu_complete(numbered_actions[choice]["submenu"]):
                    if not action_status[choice]:
                        action_status[choice] = True
                        GLOBAL_TOTAL_PROGRESS["completed"] += 1
        elif choice == "0":
            break
        elif choice in ["Q", "EXIT", "QUIT"]:
            sys.exit()
        else:
            print(f"{Fore.RED}Invalid choice. Press Enter to continue...")
            input()




def run_all_actions(actions, module, action_status):
    """
    Executes all actions in the submenu sequentially, excluding the "ALL" option itself.
    """
    for key, action in actions.items():
        if action["description"].startswith("Run All"):  # Skip the "ALL" option
            continue

        if "function" in action:
            if not action_status[key]:
                action["function"]()
                action_status[key] = True
                GLOBAL_TOTAL_PROGRESS["completed"] += 1
        elif "submenu" in action:  # Recursively handle submenus
            actions[key]["submenu"]()
            if is_submenu_complete(actions[key]["submenu"]):
                if not action_status[key]:
                    action_status[key] = True
                    GLOBAL_TOTAL_PROGRESS["completed"] += 1


def is_submenu_complete(submenu_function):
    """
    Check if all actions in a submenu are complete.
    """
    submenu_title = get_submenu_title(submenu_function)
    if submenu_title in GLOBAL_ACTION_STATUS:
        return all(GLOBAL_ACTION_STATUS[submenu_title].values())
    return False


def propagate_status_to_parent(menu_title):
    """
    Propagates the status to the parent menu when all submenu items are completed.
    """
    for parent_title, statuses in GLOBAL_ACTION_STATUS.items():
        if menu_title in statuses:
            statuses[menu_title] = all(
                GLOBAL_ACTION_STATUS[menu_title].values()
            )  # Mark parent complete if all are done


def get_submenu_title(submenu_function):
    """
    Converts a submenu function name to a readable submenu title.
    """
    return submenu_function.__name__.replace("_submenu", "").replace("_", " ").title()