# 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 typing import List
from omni.kit.window.property.templates import SimplePropertyWidget
from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidget, UsdPropertyUiEntry
from .subIdentifier_utils import SubIdentifierUtils

import copy
import asyncio
import carb
import omni.ui as ui
import omni.usd

from pxr import Usd, Tf, Vt, Sdf, UsdShade, UsdUI


class UsdMaterialAttributeWidget(UsdPropertiesWidget):
    def __init__(self, schema: Usd.SchemaBase, title: str, include_names: List[str], exclude_names: List[str], schema_ignore: Usd.SchemaBase=None):
        """
        Constructor.

        Args:
            schema (Usd.SchemaBase): schema class
            schema_ignore (List[Usd.SchemaBase]): ignored schema class
            title (str): Title of the widgets on the Collapsable Frame.
            include_names (list): list of names to be included
            exclude_names (list): list of names to be excluded
        """
        super().__init__(title=title, collapsed=False)
        self._title = title
        self._schema = schema
        self._schema_ignore = schema_ignore
        self._schema_attr_names = self._schema.GetSchemaAttributeNames(True)
        self._schema_attr_names = self._schema_attr_names + include_names
        self._schema_attr_names = set(self._schema_attr_names) - set(exclude_names)
        self._additional_paths = []
        self._lookup_table = {
            # materials - should be set
            "inputs:diffuse_color_constant": {"name": "Base Color", "group": "Albedo"},
            "inputs:diffuse_texture": {"name": "Albedo Map", "group": "Albedo"},
            "inputs:albedo_desaturation": {"name": "Albedo Desaturation", "group": "Albedo"},
            "inputs:albedo_add": {"name": "Albedo Add", "group": "Albedo"},
            "inputs:albedo_brightness": {"name": "Albedo Brightness", "group": "Albedo"},
            "inputs:diffuse_tint": {"name": "Color Tint", "group": "Albedo"},
            "inputs:reflection_roughness_constant": {"name": "Roughness Amount", "group": "Reflectivity"},
            "inputs:reflection_roughness_texture_influence": {
                "name": "Roughness Map Influence",
                "group": "Reflectivity",
            },
            "inputs:reflectionroughness_texture": {"name": "Roughness Map", "group": "Reflectivity"},
            "inputs:metallic_constant": {"name": "Metallic Amount", "group": "Reflectivity"},
            "inputs:metallic_texture_influence": {"name": "Metallic Map Influence", "group": "Reflectivity"},
            "inputs:metallic_texture": {"name": "Metallic Map", "group": "Reflectivity"},
            "inputs:specular_level": {"name": "Specular", "group": "Reflectivity"},
            "inputs:enable_ORM_texture": {"name": "Enable ORM Texture", "group": "Reflectivity"},
            "inputs:ORM_texture": {"name": "ORM Map", "group": "Reflectivity"},
            "inputs:ao_to_diffuse": {"name": "AO to Diffuse", "group": "AO"},
            "inputs:ao_texture": {"name": "Ambient Occlusion Map", "group": "AO"},
            "inputs:enable_emission": {"name": "Enable Emission", "group": "Emissive"},
            "inputs:emissive_color": {"name": "Emissive Color", "group": "Emissive"},
            "inputs:emissive_color_texture": {"name": "Emissive Color map", "group": "Emissive"},
            "inputs:emissive_mask_texture": {"name": "Emissive Mask map", "group": "Emissive"},
            "inputs:emissive_intensity": {"name": "Emissive Intensity", "group": "Emissive"},
            "inputs:bump_factor": {"name": "Normal Map Strength", "group": "Normal"},
            "inputs:normalmap_texture": {"name": "Normal Map", "group": "Normal"},
            "inputs:detail_bump_factor": {"name": "Detail Normal Strength", "group": "Normal"},
            "inputs:detail_normalmap_texture": {"name": "Detail Normal Map", "group": "Normal"},
            "inputs:project_uvw": {"name": "Enable Project UVW Coordinates", "group": "UV"},
            "inputs:world_or_object": {"name": "Enable World Space", "group": "UV"},
            "inputs:uv_space_index": {"name": "UV Space Index", "group": "UV"},
            "inputs:texture_translate": {"name": "Texture Translate", "group": "UV"},
            "inputs:texture_rotate": {"name": "Texture Rotate", "group": "UV"},
            "inputs:texture_scale": {"name": "Texture Scale", "group": "UV"},
            "inputs:detail_texture_translate": {"name": "Detail Texture Translate", "group": "UV"},
            "inputs:detail_texture_rotate": {"name": "Detail Texture Rotate", "group": "UV"},
            "inputs:detail_texture_scale": {"name": "Detail Texture Scale", "group": "UV"},
            "inputs:excludeFromWhiteMode": {"group": "Material Flags", "name": "Exclude from White Mode"},
            # UsdUVTexture
            "inputs:sourceColorSpace": {"group": "UsdUVTexture", "name": "inputs:sourceColorSpace"},
            "inputs:fallback": {"group": "UsdUVTexture", "name": "inputs:fallback"},
            "inputs:file": {"group": "UsdUVTexture", "name": "inputs:file"},
            "inputs:sdrMetadata": {"group": "UsdUVTexture", "name": "inputs:sdrMetadata"},
            "inputs:rgb": {"group": "UsdUVTexture", "name": "inputs:rgb"},
            "inputs:scale": {"group": "UsdUVTexture", "name": "inputs:scale"},
            "inputs:bias": {"group": "UsdUVTexture", "name": "inputs:bias"},
            "inputs:wrapS": {"group": "UsdUVTexture", "name": "inputs:wrapS"},
            "inputs:wrapT": {"group": "UsdUVTexture", "name": "inputs:wrapT"},
            # preview surface inputs
            "inputs:diffuseColor": {"name": "Diffuse Color", "group": "Inputs"},
            "inputs:emissiveColor": {"name": "Emissive Color", "group": "Inputs"},
            "inputs:useSpecularWorkflow": {"name": "Use Specular Workflow", "group": "Inputs"},
            "inputs:specularColor": {"name": "Specular Color", "group": "Inputs"},
            "inputs:metallic": {"name": "Metallic", "group": "Inputs"},
            "inputs:roughness": {"name": "Roughness", "group": "Inputs"},
            "inputs:clearcoat": {"name": "Clearcoat", "group": "Inputs"},
            "inputs:clearcoatRoughness": {"name": "Clearcoat Roughness", "group": "Inputs"},
            "inputs:opacity": {"name": "Opacity", "group": "Inputs"},
            "inputs:opacityThreshold": {"name": "Opacity Threshold", "group": "Inputs"},
            "inputs:ior": {"name": "Index Of Refraction", "group": "Inputs"},
            "inputs:normal": {"name": "Normal", "group": "Inputs"},
            "inputs:displacement": {"name": "Displacement", "group": "Inputs"},
            "inputs:occlusion": {"name": "Occlusion", "group": "Inputs"},
            # info
            "info:id": {"name": "ID", "group": "Info"},
            "info:implementationSource": {"name": "Implementation Source", "group": "Info"},
            "info:mdl:sourceAsset": {"name": "", "group": "Info"},
            "info:mdl:sourceAsset:subIdentifier": {"name": "", "group": "Info"},
            # nodegraph
            "ui:nodegraph:node:displayColor": {"name": "Display Color", "group": "UI Properties"},
            "ui:nodegraph:node:expansionState": {"name": "Expansion State", "group": "UI Properties"},
            "ui:nodegraph:node:icon": {"name": "Icon", "group": "UI Properties"},
            "ui:nodegraph:node:pos": {"name": "Position", "group": "UI Properties"},
            "ui:nodegraph:node:size": {"name": "Size", "group": "UI Properties"},
            "ui:nodegraph:node:stackingOrder": {"name": "Stacking Order", "group": "UI Properties"},
            # backdrop
            "ui:description": {"name": "Description", "group": "UI Properties"},
        }

    def on_new_payload(self, payload):
        """
        See PropertyWidget.on_new_payload
        """
        self._anchor_prim = None
        self._shader_paths = []
        self._add_source_color_space_placeholder = False
        self._material_annotations = {}
        if not super().on_new_payload(payload):
            return False

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

        anchor_prim = None
        shader_paths = []
        mtl_paths = []
        for prim_path in self._payload:
            prim = self._get_prim(prim_path)
            if not prim or not prim.IsA(self._schema):
                return False
            if self._schema_ignore and prim.IsA(self._schema_ignore):
                return False
            shader_prim = omni.usd.get_shader_from_material(prim, True)
            if shader_prim:
                shader_paths.append(shader_prim.GetPath())

            shader = UsdShade.Shader(shader_prim if shader_prim else prim)
            asset = shader.GetSourceAsset("mdl") if shader else None
            mdl_file = asset.resolvedPath if asset else None
            if mdl_file:
                def loaded_mdl_subids(mtl_list, filename):
                    mdl_dict = {}
                    for mtl in mtl_list:
                        mdl_dict[mtl.name] = mtl.annotations

                    self._material_annotations[filename] = mdl_dict

                mtl_paths.append(mdl_file)
                asyncio.ensure_future(omni.kit.material.library.get_subidentifier_from_mdl(mdl_file=mdl_file, on_complete_fn=lambda l, f=mdl_file: loaded_mdl_subids(mtl_list=l, filename=f)))

            anchor_prim = prim

        if anchor_prim:
            stage = payload.get_stage()
            shader = UsdShade.Shader.Get(stage, anchor_prim.GetPath())
            if shader and shader.GetShaderId() == "UsdUVTexture":
                attr = anchor_prim.GetAttribute("inputs:sourceColorSpace")
                if not attr:
                    self._add_source_color_space_placeholder = True
                elif attr.GetTypeName() == Sdf.ValueTypeNames.Token:
                    tokens = attr.GetMetadata("allowedTokens")
                    if not tokens:
                        # fix missing tokens on attribute
                        attr.SetMetadata("allowedTokens", ["auto", "raw", "sRGB"])

        self._mtl_identical = True
        if mtl_paths:
            self._mtl_identical = mtl_paths.count(mtl_paths[0]) == len(mtl_paths)

        self._anchor_prim = anchor_prim
        self._shader_paths = shader_paths
        if anchor_prim:
            return True
        return False

    def _filter_props_to_build(self, attrs):
        if not self._mtl_identical:
            for attr in copy.copy(attrs):
                if attr.GetName() == "info:mdl:sourceAsset:subIdentifier":
                    attrs.remove(attr)
        return [
            attr
            for attr in attrs
            if isinstance(attr, Usd.Attribute) and (attr.GetName() in self._schema_attr_names and not attr.IsHidden()
            or UsdShade.Input.IsInput(attr) and not attr.IsHidden()
            or UsdShade.Output.IsOutput(attr) and not attr.IsHidden())
        ]

    def _customize_props_layout(self, attrs):
        from omni.kit.property.usd.custom_layout_helper import (
            CustomLayoutFrame,
            CustomLayoutGroup,
            CustomLayoutProperty,
        )
        from omni.kit.window.property.templates import (
            SimplePropertyWidget,
            LABEL_WIDTH,
            LABEL_HEIGHT,
            HORIZONTAL_SPACING,
        )

        self._additional_paths = []
        if not self._anchor_prim:
            return []

        ignore_list = ["outputs:out", "info:id", "info:implementationSource", "info:mdl:sourceAsset", "info:mdl:sourceAsset:subIdentifier"]

        # check for duplicate attributes in material/shader
        stage = self._payload.get_stage()
        add_source_color_space_placeholder = len(self._payload)
        for path in self._payload:
            prim = stage.GetPrimAtPath(path)
            shader_prim = omni.usd.get_shader_from_material(prim, True)
            if shader_prim:
                shader = UsdShade.Shader.Get(stage, shader_prim.GetPath())

                if shader and shader.GetShaderId() == "UsdUVTexture":
                    attr = shader_prim.GetAttribute("inputs:sourceColorSpace")
                    if not attr:
                        add_source_color_space_placeholder -= 1

        if add_source_color_space_placeholder == 0:
            self._add_source_color_space_placeholder = True

        # get child attributes
        attr_added = {}
        for path in self._payload:
            if not path in attr_added:
                attr_added[path.pathString] = {}
            material_prim = stage.GetPrimAtPath(path)
            shader_prim = omni.usd.get_shader_from_material(material_prim, True)

            if shader_prim:
                asyncio.ensure_future(omni.usd.get_context().load_mdl_parameters_for_prim_async(shader_prim))
                for shader_attr in shader_prim.GetAttributes():
                    if shader_attr.IsHidden():
                        continue
                    # check if shader attribute already exists in material
                    material_attr = material_prim.GetAttribute(shader_attr.GetName())
                    if material_attr.IsHidden():
                        continue
                    if material_attr:
                        # use material attribute NOT shader attribute
                        attr_added[path.pathString][shader_attr.GetName()] = material_attr
                        continue

                    # add paths to usd_changed watch list
                    prim_path = shader_attr.GetPrimPath()
                    if not prim_path in self._payload and not prim_path in self._additional_paths:
                        self._additional_paths.append(prim_path)

                    if not any(name in shader_attr.GetName() for name in ignore_list):
                        attr_added[path.pathString][shader_attr.GetName()] = shader_attr

        if not self._anchor_prim:
            return []

        # remove uncommon keys
        anchor_path = self._anchor_prim.GetPath().pathString
        anchor_attr = attr_added[anchor_path]
        common_keys = anchor_attr.keys()
        for key in attr_added:
            if anchor_path != key:
                common_keys = common_keys & attr_added[key].keys()

        def compare_metadata(meta1, meta2) -> bool:
            ignored_metadata = {"default", "colorSpace"}
            for key, value in meta1.items():
                if key not in ignored_metadata and (key not in meta2 or meta2[key] != value):
                    return False
            for key, value in meta2.items():
                if key not in ignored_metadata and (key not in meta1 or meta1[key] != value):
                    return False
            return True

        # remove any keys that metadata's differ
        for key in copy.copy(list(common_keys)):
            attr1 = attr_added[anchor_path][key]
            for path in self._payload:
                if path == anchor_path:
                    continue

                attr2 = attr_added[path.pathString][key]
                if not (
                    attr1.GetName() == attr2.GetName() and
                    attr1.GetDisplayGroup() == attr2.GetDisplayGroup() and
                    attr1.GetConnections() == attr2.GetConnections() and
                    compare_metadata(attr1.GetAllMetadata(), attr2.GetAllMetadata()) and
                    type(attr1) == type(attr2)
                ):
                    common_keys.remove(key)
                    break

        # add remaining common keys
        # use anchor_attr.keys() to keep the original order as common_keys order has changed
        for key in anchor_attr.keys():
            if key not in common_keys:
                continue

            attr = anchor_attr[key]
            attrs.append(
                UsdPropertyUiEntry(
                    attr.GetName(),
                    attr.GetDisplayGroup(),
                    attr.GetAllMetadata(),
                    type(attr),
                    prim_paths=self._shader_paths,
                ))

        if self._add_source_color_space_placeholder:
            attrs.append(
                UsdPropertyUiEntry(
                    "inputs:sourceColorSpace",
                    "UsdUVTexture",
                    {
                        Sdf.PrimSpec.TypeNameKey: "token",
                        "allowedTokens": Vt.TokenArray(3, ("auto", "raw", "sRGB")),
                        "customData": {"default": "auto"},
                    },
                    Usd.Attribute,
                )
            )

        # move inputs/outputs to own group
        for attr in attrs:
            if attr.attr_name and attr.attr_name in self._lookup_table:
                lookup = self._lookup_table[attr.attr_name]
                displayName = attr.metadata.get("displayName", None)
                if not displayName and lookup["name"]:
                    attr.override_display_name(lookup["name"])
                if not attr.display_group and lookup["group"]:
                    attr.override_display_group(lookup["group"])
            if not attr.display_group:
                if attr.prop_name.startswith("output"):
                    attr.override_display_group("Outputs")
                elif attr.prop_name.startswith("inputs"):
                    attr.override_display_group("Inputs")

            if attr.display_group == "UI Properties":
                if self._anchor_prim.IsA(UsdUI.Backdrop):
                    attr.display_group_collapsed = False
                else:
                    attr.display_group_collapsed = True
            elif attr.display_group in ["Outputs", "Material Flags"]:
                attr.display_group_collapsed = True

            #remove inputs: and outputs: prefixes from items w/o explicit display name metadata
            if Sdf.PropertySpec.DisplayNameKey not in attr.metadata:
                attr.metadata[Sdf.PropertySpec.DisplayNameKey] = attr.attr_name.split(":")[-1]

        # move groups to top...
        frame = CustomLayoutFrame(hide_extra=False)
        with frame:
            with CustomLayoutGroup("Info"):
                CustomLayoutProperty("info:implementationSource", "implementationSource")
                CustomLayoutProperty("info:mdl:sourceAsset", "mdl:sourceAsset")
                CustomLayoutProperty("info:mdl:sourceAsset:subIdentifier", "mdl:sourceAsset:subIdentifier")
                CustomLayoutProperty("info:id", "id")

            material_annotations = SubIdentifierUtils.get_annotations("description", prim, self._material_annotations)
            if material_annotations:
                with CustomLayoutGroup("Description", collapsed=True):
                    CustomLayoutProperty(None, None, build_fn=lambda s, a, m, pt, pp, al, aw, mal=material_annotations:
                                         SubIdentifierUtils.annotations_build_fn(stage=s,
                                                                                 attr_name=a,
                                                                                 metadata=m,
                                                                                 property_type=pt,
                                                                                 prim_paths=pp,
                                                                                 additional_label_kwargs=al,
                                                                                 additional_widget_kwarg=aw,
                                                                                 material_annotations=mal))

            with CustomLayoutGroup("Outputs", collapsed=True):
                CustomLayoutProperty("outputs:mdl:surface", "mdl:surface")
                CustomLayoutProperty("outputs:mdl:displacement", "mdl:displacement")
                CustomLayoutProperty("outputs:mdl:volume", "mdl:volume")
                CustomLayoutProperty("outputs:surface", "surface")
                CustomLayoutProperty("outputs:displacement", "displacement")
                CustomLayoutProperty("outputs:volume", "volume")
        attrs =  frame.apply(attrs)

        # move Outputs to last, anything appearing after needs adding to CustomLayout above
        def output_sort(attr):
            if attr.prop_name.startswith("output"):
                if attr.prop_name == "outputs:out":
                    return 2
                return 1
            return 0

        attrs.sort(key=output_sort)
        return attrs

    def build_property_item(self, stage, ui_attr: UsdPropertyUiEntry, prim_paths: List[Sdf.Path]):
        if ui_attr.prop_name == "info:mdl:sourceAsset:subIdentifier":
            if self._anchor_prim:
                attr = self._anchor_prim.GetAttribute(ui_attr.prop_name)
                material_annotations = None
                if attr:
                    shader = UsdShade.Shader(self._anchor_prim)
                    asset = shader.GetSourceAsset("mdl") if shader else None
                    mdl_file = asset.resolvedPath if asset else None
                    if mdl_file and mdl_file in self._material_annotations:
                        material_annotations = self._material_annotations[mdl_file]

                    model = SubIdentifierUtils.build_subidentifier(
                        stage,
                        attr_name=ui_attr.prop_name,
                        type_name=ui_attr.property_type,
                        metadata=ui_attr.metadata,
                        attr=attr,
                        prim_paths=prim_paths,
                    )
                    return model

        return super().build_property_item(stage, ui_attr, prim_paths)

    def _on_usd_changed(self, notice, stage):
        if stage != self._payload.get_stage():
            return

        if not self._collapsable_frame:
            return

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

        # Widget is pending rebuild, no need to check for dirty
        if self._pending_rebuild_task is not None:
            return
            
        for path in notice.GetChangedInfoOnlyPaths():
            if path.name.startswith("info:"):
                self.request_rebuild()
                return

        for path in notice.GetResyncedPaths():
            if path in self._additional_paths or path in self._payload:
                self.request_rebuild()
                return
            elif (
                path.GetPrimPath() in self._additional_paths
                or path.GetPrimPath() in self._shader_paths
                or path.GetPrimPath() in self._payload
            ):
                # If prop is added or removed, rebuild frame
                # TODO only check against the attributes this widget cares about
                if bool(stage.GetPropertyAtPath(path)) != bool(self._models.get(path)):
                    self.request_rebuild()
                    return
                else:
                    # OM-75480: For material prim, it needs special treatment. So modifying shader prim
                    # needs to refresh material prim widget also.
                    # How to avoid accessing private variable of parent class?
                    self._pending_dirty_paths.add(path)

        super()._on_usd_changed(notice=notice, stage=stage)
