# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

from __future__ import annotations

import asyncio
import weakref
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Callable, Dict
from weakref import ProxyType

import carb
import omni.kit.app
import omni.kit.context_menu
import omni.ui as ui

from .manipulator import TransformManipulator
from .types import Operation

ICON_FOLDER_PATH = Path(
    f"{omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)}/data/icons"
)


class ToolbarTool(ABC):
    def __init__(
        self,
        manipulator: ProxyType[TransformManipulator],
        operation: Operation,
        toolbar_height: int,
        toolbar_payload: Dict[str, Any] = {},
        tooltip_update_fn: Callable[[str, bool, float], None] = None,
    ):
        self._manipulator = manipulator
        self._operation = operation
        self._model = manipulator.model
        self._toolbar_height = toolbar_height
        self._toolbar_payload = toolbar_payload
        self._tooltip_update_fn = tooltip_update_fn

    def destroy(self):
        self._manipulator = None
        self._model = None
        self._toolbar_payload = None
        self._tooltip_update_fn = None

    def __del__(self):
        self.destroy()

    @classmethod
    @abstractmethod
    def can_build(cls, manipulator: TransformManipulator, operation: Operation) -> bool:
        """
        Called right before a tool instance is to be instantiated to determine if this tool can be built on current toolbar.

        Args:
            manipulator (TransformManipulator): manipulator that hosts the toolbar
            operation (Operation): The transform Operation the tool will be built for.

        Return:
            True if the tool can be built. Its constructor will be called.
            False if not. The tool will be skipped and not placed on toolbar.
        """
        raise NotImplementedError("You must override `can_build` in your derived class!")
        return False


class DefaultMenuDelegate(ui.MenuDelegate):
    def get_style(self):
        from omni.kit.context_menu import style

        return style.MENU_STYLE


