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

import carb
import omni.kit.commands
import omni.kit.context_menu
import omni.kit.undo
import omni.usd
from omni.kit.usd.layers import LayerEditMode, get_layers
from pxr import Sdf

from .usd_attribute_model import (
    GfVecAttributeModel,
    MdlEnumAttributeModel,
    TfTokenAttributeModel,
    UsdAttributeModel,
    UsdBase,
)


class AttributeContextMenuEvent:
    def __init__(self, widget, attribute_paths, stage, time_code, model, comp_index):
        self.widget = widget
        self.attribute_paths = attribute_paths
        self.stage = stage
        self.time_code = time_code
        self.model = model
        self.comp_index = comp_index
        self.type = 0


class AttributeContextMenu:
    _instance = None

    @classmethod
    def get_instance(cls):
        return weakref.proxy(cls._instance)

    def __init__(self):
        self._register_context_menus()

        AttributeContextMenu._instance = self

    def __del__(self):
        self.destroy()

    def destroy(self):
        self._copy_menu_entry = None
        self._paste_menu_entry = None
        self._delete_menu_entry = None
        self._copy_path_menu_entry = None

        AttributeContextMenu._instance = None

    def on_mouse_event(self, event: AttributeContextMenuEvent):
        # check its expected event
        if event.type != int(omni.kit.ui.MenuEventType.ACTIVATE):
            return

        # setup objects, this is passed to all functions
        objects = {
            "widget": event.widget,
            "stage": event.stage,
            "attribute_paths": event.attribute_paths,
            "time_code": event.time_code,
            "model": event.model,
            "comp_index": event.comp_index,
        }

        menu_list = omni.kit.context_menu.get_menu_dict("attribute", "omni.kit.property.usd")

        stage = event.stage
        if stage:
            usd_context = omni.usd.get_context()
            layers = get_layers(usd_context)
            edit_mode = layers.get_edit_mode()
            auto_authoring = layers.get_auto_authoring()
            layers_state = layers.get_layers_state()
            if edit_mode == LayerEditMode.SPECS_LINKING:
                per_layer_link_menu = []
                for layer in stage.GetLayerStack():
                    if auto_authoring.is_auto_authoring_layer(layer.identifier):
                        continue

                    layer_name = layers_state.get_layer_name(layer.identifier)
                    per_layer_link_menu.append(
                        {
                            "name": {
                                f"{layer_name}": [
                                    {
                                        "name": "Link",
                                        "show_fn": [],
                                        "onclick_fn": partial(self._link_attributes, True, layer.identifier),
                                    },
                                    {
                                        "name": "Unlink",
                                        "show_fn": [],
                                        "onclick_fn": partial(self._link_attributes, False, layer.identifier),
                                    },
                                ]
                            }
                        }
                    )

                menu_list.append(
                    {
                        "name": {"Layers": per_layer_link_menu},
                        "glyph": "",
                    }
                )

            menu_list.extend([{"name": ""}])
            menu_list.append(
                {
                    "name": {
                        "Locks": [
                            {
                                "name": "Lock",
                                "show_fn": [],
                                "onclick_fn": partial(self._lock_attributes, True),
                            },
                            {
                                "name": "Unlock",
                                "show_fn": [],
                                "onclick_fn": partial(self._lock_attributes, False),
                            },
                        ]
                    },
                    "glyph": "",
                }
            )

        omni.kit.context_menu.get_instance().show_context_menu("attribute", objects, menu_list)

    def _link_attributes(self, link_or_unlink, layer_identifier, objects):
        attribute_paths = objects.get("attribute_paths", None)
        if not attribute_paths:
            return

        if link_or_unlink:
            command = "LinkSpecs"
        else:
            command = "UnlinkSpecs"

        omni.kit.commands.execute(command, spec_paths=attribute_paths, layer_identifiers=layer_identifier)

    def _lock_attributes(self, lock_or_unlock, objects):
        attribute_paths = objects.get("attribute_paths", None)
        if not attribute_paths:
            return

        if lock_or_unlock:
            command = "LockSpecs"
        else:
            command = "UnlockSpecs"

        omni.kit.commands.execute(command, spec_paths=attribute_paths)

    def _register_context_menus(self):
        self._register_copy_paste_menus()
        self._register_delete_prop_menu()
        self._register_copy_prop_path_menu()

    def _register_copy_paste_menus(self):
        def can_show_copy_paste(object):
            model = object.get("model", None)
            if not model:
                return False

            return isinstance(model, UsdBase)

        def can_copy(object):
            model: UsdBase = object.get("model", None)
            if not model:
                return False

            # Only support copying during single select per OM-20206
            paths = model.get_property_paths()
            if len(paths) > 1:
                return False

            comp_index = object.get("comp_index", -1)
            return not model.is_comp_ambiguous(comp_index) and (
                isinstance(model, UsdAttributeModel)
                or isinstance(model, TfTokenAttributeModel)
                or isinstance(model, MdlEnumAttributeModel)
                or isinstance(model, GfVecAttributeModel)
            )

        def on_copy(object):
            model = object.get("model", None)
            if not model:
                return

            if isinstance(model, UsdAttributeModel):
                value_str = model.get_value_as_string(elide_big_array=False)
            elif isinstance(model, TfTokenAttributeModel):
                value_str = model.get_value_as_token()
            elif isinstance(model, MdlEnumAttributeModel):
                value_str = model.get_value_as_string()
            elif isinstance(model, GfVecAttributeModel):
                value_str = str(model._construct_vector_from_item())
            else:
                carb.log_warn("Unsupported type to copy")
                return

            omni.kit.clipboard.copy(value_str)

        menu = {
            "name": "Copy",
            "show_fn": can_show_copy_paste,
            "enabled_fn": can_copy,
            "onclick_fn": on_copy,
        }
        self._copy_menu_entry = omni.kit.context_menu.add_menu(menu, "attribute", "omni.kit.property.usd")

        # If type conversion fails, function raise exception and disables menu entry/quit paste
        def convert_type(value_type, value_str: str):
            if value_type == str or value_type == Sdf.AssetPath or value_type == Sdf.Path:
                return value_str
            if value_type == Sdf.AssetPathArray:
                # Copied AssetPathArray is in this format:
                # [@E:/USD/foo.usd@, @E:/USD/bar.usd@]
                # parse it manually
                value_str = value_str.strip("[] ")
                paths = value_str.split(", ")
                paths = [path.strip("@") for path in paths]
                return paths
            else:
                # Use an eval here so tuple/list/etc correctly converts from string. May need revisit.
                return value_type(eval(value_str))

        def can_paste(object):
            model = object.get("model", None)
            widget = object.get("widget", None)
            if not model:
                return False

            if not widget or not widget.enabled:
                return False

            comp_index = object.get("comp_index", -1)

            paste = omni.kit.clipboard.paste()
            if paste:
                try:
                    ret = True
                    if isinstance(model, UsdAttributeModel):
                        value = model.get_value_by_comp(comp_index)
                        convert_type(type(value), paste)
                    elif isinstance(model, TfTokenAttributeModel):
                        if not model.is_allowed_token(paste):
                            raise ValueError(f"Token {paste} is not allowed on this attribute")
                    elif isinstance(model, MdlEnumAttributeModel):
                        if not model.is_allowed_enum_string(paste):
                            raise ValueError(f"Enum {paste} is not allowed on this attribute")
                    elif isinstance(model, GfVecAttributeModel):
                        value = model.get_value()
                        convert_type(type(value), paste)
                    else:
                        carb.log_warn("Unsupported type to paste to")
                        ret = False
                    return ret
                except Exception as e:
                    carb.log_warn(f"{e}")
                    pass

            return False

        def on_paste(object):
            model = object.get("model", None)
            if not model:
                return

            comp_index = object.get("comp_index", -1)
            value = model.get_value_by_comp(comp_index)

            paste = omni.kit.clipboard.paste()
            if paste:
                try:
                    if (
                        isinstance(model, UsdAttributeModel)
                        or isinstance(model, TfTokenAttributeModel)
                        or isinstance(model, GfVecAttributeModel)
                    ):
                        typed_value = convert_type(type(value), paste)
                        model.set_value(typed_value)
                    elif isinstance(model, MdlEnumAttributeModel):
                        model.set_from_enum_string(paste)

                except Exception as e:
                    carb.log_warn(f"Failed to paste: {e}")

        menu = {
            "name": "Paste",
            "show_fn": can_show_copy_paste,
            "enabled_fn": can_paste,
            "onclick_fn": on_paste,
        }
        self._paste_menu_entry = omni.kit.context_menu.add_menu(menu, "attribute", "omni.kit.property.usd")

    def _register_delete_prop_menu(self):
        def can_show_delete(object):
            model = object.get("model", None)
            if not model:
                return False

            stage = model.stage
            if not stage:
                return False

            # OM-49324 Hide the Remove context menu for xformOp attributes
            # Cannot simply remove an xformOp attribute. Leave it to transform widget context menu
            paths = model.get_attribute_paths()
            for path in paths:
                prop = stage.GetPropertyAtPath(path)
                if not prop or prop.GetName().startswith("xformOp:"):
                    return False

            return isinstance(model, UsdBase)

        def can_delete(object):
            model = object.get("model", None)
            if not model:
                return False

            stage = model.stage
            if not stage:
                return False

            paths = model.get_attribute_paths()
            for path in paths:
                prop = stage.GetPropertyAtPath(path)
                if prop:
                    prim = prop.GetPrim()
                    prim_definition = prim.GetPrimDefinition()

                    # If the property is part of a schema, it cannot be completely removed. Removing it will simply reset to default value
                    prop_spec = prim_definition.GetSchemaPropertySpec(prop.GetPath().name)
                    if prop_spec:
                        return False

                else:
                    return False
            return True

        def on_delete(object):
            model = object.get("model", None)
            if not model:
                return

            paths = model.get_attribute_paths()
            with omni.kit.undo.group():
                for path in paths:
                    omni.kit.commands.execute("RemoveProperty", prop_path=path)

        menu = {
            "name": "Remove",
            "show_fn": can_show_delete,
            "enabled_fn": can_delete,
            "onclick_fn": on_delete,
        }
        self._delete_menu_entry = omni.kit.context_menu.add_menu(menu, "attribute", "omni.kit.property.usd")

    def _register_copy_prop_path_menu(self):
        def can_show_copy_path(object):
            model = object.get("model", None)
            if not model:
                return False

            stage = model.stage
            if not stage:
                return False

            return isinstance(model, UsdBase)

        def can_copy_path(object):
            model = object.get("model", None)
            if not model:
                return False

            stage = model.stage
            if not stage:
                return False

            paths = model.get_attribute_paths()
            return len(paths) == 1

        def on_copy_path(object):
            model = object.get("model", None)
            if not model:
                return

            paths = model.get_attribute_paths()

            omni.kit.clipboard.copy(paths[-1].pathString)

        menu = {
            "name": "Copy Property Path",
            "show_fn": can_show_copy_path,
            "enabled_fn": can_copy_path,
            "onclick_fn": on_copy_path,
        }
        self._copy_path_menu_entry = omni.kit.context_menu.add_menu(menu, "attribute", "omni.kit.property.usd")
