from typing import Union, Optional, Callable, List
import omni.ui as ui

STATUS_ICON_WIDTH = 20
MENU_ARROW_SIZE = 8

__all__ = ["OptionBoxStyle", "ViewportMenuDelegate"]


class ViewportMenuDelegate(ui.MenuDelegate):
    """The menu delegate that draws in the menu or menuitem the given status.

    kwargs:
        icon_name (str): Show icon after menu text. Default "" to do not show.
        icon_width (int): Menu icon width. Defautl 16
        reserve_status (bool): Reserve space for status icon. Default True.
        force_checked (bool): Force to show checked status. It is strange that when a menu is checked, sub menu items
        checked flag will be cleared. Use this flag to show correct sub menu item status.

    """

    def __init__(
        self,
        icon_name: str = "",
        icon_width: int = 16,
        reserve_status: bool = True,
        force_checked: bool = False,
        icon_clicked_fn: Callable[[None], None] = None,
        build_custom_widgets: Callable[[ui.MenuDelegate, Union[ui.MenuItem, ui.Menu]], None] = None
    ):
        self._icon_name = icon_name
        self._icon_width = icon_width
        self._reserve_status = reserve_status
        self._force_checked = force_checked
        self._icon_clicked_fn = icon_clicked_fn
        self.icon: Optional[ui.Widget] = None
        self._container: Optional[ui.Widget] = None
        self._build_custom_widgets = build_custom_widgets

        super().__init__(propagate=False)

    def __icon_clicked(self, x, y, *arg, **kwargs):
        if self._icon_clicked_fn is None or self.icon is None:
            return
        left = self.icon.screen_position_x
        top = self.icon.screen_position_y
        if (x < left) or (y < top):
            return
        right = left + self.icon.computed_width
        bottom = top + self.icon.computed_height
        if (x > right) or (y > bottom):
            return
        self._icon_clicked_fn(x, y, *arg, **kwargs)

    def build_item(self, item: Union[ui.MenuItem, ui.Menu]):
        icon_type = "Menu.Item.Status"
        if self._force_checked:
            item.checked = self._force_checked
        self._container = ui.HStack(
            style_type_name_override="MenuBar.Item", selected=item.selected, checked=(item.checkable and item.checked)
        )
        with self._container:
            selected = item.selected or (item.checkable and item.checked)
            # Status icon before menu text
            if selected or self._reserve_status:
                self.icon = ui.ImageWithProvider(style_type_name_override=icon_type, width=STATUS_ICON_WIDTH, mouse_released_fn=self.__icon_clicked)

            # Text
            ui.Label(item.text, style_type_name_override="Menu.Item.Text")

            if self._icon_name or self._build_custom_widgets or isinstance(item, ui.Menu):
                ui.Spacer()

            # Icon after menu text
            if self._icon_name:
                ui.Spacer(width=10)
                ui.ImageWithProvider(
                    style_type_name_override="Menu.Item.Icon", name=self._icon_name, width=self._icon_width
                )

            if self._build_custom_widgets:
                self._build_custom_widgets(self, item)

            if isinstance(item, ui.Menu):
                # For menu, show arrow for children
                ui.Spacer(width=10)
                with ui.VStack(width=MENU_ARROW_SIZE / 2):
                    ui.Spacer()
                    ui.Triangle(
                        height=MENU_ARROW_SIZE,
                        alignment=ui.Alignment.RIGHT_CENTER,
                        style_type_name_override="MenuBar.Item.Triangle",
                    )
                    ui.Spacer()
                ui.Spacer(width=3)

    def get_selected(self) -> bool:
        return self._container.selected if self._container else False

    def set_selected(self, value: bool) -> bool:
        '''Set the 'selected' state and return whether it changed or not.'''
        value = bool(value)
        if self._container and (self.get_selected() != value):
            self._container.selected = value
            return True
        return False

    def get_checked(self) -> bool:
        return self._container.checked if self._container else False

    def set_checked(self, value: bool) -> bool:
        '''Set the 'checked' state and return whether it changed or not.'''
        value = bool(value)
        if self._container and (self.get_checked() != value):
            self._container.checked = value
            return True
        return False

    def __on_additional_btn_clicked(self, button: ui.Button) -> None:
        if self._on_additional_btn_clicked_fn:
            self._on_additional_btn_clicked_fn(button)

    @property
    def selected(self) -> bool:
        return self.get_selected()

    @selected.setter
    def selected(self, value: bool) -> None:
        self.set_selected(value)

    @property
    def checked(self) -> bool:
        return self.get_checked()

    @checked.setter
    def checked(self, value: bool) -> None:
        self.set_checked(value)