class SimpleToolButton(ToolbarTool):
    def __init__(self, menu_delegate: ui.MenuDelegate = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._show_menu_task = None
        self._cancel_next_value_changed = False  # workaround context menu triggering button state change
        self._ignore_model_change = False
        self._button = None
        self._stack = None
        self._menu_delegate = menu_delegate if menu_delegate else DefaultMenuDelegate()

    def destroy(self):
        self._menu_delegate = None

        if self._show_menu_task is not None:
            self._show_menu_task.cancel()
            self._show_menu_task = None

        if self._button:
            self._button.set_mouse_hovered_fn(None)

        self._sub = None
        self._button = None
        self._stack = None

        super().destroy()

    def _get_style(self) -> Dict:
        return {
            "Button": {"background_color": 0x0},
            "Button:checked": {"background_color": 0x8FD1912E},
            "Button:hovered": {"background_color": 0x0},
            "Button:pressed": {"background_color": 0x0},
        }

    def _build_widget(
        self,
        button_name: str,
        model: ui.AbstractValueModel,
        enabled_img_url: str,
        disabled_img_url: str = None,
        menu_index: str = None,
        menu_extension_id: str = None,
        no_toggle: bool = False,
        menu_on_left_click: bool = False,
        menu_payload: Dict[str, Any] = {},
        tooltip: str = "",
        disabled_tooltip: str = "",
    ):
        self._button_name = button_name
        self._model = model
        self._enabled_img_url = enabled_img_url
        self._disabled_img_url = disabled_img_url
        self._menu_index = menu_index
        self._menu_extension_id = menu_extension_id
        self._no_toggle = no_toggle
        self._menu_on_left_click = menu_on_left_click
        self._menu_payload = menu_payload
        self._tooltip = tooltip
        self._disabled_tooltip = disabled_tooltip
        self._button_hovered = False

        style = {}
        if self._disabled_img_url:
            style[f"Button.Image::{self._button_name}_enabled"] = {"image_url": f"{self._enabled_img_url}"}
            style[f"Button.Image::{self._button_name}_disabled"] = {"image_url": f"{self._disabled_img_url}"}
        else:
            style[f"Button.Image::{self._button_name}"] = {"image_url": f"{self._enabled_img_url}"}

        style.update(self._get_style())

        self._stack = ui.ZStack(width=0, height=0, style=style)
        with self._stack:
            dimension = self._toolbar_height
            self._button = ui.ToolButton(
                name=self._button_name,
                model=self._model,
                image_width=dimension,
                image_height=dimension,
            )
            self._button.set_mouse_hovered_fn(self._on_hovered)
            if self._menu_index is not None and self._menu_extension_id is not None:
                self._button.set_mouse_pressed_fn(lambda x, y, b, _: self._on_mouse_pressed(b, self._menu_index))
                self._button.set_mouse_released_fn(lambda x, y, b, _: self._on_mouse_released(b))
                self._build_flyout_indicator(dimension, dimension, self._menu_index, self._menu_extension_id)

        # only update name if enabled and disabled imgs are different
        if self._disabled_img_url or self._tooltip != self._disabled_tooltip:
            self._update_name(self._button, self._button.model.as_bool)

        self._sub = self._button.model.subscribe_value_changed_fn(self._on_model_changed)

    def _on_hovered(self, state: bool):
        self._button_hovered = state
        if self._tooltip and self._tooltip_update_fn:
            tooltip = self._disabled_tooltip if not self._button.model.as_bool and self._disabled_tooltip else self._tooltip
            self._tooltip_update_fn(tooltip, state, self._button.screen_position_x)

    def _on_model_changed(self, model):
        if self._ignore_model_change:
            return

        if self._no_toggle:
            self._ignore_model_change = True
            model.set_value(not model.as_bool)
            self._ignore_model_change = False
            return

        if self._cancel_next_value_changed:
            self._cancel_next_value_changed = False
            model.set_value(not model.as_bool)

        if self._disabled_img_url or self._tooltip != self._disabled_tooltip:
            self._update_name(self._button, model.as_bool)

    def _update_name(self, button: ui.ToolButton, enabled: bool):
        if self._disabled_img_url:
            button.name = f"{self._button_name}_{'enabled' if enabled else 'disabled'}"
        self._on_hovered(self._button_hovered)
        self._manipulator.refresh_toolbar()

    def _invoke_context_menu(self, button_id: str, right_click: bool, min_menu_entries: int = 1):
        context_menu = omni.kit.context_menu.get_instance()

        objects = {"widget_name": button_id, "manipulator": self._manipulator, "model": self._model}
        objects.update(self._toolbar_payload)
        objects.update(self._menu_payload)

        menu_list = omni.kit.context_menu.get_menu_dict(self._menu_index, self._menu_extension_id)
        context_menu.show_context_menu(
            self._menu_index, objects, menu_list, min_menu_entries, delegate=self._menu_delegate
        )

        # if from long LMB hold
        if not right_click:
            self._cancel_next_value_changed = True

    def _on_mouse_pressed(self, button, button_id: str, min_menu_entries: int = 1):
        """
        Function to handle flyout menu. Either with LMB long press or RMB click.

        Args:
            button_id: button_id of the context menu to be invoked.
            min_menu_entries: minimal number of menu entries required for menu to be visible (default 1).
        """
        if button == 1 or button == 0 and self._menu_on_left_click:  # show immediately
            # We cannot call self._invoke_context_menu directly inside of sc.Widget's context, otherwise it will be draw
            # on the auxiliary window. Schedule a async task so the menu is still in main window
            self._show_menu_task = asyncio.ensure_future(
                self._schedule_show_menu(button_id, wait_seconds=0.0, min_menu_entries=min_menu_entries)
            )
        elif button == 0:  # Schedule a task if hold LMB long enough
            self._show_menu_task = asyncio.ensure_future(
                self._schedule_show_menu(button_id, min_menu_entries=min_menu_entries)
            )

    def _on_mouse_released(self, button):
        if button == 0:
            if self._show_menu_task:
                self._show_menu_task.cancel()

    async def _schedule_show_menu(self, button_id: str, min_menu_entries: int = 1, wait_seconds: float = 0.3):
        if wait_seconds > 0.0:
            await asyncio.sleep(wait_seconds)
        try:
            self._invoke_context_menu(button_id, False, min_menu_entries)
        except Exception:
            import traceback

            carb.log_error(traceback.format_exc())
        self._show_menu_task = None

    def _build_flyout_indicator(self, width, height, index: str, extension_id: str, padding=-8, min_menu_count=1):
        indicator_size = 4
        with ui.Placer(offset_x=width - indicator_size - padding, offset_y=height - indicator_size - padding):
            indicator = ui.Image(
                f"{ICON_FOLDER_PATH}/flyout_indicator_dark.svg",
                width=indicator_size,
                height=indicator_size,
            )

        def on_menu_changed(evt: carb.events.IEvent):
            try:
                menu_list = omni.kit.context_menu.get_menu_dict(index, extension_id)
                # TODO check the actual menu entry visibility with show_fn
                indicator.visible = len(menu_list) >= min_menu_count
            except AttributeError as exc:
                carb.log_warn(f"on_menu_changed error {exc}")

        # Check initial state
        on_menu_changed(None)

        event_stream = omni.kit.context_menu.get_menu_event_stream()
        return event_stream.create_subscription_to_pop(on_menu_changed)
