# Copyright (c) 2020, 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.
#

import weakref
from functools import partial
from typing import Callable, List, Optional

import omni.kit.commands
import omni.ui as ui
from omni.kit.widget.stage import StageWidget
from pxr import Sdf


class SelectionWatch(object):
    def __init__(self, stage, on_selection_changed_fn, filter_type_list, filter_lambda, tree_view=None):
        self._stage = weakref.ref(stage)
        self._last_selected_prim_paths = None
        self._filter_type_list = filter_type_list
        self._filter_lambda = filter_lambda
        self._on_selection_changed_fn = on_selection_changed_fn
        self._targets_limit = 0
        if tree_view:
            self.set_tree_view(tree_view)

    def reset(self, targets_limit):
        self._targets_limit = targets_limit
        self.clear_selection()

    def set_tree_view(self, tree_view):
        self._tree_view = tree_view
        self._tree_view.set_selection_changed_fn(self._on_widget_selection_changed)
        self._last_selected_prim_paths = None

    def clear_selection(self):
        if not self._tree_view:
            return

        self._tree_view.model.update_dirty()
        self._tree_view.selection = []
        if self._on_selection_changed_fn:
            self._on_selection_changed_fn([])

    def _on_widget_selection_changed(self, selection):
        stage = self._stage()
        if not stage:
            return

        prim_paths = [str(item.path) for item in selection if item]

        # Deselect instance proxy items if they were selected
        selection = [item for item in selection if item and not item.instance_proxy]

        # Although the stage view has filter, you can still select the ancestor of filtered prims, which might not match the type.
        if len(self._filter_type_list) != 0:
            filtered_selection = []
            for item in selection:
                prim = stage.GetPrimAtPath(item.path)
                if prim:
                    for type in self._filter_type_list:
                        if prim.IsA(type):
                            filtered_selection.append(item)
                            break
            if filtered_selection != selection:
                selection = filtered_selection

        # or the ancestor might not match the lambda filter
        if self._filter_lambda is not None:
            selection = [item for item in selection if self._filter_lambda(stage.GetPrimAtPath(item.path))]

        # Deselect if over the limit
        if self._targets_limit > 0 and len(selection) > self._targets_limit:
            selection = selection[: self._targets_limit]

        if self._tree_view.selection != selection:
            self._tree_view.selection = selection
            prim_paths = [str(item.path) for item in selection]

        if prim_paths == self._last_selected_prim_paths:
            return

        self._last_selected_prim_paths = prim_paths
        if self._on_selection_changed_fn:
            self._on_selection_changed_fn(self._last_selected_prim_paths)

    def enable_filtering_checking(self, enable: bool):
        """
        It is used to prevent selecting the prims that are filtered out but
        still displayed when such prims have filtered children. When `enable`
        is True, SelectionWatch should consider filtering when changing Kit's
        selection.
        """
        pass

    def set_filtering(self, filter_string: Optional[str]):
        pass


class RelationshipTargetPicker:
    def __init__(
        self, stage, relationship_widget, filter_type_list, filter_lambda, on_add_targets: Optional[Callable] = None
    ):
        self._weak_stage = weakref.ref(stage)
        self._relationship_widget = relationship_widget
        self._filter_lambda = filter_lambda
        self._selected_paths = []
        self._filter_type_list = filter_type_list
        self._on_add_targets = on_add_targets

        def on_window_visibility_changed(visible):
            if not visible:
                self._stage_widget.open_stage(None)
            else:
                # Only attach the stage when picker is open. Otherwise the Tf notice listener in StageWidget kills perf
                self._stage_widget.open_stage(self._weak_stage())

        self._window = ui.Window(
            "Select Target(s)",
            width=400,
            height=400,
            visible=False,
            flags=0,
            visibility_changed_fn=on_window_visibility_changed,
        )
        with self._window.frame:
            with ui.VStack():
                with ui.Frame():
                    self._stage_widget = StageWidget(None, columns_enabled=["Type"])
                    self._selection_watch = SelectionWatch(
                        stage=stage,
                        on_selection_changed_fn=self._on_selection_changed,
                        filter_type_list=filter_type_list,
                        filter_lambda=filter_lambda,
                    )
                    self._stage_widget.set_selection_watch(self._selection_watch)

                def on_select(weak_self):
                    weak_self = weak_self()
                    if not weakref:
                        return

                    pending_add = []
                    for relationship in weak_self._relationship_widget._relationships:
                        if relationship:
                            existing_targets = relationship.GetTargets()
                            for selected_path in weak_self._selected_paths:
                                selected_path = Sdf.Path(selected_path)
                                if selected_path not in existing_targets:
                                    pending_add.append((relationship, selected_path))

                    if len(pending_add):
                        omni.kit.undo.begin_group()
                        for add in pending_add:
                            omni.kit.commands.execute("AddRelationshipTarget", relationship=add[0], target=add[1])
                        omni.kit.undo.end_group()
                        if self._on_add_targets:
                            self._on_add_targets(pending_add)

                    weak_self._window.visible = False

                with ui.VStack(
                    height=0, style={"Button.Label:disabled": {"color": 0xFF606060}}
                ):  # TODO consolidate all styles
                    self._label = ui.Label("Selected Path(s):\n\tNone")
                    self._button = ui.Button(
                        "Add",
                        height=10,
                        clicked_fn=partial(on_select, weak_self=weakref.ref(self)),
                        enabled=False,
                        identifier="add_button"
                    )

    def clean(self):
        self._window.set_visibility_changed_fn(None)
        self._window = None
        self._selection_watch = None
        self._stage_widget.open_stage(None)
        self._stage_widget.destroy()
        self._stage_widget = None
        self._filter_type_list = None
        self._filter_lambda = None
        self._on_add_targets = None

    def show(self, targets_limit):
        self._targets_limit = targets_limit
        self._selection_watch.reset(targets_limit)
        self._window.visible = True
        if self._filter_lambda is not None:
            self._stage_widget._filter_by_lambda({"relpicker_filter": self._filter_lambda}, True)
        if self._filter_type_list:
            self._stage_widget._filter_by_type(self._filter_type_list, True)
            self._stage_widget.update_filter_menu_state(self._filter_type_list)

    def _on_selection_changed(self, paths):
        self._selected_paths = paths
        if self._button:
            self._button.enabled = len(self._selected_paths) > 0
        if self._label:
            text = "\n\t".join(self._selected_paths)
            label_text = "Selected Path(s)"
            if self._targets_limit > 0:
                label_text += f" ({len(self._selected_paths)}/{self._targets_limit})"
            label_text += f":\n\t{text if len(text) else 'None'}"
            self._label.text = label_text


