from typing import Optional, Union
import omni.ui as ui
import carb


class InspectorQuery:
    def __init__(self) -> None:
        pass

    @classmethod
    def child_widget(cls, widget: ui.Widget, predicate: str, last=False) -> Union[ui.Widget, None]:
        adjusted_predicate = predicate
        # when we don't have index we assume it is the first one
        index = None
        if "[" in adjusted_predicate:
            index = int(adjusted_predicate.split("[")[-1].split("]")[0])
            adjusted_predicate = adjusted_predicate.split("[")[0]

        if last and index is None:
            if widget.__class__.__name__ == adjusted_predicate:
                return widget
            else:
                carb.log_warn(f"Widget {widget} is not matching predicate {adjusted_predicate}")
                return None

        if not widget.__class__.__name__ == adjusted_predicate:
            carb.log_warn(f"Widget {widget} is not matching predicate {adjusted_predicate}")
            return None

        children = ui.Inspector.get_children(widget)
        if not index:
            index = 0

        if index >= len(children):
            carb.log_warn(f"Widget {widget} only as {len(children)}, asked for index {index}")
            return None

        return children[index]

    @classmethod
    def find_widget_path(cls, window: ui.Window, widget: ui.Widget) -> Union[str, None]:
        def traverse_widget_tree_for_match(
            location: ui.Widget, current_path: str, searched_widget: ui.Widget
        ) -> Union[str, None]:
            current_index = 0
            current_children = ui.Inspector.get_children(location)
            for a_child in current_children:
                class_name = a_child.__class__.__name__
                if a_child == widget:
                    return f"{current_path}[{current_index}]/{class_name}"

                if isinstance(a_child, ui.Container):
                    path = traverse_widget_tree_for_match(
                        a_child, f"{current_path}[{current_index}]/{class_name}", searched_widget
                    )
                    if path:
                        return path

                current_index = current_index + 1

            return None

        if window:
            start_path = f"{window.title}/Frame"
            return traverse_widget_tree_for_match(window.frame, start_path, widget)
        else:
            return None

    @classmethod
    def find_menu_item(cls, query: str) -> Optional[ui.Widget]:

        from omni.kit.mainwindow import get_main_window

        main_menu_bar = get_main_window().get_main_menu_bar()
        if not main_menu_bar:
            carb.log_warn("Current App doesn't have a MainMenuBar")
            return None

        tokens = query.split("/")
        current_children = main_menu_bar
        for i in range(1, len(tokens)):
            last = i == len(tokens) - 1
            child = InspectorQuery.child_widget(current_children, tokens[i], last)
            if not child:
                carb.log_warn(f"Failed to find Child Widget at level {tokens[i]}")
                return None

            current_children = child

        return current_children

    @classmethod
    def find_context_menu_item(cls, query: str, menu_root: ui.Menu) -> Optional[ui.Widget]:
        menu_items = ui.Inspector.get_children(menu_root)
        for menu_item in menu_items:
            if isinstance(menu_item, ui.MenuItem):
                if query in menu_item.text:
                    return menu_item
        return None

    @classmethod
    def find_widget(cls, query: str) -> Union[ui.Widget, None]:
        if not query:
            return None

        tokens = query.split("/")
        window = ui.Workspace.get_window(tokens[0])
        if not window:
            carb.log_warn(f"Failed to find window: {tokens[0]}")
            return None

        if not isinstance(window, ui.Window):
            carb.log_warn(f"window: {tokens[0]} is not a ui.Window, query only work on ui.Window")
            return None

        if not (tokens[1] == "Frame" or tokens[1] == "Frame[0]"):
            carb.log_warn("Query currently only Support '<WindowName>/Frame/* or <WindowName>/Frame[0]/ type of query")
            return None

        frame = window.frame

        if len(tokens) == 2:
            return frame

        if tokens[-1] == "":
            tokens = tokens[:-1]

        current_children = ui.Inspector.get_children(frame)[0]
        for i in range(2, len(tokens)):
            last = i == len(tokens) - 1
            child = InspectorQuery.child_widget(current_children, tokens[i], last)
            if not child:
                carb.log_warn(f"Failed to find Child Widget at level {tokens[i]}")
                return None

            current_children = child

        return current_children
