# 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.
#

from omni.kit.window.property.templates import SimplePropertyWidget, LABEL_WIDTH, LABEL_HEIGHT
from omni.kit.tagging import OmniKitTaggingDelegate, get_tagging_instance
from pxr import UsdShade, UsdUI

import os.path
import carb
import carb.settings
import omni.ui as ui
import omni.kit.app
import omni.kit.commands
import omni.client


EXCLUDED_NAMESPACE = "excluded"
GENERATED_NAMESPACE = "generated"
DOT_EXCLUDED_NAMESPACE = ".excluded"
DOT_GENERATED_NAMESPACE = ".generated"
EMPTY_NAMESPACE_DISPLAY_VALUE = "<empty namespace>"
SPACING = 5


class TaggingAttributesWidget(SimplePropertyWidget, OmniKitTaggingDelegate):
    def __init__(self, style):
        super().__init__(title="Tagging", collapsed=False)

        self._style = style
        self._tagging = get_tagging_instance()
        self._tagging.set_delegate(self)
        self._advanced_view = False
        self._refresh_tags = True
        self._selected_asset = None
        self._path_names = []
        self._current_tags = {}
        self._configured_namespaces = ["appearance"]
        self._init_settings()
        self._read_settings()

    def on_new_payload(self, payload):
        self._selected_asset = None
        self._refresh_tags = True

        # redraw if rebuild flag is true, or we have a new payload
        if not super().on_new_payload(payload):
            return False

        # exclude material/shaders
        for prim_path in self._payload:
            prim = self._get_prim(prim_path)
            if prim and (prim.IsA(UsdShade.Material) or
                         prim.IsA(UsdShade.Shader) or
                         prim.IsA(UsdShade.NodeGraph) or
                         prim.IsA(UsdUI.Backdrop)):
                return False

        return payload is not None and len(payload) > 0

    def clean(self):
        """Inherited function from subclass
        """
        self._tag_event_sub = None
        self._selected_asset = None
        self._sub_show_advanced_view = None
        self._sub_hidden_tags = None
        self._sub_mod_hidden_tags = None
        self._stage_event = None
        self._tagging.set_delegate(None)
        self._tagging = None

    def _get_prim(self, prim_path):
        if prim_path:
            stage = self._payload.get_stage()
            if stage:
                return stage.GetPrimAtPath(prim_path)
        return None

    def _init_settings(self):
        settings = carb.settings.get_settings()
        self._sub_show_advanced_view = omni.kit.app.SettingChangeSubscription(
            "/persistent/exts/omni.kit.property.tagging/showAdvancedTagView", lambda *_: self._on_settings_change()
        )
        self._sub_hidden_tags = omni.kit.app.SettingChangeSubscription(
            "/persistent/exts/omni.kit.property.tagging/showHiddenTags", lambda *_: self._on_settings_change()
        )
        self._sub_mod_hidden_tags = omni.kit.app.SettingChangeSubscription(
            "/persistent/exts/omni.kit.property.tagging/modifyHiddenTags", lambda *_: self._on_settings_change()
        )

    def _read_settings(self):
        settings = carb.settings.get_settings()
        self._allow_advanced_view = settings.get("/persistent/exts/omni.kit.property.tagging/showAdvancedTagView")
        self._show_hidden_tags = settings.get("/persistent/exts/omni.kit.property.tagging/showHiddenTags")
        self._modify_hidden_tags = settings.get("/persistent/exts/omni.kit.property.tagging/modifyHiddenTags")
        self._show_and_modify_hidden_tags = self._show_hidden_tags and self._modify_hidden_tags

    def _on_settings_change(self):
        self._read_settings()
        self.request_rebuild()

    def _get_tags_callback(self):
        self.request_rebuild()

    def updated_tags_for_url(self, url):
        if url in self._path_names:
            self._refresh_tags = True
            self.request_rebuild()

    def _get_prim_paths(self, payload):
        """Get the absolute paths of the full prim stack of the payload.
        """
        stage = payload.get_stage()
        if stage is None:
            return [], []

        prims = []
        path_names = []
        for path in payload:
            prim = stage.GetPrimAtPath(path)
            if prim:
                for prim_spec in prim.GetPrimStack():
                    identifier = prim_spec.layer.identifier
                    broken_url = omni.client.break_url(identifier)
                    if broken_url.scheme == "omniverse" and identifier not in path_names:
                        path_names.append(identifier)
                prims.append(prim)

        return prims, path_names

    def _prim_selection_view(self):
        """Build the combo box that can select the absolute path of the asset whose tags we will display
        """
        with ui.HStack():
            ui.Label("Reference(s):", name="label", word_wrap=True, width=LABEL_WIDTH, height=LABEL_HEIGHT)

            with ui.VStack(spacing=SPACING):

                def current_index_changed(model, _index):
                    index = model.get_item_value_model().get_value_as_int()
                    self._selected_asset = self._path_names[index]
                    self._refresh_tags = True
                    self.request_rebuild()

                if self._selected_asset is None and len(self._path_names) > 0:
                    # default - choose last asset path
                    self._selected_asset = self._path_names[-1]

                index = self._path_names.index(self._selected_asset)

                # parse identifiers and create short readable names
                shorter_path_names = []
                for path in self._path_names:
                    broken_url = omni.client.break_url(path)
                    if broken_url.path not in shorter_path_names:
                        shorter_path_names.append(os.path.basename(broken_url.path))

                if len(shorter_path_names) < len(self._path_names):
                    shorter_path_names = self._path_names

                model = ui.ComboBox(index, *shorter_path_names).model
                model.add_item_changed_fn(lambda m, i: current_index_changed(m, i))
        return shorter_path_names[index], index

    def _advanced_view_toggle(self):
        if self._allow_advanced_view:

            def toggle_advanced_view(model):
                val = model.get_value_as_bool()
                model.set_value(val)
                self._advanced_view = val
                self.request_rebuild()

            with ui.HStack(width=0, spacing=SPACING):
                ui.Label("Advanced view", name="label", word_wrap=True, width=LABEL_WIDTH, height=LABEL_HEIGHT)
                with ui.VStack(width=10):
                    ui.Spacer()
                    bool_check = ui.CheckBox(width=10, height=0, name="greenCheck")
                    bool_check.model.set_value(self._advanced_view)
                    bool_check.model.add_value_changed_fn(toggle_advanced_view)
                    ui.Spacer()
        else:
            self._advanced_view = False

    def build_items(self):
        """Main function that builds the tagging ui
        """
        self._sub_to_prim = None

        if len(self._payload) == 0:
            return

        with ui.VStack(spacing=SPACING, height=0):
            prims, path_names = self._get_prim_paths(self._payload)
            self._path_names = path_names

            if len(prims) == 0 or not prims[0].GetPath():
                ui.Label("Nothing taggable selected")
                return

            if len(self._path_names) == 0:
                ui.Label("No asset paths found.", name="label", alignment=ui.Alignment.LEFT_TOP)
                return

            if len(self._path_names) == 0:
                with ui.VStack(height=0, spacing=SPACING):
                    ui.Label(
                        "No tagging service found", name="label", word_wrap=True, width=LABEL_WIDTH, height=LABEL_HEIGHT
                    )
                    return

            if self._refresh_tags:
                # send a tag query to the tagging service asyncronously
                self._tagging.get_tags(self._path_names, self._get_tags_callback)
                self._refresh_tags = False

            path_tags = {}
            for path in self._path_names:
                # get the tags already cached from the tagging service
                results = self._tagging.tags_for_url(path)
                if results is None:
                    results = []

                path_tags[path] = self._tagging.unpack_raw_tags(results)

            # prim selection combo box
            selected_name, selected_index = self._prim_selection_view()

            # checkbox that enables advanced view (if enabled in preferences)
            self._advanced_view_toggle()

            if self._selected_asset is None or self._selected_asset not in path_tags:
                return

            self._current_tags = path_tags[self._selected_asset]
            if not self._current_tags:
                self._current_tags = {}
            for namespace in self._configured_namespaces:
                if namespace not in self._current_tags:
                    self._current_tags[namespace] = {}

            # sort by namespace
            sorted_tags = [t for t in self._current_tags]
            sorted_tags = sorted(sorted_tags)
            if "" in sorted_tags:
                sorted_tags = sorted_tags[1:] + [""]

            with ui.CollapsableFrame(title=f"Tags for {selected_name}", name="groupFrame"):
                with ui.VStack(spacing=SPACING, height=0):
                    for namespace in sorted_tags:
                        if self._advanced_view:
                            self._create_advanced_view(namespace)
                        else:
                            self._create_default_view(namespace)

    def _advanced_tag_row(self, namespace, hidden_label, key, value, editable):
        """Show one tag in advanced view along with buttons.
           A user tag will be an editable text field.
           A generated tag will be either hidden, non-editable, or editable, depending on preferences.
        """
        full_tag = self._tagging.pack_tag(namespace, hidden_label, key=key, value=value)
        if not editable:
            with ui.HStack(spacing=SPACING, width=ui.Percent(100)):
                ui.Label(full_tag, name="hidden", style=self._style)
            return

        # remove a tag or revert (after an update) a change.
        def revert_tag(field, remove_btn):
            field.model.set_value(remove_btn.name)
            remove_btn.set_tooltip("Remove")

        # toggle revert and update buttons
        def on_tag_changed(field, update_btn, remove_btn, value):
            old_value = remove_btn.name
            valid = self._tagging.validate_tag(value, self._show_and_modify_hidden_tags)
            changed = value != old_value and valid
            if not valid:
                remove_btn.set_clicked_fn(lambda f=field, rb=remove_btn: revert_tag(f, rb))
                remove_btn.set_tooltip("Revert")
                remove_btn.enabled = True
                update_btn.visible = False
            elif changed:
                remove_btn.set_clicked_fn(lambda f=field, rb=remove_btn: revert_tag(f, rb))
                remove_btn.set_tooltip("Revert")
                remove_btn.enabled = True
                update_btn.visible = True
            elif not changed:
                remove_btn.set_tooltip("Remove")
                remove_btn.set_clicked_fn(lambda f=field, rb=remove_btn: remove_tag(f, rb))
                remove_btn.enabled = True
                update_btn.visible = False

        # update a tag via changing text field and clicking "update" button.
        def update_value(field, update_button, remove_button):
            full_tag = field.model.get_value_as_string()
            if self._tagging.validate_tag(full_tag, self._show_and_modify_hidden_tags):
                old_value = remove_button.name
                omni.kit.commands.execute(
                    "UpdateTag", asset_path=self._selected_asset, old_tag=old_value, new_tag=full_tag
                )
                update_button.visible = False
                # update names so we know when the value changes
                remove_button.name = full_tag

        # remove a tag or revert (after an update) a change.
        def remove_tag(field, remove_btn):
            full_tag = remove_btn.name
            if self._tagging.validate_tag(full_tag, self._show_and_modify_hidden_tags):
                omni.kit.commands.execute("RemoveTag", asset_path=self._selected_asset, tag=full_tag)
            field.visible = False
            remove_btn.visible = False

        with ui.HStack(spacing=SPACING, width=ui.Percent(100)):
            field = ui.StringField()
            update_button = ui.Button(
                tooltip="Update",
                style=self._style["check"],
                image_height=self._style["check"]["height"],
                image_width=self._style["check"]["width"],
                width=1,
            )
            remove_button = ui.Button(
                tooltip="Remove",
                style=self._style["remove"],
                image_height=self._style["remove"]["height"],
                image_width=self._style["remove"]["width"],
                width=1,
                name=full_tag,
            )
            ui.Spacer(width=2)
            remove_button.set_clicked_fn(lambda f=field, rb=remove_button: remove_tag(f, rb))
            update_button.set_clicked_fn(lambda f=field, ub=update_button, rb=remove_button: update_value(f, ub, rb))
            update_button.visible = False

            field.model.set_value(full_tag)
            field.model.add_value_changed_fn(
                lambda m, f=field, ub=update_button, rb=remove_button: on_tag_changed(
                    f, ub, rb, m.get_value_as_string()
                )
            )

    def _advanced_new_tag_row(self, namespace):
        """This allows adding a new tag, but the text box must include the namespace (unlike in default view).
           For example: appearance.user.car
        """
        # add a new tag (must be valid)
        def add_tag(model):
            full_tag = model.get_value_as_string()
            if self._tagging.validate_tag(full_tag, self._show_and_modify_hidden_tags):
                omni.kit.commands.execute("AddTag", asset_path=self._selected_asset, tag=full_tag)
                model.set_value("")

        # toggle add button enabled
        def on_new_tag_changed(btn, value):
            valid = self._tagging.validate_tag(value, self._show_and_modify_hidden_tags)
            if valid:
                btn.enabled = True
            else:
                btn.enabled = False

        with ui.HStack(spacing=SPACING, width=ui.Percent(100)):
            model = ui.StringField(name="new_tag:" + namespace).model
            b = ui.Button(
                "Add",
                tooltip="Add new tag",
                style=self._style["plus"],
                image_height=self._style["plus"]["height"],
                image_width=self._style["plus"]["height"],
                width=60,
                clicked_fn=lambda m=model: add_tag(m),
            )
            b.enabled = False
            model.add_value_changed_fn(lambda m, btn=b: on_new_tag_changed(btn, m.get_value_as_string()))

    def _create_advanced_view(self, namespace):
        """Create the full advanced view for each namespace.
        """

        namespace_name = namespace or EMPTY_NAMESPACE_DISPLAY_VALUE
        with ui.CollapsableFrame(title=f"Advanced tag view: {namespace_name}", name="subFrame"):
            with ui.VStack(spacing=SPACING, height=0):

                # show non-hidden tags first
                for key in self._current_tags[namespace]:
                    if not key[0] == ".":
                        value = self._current_tags[namespace][key]
                        self._advanced_tag_row(namespace, None, key, value, True)

                # show hidden tags
                if self._show_hidden_tags:
                    for hidden_label in self._current_tags[namespace]:
                        if hidden_label[0] == ".":
                            for key, value in self._current_tags[namespace][hidden_label].items():
                                self._advanced_tag_row(
                                    namespace, hidden_label[1:], key, value, self._modify_hidden_tags
                                )

                self._advanced_new_tag_row(namespace)

    def _default_user_tag_row(self, namespace, key, generated_tags):
        """Show a user created tag and the button to remove it.
        """
        with ui.HStack(spacing=SPACING):

            ui.Label(key, name="label")
            ui.Spacer()
            if namespace == "":
                full_tag = key
            else:
                full_tag = ".".join([namespace, key])

            with ui.HStack(width=60):
                ui.Spacer(width=2)
                if key in generated_tags:
                    if namespace == "":
                        excluded_tag = "." + ".".join([EXCLUDED_NAMESPACE, key])
                    else:
                        excluded_tag = "." + ".".join([namespace, EXCLUDED_NAMESPACE, key])
                    b = ui.Button(
                        style=self._style["remove"],
                        width=1,
                        image_height=self._style["remove"]["height"],
                        image_width=self._style["remove"]["width"],
                        tooltip="Reject",
                        clicked_fn=lambda t=full_tag, ap=self._selected_asset: omni.kit.commands.execute(
                            "AddTag", asset_path=ap, tag=excluded_tag, excluded=t
                        ),
                    )
                else:
                    b = ui.Button(
                        style=self._style["remove"],
                        width=1,
                        image_height=self._style["remove"]["height"],
                        image_width=self._style["remove"]["width"],
                        tooltip="Remove",
                        clicked_fn=lambda t=full_tag, ap=self._selected_asset: omni.kit.commands.execute(
                            "RemoveTag", asset_path=ap, tag=t
                        ),
                    )

                ui.Spacer(width=2)

    def _default_generated_tag_row(self, namespace, key):
        """Show a generated created tag and the buttons to confirm it or reject it.
        """
        with ui.HStack(spacing=SPACING):
            ui.Label(key, name="generated", style=self._style)
            ui.Spacer()
            if namespace == "":
                user_tag = key
                exclude_tag = "." + ".".join([EXCLUDED_NAMESPACE, key])
            else:
                user_tag = ".".join([namespace, key])
                exclude_tag = "." + ".".join([namespace, EXCLUDED_NAMESPACE, key])

            with ui.HStack(width=60):
                ui.Spacer(width=2)
                a = ui.Button(
                    width=1,
                    style=self._style["check"],
                    tooltip="Confirm",
                    image_height=self._style["check"]["height"],
                    image_width=self._style["check"]["width"],
                    clicked_fn=lambda t=user_tag, ap=self._selected_asset: omni.kit.commands.execute(
                        "AddTag", asset_path=ap, tag=t
                    ),
                )
                ui.Spacer(width=15)
                b = ui.Button(
                    width=1,
                    style=self._style["remove"],
                    tooltip="Reject",
                    image_height=self._style["remove"]["height"],
                    image_width=self._style["remove"]["width"],
                    clicked_fn=lambda t=exclude_tag, ap=self._selected_asset: omni.kit.commands.execute(
                        "AddTag", asset_path=ap, tag=t
                    ),
                )
                ui.Spacer(width=2)

    def _default_new_tag_row(self, namespace):
        """Show a text box and button to create a new user tag in the given namespace.
        """

        def add_tag(namespace, model):
            value = model.get_value_as_string()
            if namespace == "":
                full_tag = value
            else:
                full_tag = ".".join([namespace, value])

            # if the tag we are adding is excluded, we want to remove the exclusion (this is undoable)
            excluded = ""
            if (
                namespace in self._current_tags
                and DOT_EXCLUDED_NAMESPACE in self._current_tags[namespace]
                and value in self._current_tags[namespace][DOT_EXCLUDED_NAMESPACE]
            ):
                if namespace == "":
                    excluded = "." + ".".join([EXCLUDED_NAMESPACE, value])
                else:
                    excluded = "." + ".".join([namespace, EXCLUDED_NAMESPACE, value])

            if self._tagging.validate_tag(full_tag, self._show_and_modify_hidden_tags):
                omni.kit.commands.execute(
                    "AddTag", asset_path=self._selected_asset, tag=full_tag, excluded=excluded
                )
                model.set_value("")

        def on_new_tag_changed(btn, value):
            if not namespace == "":
                value = ".".join([namespace, value])
            valid = self._tagging.validate_tag(value, self._show_and_modify_hidden_tags)
            if valid:
                btn.enabled = True
            else:
                btn.enabled = False

        with ui.HStack(spacing=SPACING, width=ui.Percent(100)):
            field = ui.StringField(name="new_tag:" + namespace)
            model = field.model
            with ui.HStack(spacing=SPACING, width=60):
                b = ui.Button(
                    "Add",
                    style=self._style["plus"],
                    image_height=self._style["plus"]["height"],
                    image_width=self._style["plus"]["height"],
                    tooltip="Add new tag",
                    width=60,
                    clicked_fn=lambda m=model, n=namespace: add_tag(n, m),
                )
                b.enabled = False
                model.add_value_changed_fn(lambda m, btn=b: on_new_tag_changed(btn, m.get_value_as_string()))

    def _create_default_view(self, namespace):
        """ This is the default tag view, which shows tags both added by the user and generated hidden tags.
            When generated tags are "removed" we actually do not delete them, but add another tag with
            the namespace ".excluded" instead of ".generated"
        """

        def safe_int(value):
            try:
                return int(value)
            except ValueError:
                return 0

        # populate special namespaces
        user_tags = {}
        excluded_tags = []
        generated_tags = {}
        if namespace in self._current_tags:
            user_tags = {n: self._current_tags[namespace][n] for n in self._current_tags[namespace] if n[0] != "."}
        if DOT_EXCLUDED_NAMESPACE in self._current_tags[namespace]:
            excluded_tags = self._current_tags[namespace][DOT_EXCLUDED_NAMESPACE]
        if DOT_GENERATED_NAMESPACE in self._current_tags[namespace]:
            generated_tags = {
                key: safe_int(value)
                for key, value in self._current_tags[namespace][DOT_GENERATED_NAMESPACE].items()
                if key not in excluded_tags
            }

        sorted_by_confidence = [key for key in generated_tags]
        sorted_by_confidence.sort(key=lambda x: -generated_tags[x])

        # skip empty blocks unless they are "appearance"
        if len(user_tags) == 0 and len(sorted_by_confidence) == 0 and namespace not in self._configured_namespaces:
            return

        namespace_name = namespace or EMPTY_NAMESPACE_DISPLAY_VALUE
        with ui.CollapsableFrame(title=f"Tags for {namespace_name}", name="subFrame"):
            with ui.VStack(spacing=SPACING):
                # first list user added tags (non-hidden tags)
                for key, _ in user_tags.items():
                    self._default_user_tag_row(namespace, key, generated_tags)

                # generated tags
                if sorted_by_confidence and len(sorted_by_confidence) > 0:
                    for key in sorted_by_confidence:
                        if key in user_tags or key in excluded_tags:
                            continue

                        self._default_generated_tag_row(namespace, key)

                self._default_new_tag_row(namespace)