class RelationshipEditWidget:
    def __init__(self, stage, attr_name, prim_paths, additional_widget_kwargs=None):
        self._id_name = f"{prim_paths[-1]}_{attr_name}".replace('/', '_')
        self._relationships = [stage.GetPrimAtPath(path).GetRelationship(attr_name) for path in prim_paths]
        self._additional_widget_kwargs = additional_widget_kwargs if additional_widget_kwargs else {}
        self._targets_limit = self._additional_widget_kwargs.get("targets_limit", 0)
        self._button = None
        self._target_picker = RelationshipTargetPicker(
            stage,
            self,
            self._additional_widget_kwargs.get("target_picker_filter_type_list", []),
            self._additional_widget_kwargs.get("target_picker_filter_lambda", None),
            self._additional_widget_kwargs.get("target_picker_on_add_targets", None),
        )
        self._frame = ui.Frame()
        self._frame.set_build_fn(self._build)
        self._on_remove_target = self._additional_widget_kwargs.get("on_remove_target", None)
        self._enabled = self._additional_widget_kwargs.get("enabled", True)
        self._shared_targets = None

    def clean(self):
        self._target_picker.clean()
        self._target_picker = None
        self._frame = None
        self._button = None
        self._label = None
        self._on_remove_target = None
        self._enabled = True

    def is_ambiguous(self) -> bool:
        return self._shared_targets is None

    def get_all_comp_ambiguous(self) -> List[bool]:
        return []

    def get_relationship_paths(self) -> List[Sdf.Path]:
        return [rel.GetPath() for rel in self._relationships]

    def get_property_paths(self) -> List[Sdf.Path]:
        return self.get_relationship_paths()

    def get_targets(self) -> List[Sdf.Path]:
        return self._shared_targets

    def set_targets(self, targets: List[Sdf.Path]):
        with omni.kit.undo.group():
            for relationship in self._relationships:
                if relationship:
                    omni.kit.commands.execute("SetRelationshipTargets", relationship=relationship, targets=targets)

    def set_value(self, targets: List[Sdf.Path]):
        self.set_targets(targets)

    def _build(self):
        self._shared_targets = None

        for relationship in self._relationships:
            targets = relationship.GetTargets()
            if self._shared_targets is None:
                self._shared_targets = targets
            elif self._shared_targets != targets:
                self._shared_targets = None
                break

        with ui.VStack(spacing=2):
            if self._shared_targets is not None:
                for target in self._shared_targets:
                    with ui.HStack(spacing=2):
                        ui.StringField(name="models", read_only=True).model.set_value(target.pathString)

                        def on_remove_target(weak_self, target):
                            weak_self = weak_self()
                            if weak_self:
                                with omni.kit.undo.group():
                                    for relationship in weak_self._relationships:
                                        if relationship:
                                            omni.kit.commands.execute(
                                                "RemoveRelationshipTarget", relationship=relationship, target=target
                                            )
                                if self._on_remove_target:
                                    self._on_remove_target(target)

                        ui.Button(
                            "-",
                            enabled=self._enabled,
                            width=ui.Pixel(14),
                            clicked_fn=partial(on_remove_target, weak_self=weakref.ref(self), target=target),
                            identifier=f"remove_relationship{target.pathString.replace('/', '_')}"
                        )

                def on_add_target(weak_self):
                    weak_self = weak_self()
                    if weak_self:
                        weak_self._target_picker.show(weak_self._targets_limit - len(weak_self._shared_targets))

                within_target_limit = self._targets_limit == 0 or len(self._shared_targets) < self._targets_limit
                button = ui.Button(
                    "Add Target(s)",
                    width=ui.Pixel(30),
                    clicked_fn=partial(on_add_target, weak_self=weakref.ref(self)),
                    enabled=within_target_limit and self._enabled,
                    identifier=f"add_relationship{self._id_name}"

                )
                if not within_target_limit:
                    button.set_tooltip(
                        f"Targets limit of {self._targets_limit} has been reached. To add more target(s), remove current one(s) first."
                    )
            else:
                ui.StringField(name="models", read_only=True).model.set_value("Mixed")

    def _set_dirty(self):
        self._frame.rebuild()
