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


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

    @classmethod
    def _child_widget(cls, widget: ui.Widget, path_segment: str, show_warnings=True) -> Union[ui.Widget, None]:
        """
        Find a widget from a path given a starting widget
        """
        # print ("looking for", path_segment, "child of, ", widget)
        adjusted_path_segment = path_segment
        # when we don't have index we assume it is the first one
        index = None
        if "[" in adjusted_path_segment:
            index = int(adjusted_path_segment.split("[")[-1].split("]")[0])
            adjusted_path_segment = adjusted_path_segment.split("[")[0]

        children = ui.Inspector.get_children(widget)

        counter = 0
        for a_child in children:
            if not a_child:
                continue
            # NOTE: Wait until we have identifier support before activating this
            if adjusted_path_segment == a_child.identifier and index == None:
                return a_child
            elif adjusted_path_segment == a_child.__class__.__name__:
                if index == counter:
                    return a_child
                else:
                    counter += 1

        if not widget.__class__.__name__ == adjusted_path_segment:
            if show_warnings:
                carb.log_warn(f"Widget {widget} does not match path segment {adjusted_path_segment}")
            return None

        if not index:
            index = 0

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

        return children[index]

    @classmethod
    def _find_children_by_type(cls, widget: ui.Widget, path_segment: str) -> List[ui.Widget]:
        """
        given a starting widget, look at it's children to find the path segment requested, where
        the path segment index has a wildcard (e.g [*])
        Assumes input path_segment is like "Type[*]"
        """

        adjusted_path_segment = path_segment[: path_segment.find("[")]
        children = ui.Inspector.get_children(widget)
        children_of_type = [c for c in children if c.__class__.__name__ == adjusted_path_segment]
        return children_of_type

    @classmethod
    def get_widget_path(cls, window: ui.Window, widget: ui.Widget) -> Union[str, None]:
        """
        given a Window and a Widget in that window, get it's path
        """

        def traverse_widget_tree_for_match(
            starting_widget: ui.Widget, current_path: str, searched_widget: ui.Widget
        ) -> Union[str, None]:
            current_index = 0
            current_children = ui.Inspector.get_children(starting_widget)
            current_children = [c for c in current_children if c]

            type_frequencies = dict(zip([a.__class__ for a in current_children], [0] * len(current_children)))

            for a_child in current_children:
                class_name = a_child.__class__.__name__
                current_index = type_frequencies[a_child.__class__]
                widget_name = a_child.identifier or f"{class_name}[{current_index}]"
                type_frequencies[a_child.__class__] += 1

                if a_child == searched_widget:
                    p = f"{current_path}/{widget_name}"
                    return p

                if isinstance(a_child, ui.Container) or isinstance(a_child, ui.TreeView):
                    path = traverse_widget_tree_for_match(a_child, f"{current_path}/{widget_name}", searched_widget)
                    if path:
                        return path
            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]:
        """
        TODO: This may now be legacy functionality after https://gitlab-master.nvidia.com/omniverse/kit/-/merge_requests/10402
        is merged so need to see if it's actually being used on 102 tests etc
        """

        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("/")
        print(tokens)
        current_children = main_menu_bar
        for i in range(1, len(tokens)):
            last = i == len(tokens) - 1
            print(tokens[i])
            child = cls._child_widget(current_children, tokens[i])
            if not child != 1:
                carb.log_warn(f"find_menu_item: Failed to find Child Widget at level {tokens[i]}")
                return None

            current_children = child

        return current_children

    @classmethod
    def _parse_input(cls, query: str) -> Tuple[bool, Union[ui.Widget, None], List[str], str]:
        """
        parse and validate the input.
        return widge tokens only ie sans window
        """
        # TODO Replace w regex
        tokens = query.split("//")
        window_name = tokens[0] if len(tokens) > 1 else None
        widget_predicate = ""
        widget_part = tokens[1] if len(tokens) > 1 else tokens[0]
        widget_part_list = widget_part.split(".", maxsplit=1)
        widget_path = widget_part_list[0]
        if len(widget_part_list) > 1:
            widget_predicate = widget_part_list[1]

        window = None

        if window_name:
            windows = ui.Workspace.get_windows()
            window_list = []
            for window in windows:
                if window.title == window_name:
                    window_list.append(window)

            if not window_list:
                carb.log_warn(f'Failed to find window: "{window_name}"')
                return False, None, [], widget_predicate

            if len(window_list) == 1:
                window = window_list[0]
            else:
                carb.log_warn(
                    f'found {len(window_list)} windows named "{window_name}". Using first visible window found'
                )
                window = None
                for win in window_list:
                    if win.visible:
                        window = win
                        break
                if not window:
                    carb.log_warn(f'Failed to find visible window: "{window_name}"')
                    return False, None, [], widget_predicate

            if not isinstance(window, ui.Window) and not isinstance(window, ui.ToolBar):
                carb.log_warn(f"window: {window_name} is not a ui.Window, query only works on ui.Window")
                return False, None, [], widget_predicate

        widget_tokens = widget_path.split("/")

        if window and not (widget_tokens[0] == "Frame" or widget_tokens[0] == "Frame[0]"):
            carb.log_warn("Query with a window currently only supports '<WindowName>//Frame/* type query")
            return False, None, [], widget_predicate

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

        return True, window, widget_tokens, widget_predicate

    @classmethod
    def find_widgets(cls, query: str, root_widgets=[]) -> List[ui.Widget]:
        """
        find a set of widgets that match our query.
        """

        def _get_descendants(widget: ui.Widget):
            kids = []
            local_kids = ui.Inspector.get_children(widget)
            kids.extend(local_kids)
            for k in local_kids:
                child_kids = _get_descendants(k)
                kids.extend(child_kids)
            return kids

        validate_status, window, widget_tokens, predicate = cls._parse_input(query)

        if not validate_status:
            return []

        if root_widgets:
            curr_kids = root_widgets
        else:
            curr_kids = [window.frame] if window else []

        if not curr_kids:
            return []

        curr_kids_tmp = []
        for i in range(0, len(widget_tokens)):
            # Skip only parent frame
            if widget_tokens[i] == "Frame":
                continue

            curr_kids_tmp = []
            for current_child in curr_kids:
                if isinstance(current_child, ui.Window):
                    current_child = current_child.frame

                if widget_tokens[i] == "*":
                    curr_kids_tmp += ui.Inspector.get_children(current_child)
                elif widget_tokens[i].find("[*]") != -1:
                    curr_kids_tmp += cls._find_children_by_type(current_child, widget_tokens[i])
                elif widget_tokens[i] == "**":
                    curr_kids_tmp += _get_descendants(current_child)
                else:
                    child = cls._child_widget(current_child, widget_tokens[i], show_warnings=False)
                    if child:
                        curr_kids_tmp += [child]

            # Filter out any leaf node widgets as we're not on the last pass and last pass deals with them
            curr_kids = [kid for kid in curr_kids_tmp if len(ui.Inspector.get_children(kid)) > 0]

        # Remove None(s)
        curr_kids_tmp = [w for w in curr_kids_tmp if w is not None]

        if predicate:
            curr_kids_tmp = [w for w in curr_kids_tmp if eval("w." + predicate)]

        return curr_kids_tmp

    @classmethod
    def find_widget(cls, query: str) -> Union[ui.Widget, None]:
        """
        find a single widget given a full widget path
        """

        validate_status, window, widget_tokens, _ = cls._parse_input(query)
        if not validate_status:
            return None

        if len(widget_tokens) == 1:
            return window.frame

        current_child = window.frame

        for i in range(1, len(widget_tokens)):
            child = cls._child_widget(current_child, widget_tokens[i])
            if not child:
                carb.log_warn(f'find_widget: from query "{query}" Failed to widget {widget_tokens[i]}')
                actual_kids = ui.Inspector.get_children(current_child)
                names = [a.identifier or f"{a.__class__.__name__}" for a in actual_kids]
                carb.log_warn(f"find_widget: did find widgets {names}")
                return None

            current_child = child

        return current_child
