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

import carb
import omni.ui as ui
import omni.usd as ou

from pxr import Sdf, Usd, UsdShade

from omni.kit.property.usd.usd_property_widget import UsdPropertiesWidgetBuilder
from omni.kit.window.property.templates import HORIZONTAL_SPACING

from carb.input import KeyboardInput as Key

from ..presets import ShaderPresets
from .index_properties_base_widget import IndexPropertiesBaseWidget

class ShaderPropertiesWidget(IndexPropertiesBaseWidget):

    def __init__(self):
        super().__init__("IndeX Shader", "Shader")

        self.xac_source_asset_model = None
        self.xac_source_code_field = None
        self.id_combo_box = None

    def switch_implementation_source(self, value):
        key = "info:implementationSource"
        # self.request_rebuild() # update UI to change mode
        self.store_value(key, value, "")

    def reload_shader(self):
        prim = self.get_prim()
        key = "info:xac:sourceAsset"
        if prim.GetAttribute("info:implementationSource").Get() == "sourceAsset" and prim.GetAttribute(key).Get():
            asset = prim.GetAttribute(key).Get()
            prim.GetAttribute(key).Clear() # force an update
            custom = False
            prim.CreateAttribute(key, Sdf.ValueTypeNames.Asset, custom, Sdf.VariabilityUniform).Set(asset)

    def convert_to_source_code(self):
        prim = self.get_prim()

        shader = UsdShade.Shader(prim)
        if not shader:
            return

        XAC_SOURCE_TYPE = "xac"
        source_asset = shader.GetSourceAsset(XAC_SOURCE_TYPE)
        if not source_asset:
            return

        xac_path = source_asset.resolvedPath
        if not xac_path:
            carb.log_error("Error: Could not resolve asset path '" + str(source_asset) + "' for prim '" + prim.GetPath().pathString + "'")
            return

        try:
            with open(xac_path, 'r') as file:
                contents = file.read()
        except IOError:
            carb.log_error("Error: Could not load file '" + str(xac_path) + "' for prim '" + prim.GetPath().pathString + "'")
            return

        self.store_value("info:xac:sourceCode", contents)
        self.switch_implementation_source("sourceCode")

    def update_shader(self):
         self.store_value("info:xac:sourceCode", self.xac_source_code_field.model.as_string)

    def load_shader_preset(self, prim, preset_name):
        if prim and preset_name:
            UsdShade.Shader(prim).SetSourceCode(ShaderPresets.get(preset_name), "xac")

    def on_key_pressed(self, key: int, key_mod: int, key_down: bool):
        key = Key(key)  # Convert to enumerated type

        if key == Key.F9:
            self.update_shader()

    def _on_prop_changed(self, path):
        e = path.elementString
        if e == ".info:id" and self.id_combo_box:
            value = self.get_prim().GetAttribute("info:id").Get() or ""
            model = self.id_combo_box.model
            if value in model.get_item_list():
                model.set_current_index(model.get_item_list().index(value))
            else:
                model.set_current_index(0)
        if e == ".info:implementationSource":
            self.request_rebuild()
        elif e == ".info:xac:sourceAsset" and self.xac_source_asset_model:
            asset = self.get_prim().GetAttribute("info:xac:sourceAsset").Get()
            self.xac_source_asset_model.set_value(asset.path if asset else "")
        elif e == ".info:xac:sourceCode" and self.xac_source_code_field:
            source_code = self.get_prim().GetAttribute("info:xac:sourceCode").Get() or ""
            self.xac_source_code_field.model.set_value(source_code)

    def build_properties(self, stage, prim, prim_path):
        with ui.VStack(spacing=5):
            #
            # implementation source: asset, source code or id
            #
            with ui.HStack():
                # self.info_id_implementation_source_model = UsdPropertiesWidgetBuilder._tftoken_builder(
                #     stage,
                #     "info:implementationSource",
                #     Sdf.ValueTypeNames.Token,
                #     { 'allowedTokens': ['id', 'sourceAsset', 'sourceCode'], 'default': 'id', 'typeName': 'token' }, # metadata
                #     [ prim_path ]
                # )

                key = "info:implementationSource"
                with ui.HStack():
                    self.build_label("Implementation Source")
                    value = prim.GetAttribute(key).Get() or ""
                    self.implementation_source_combo_box = ui.ComboBox(self.create_combo_list_model(value, ['id', 'sourceAsset', 'sourceCode']), name=key)
                    self.implementation_source_combo_box.model.add_item_changed_fn(lambda m, i: self.switch_implementation_source(m.get_current_string()))

            implementation_source = prim.GetAttribute("info:implementationSource").Get() or "id"

            #
            # source asset
            #
            if implementation_source == 'sourceAsset':
                with ui.HStack():
                    models = UsdPropertiesWidgetBuilder._sdf_asset_path_builder(
                        stage,
                        "info:xac:sourceAsset",
                        Sdf.ValueTypeNames.Asset,
                        { Sdf.PropertySpec.DisplayNameKey: 'Source Asset', 'default': '', 'typeName': 'asset' },
                        [ prim_path ]
                    )
                    
                    if isinstance(models, list):
                        self.xac_source_asset_model = models[0]
                    else:
                        self.xac_source_asset_model = models
                        
                with ui.HStack():
                    ui.Button("Reload Asset", clicked_fn=self.reload_shader, name="reloadAsset")
                with ui.HStack():
                    ui.Button("Integrate Asset Contents", clicked_fn=self.convert_to_source_code, name="integrateAssetContents",
                              tooltip="Switches 'implementationSource' to 'sourceCode'")

            #
            # source code
            #
            elif implementation_source == 'sourceCode':
                key = "info:xac:sourceCode"
                with ui.HStack():
                    self.build_label("Source Code")
                with ui.HStack():
                    source_code = prim.GetAttribute(key).Get() or ""
                    self.xac_source_code_field = ui.StringField(name=key, multiline=True, tabInputAllowed=True, height=200)
                    self.xac_source_code_field.model.set_value(source_code)

                    self.xac_source_code_field.set_key_pressed_fn(self.on_key_pressed)

                    # Don't automatically update, as that would trigger recompile on each change
                    #self.xac_source_code_field.model.add_value_changed_fn(self.set_script)

                with ui.HStack():
                    ui.Button("Save and Compile [F9]", clicked_fn=self.update_shader, name="saveAndCompile")

            #
            # id
            #
            else:
                key = "info:id"
                with ui.HStack():
                    self.build_label("Implementation ID")
                    id_value = prim.GetAttribute(key).Get() or ""
                    self.id_combo_box = ui.ComboBox(self.create_combo_list_model(id_value, ['', 'nvindex:phong_gl']), name=key)
                    self.id_combo_box.model.add_item_changed_fn(lambda m, i: self.store_value(key, m.get_current_string(), ""))

            with ui.HStack():
                self.build_label("Load Shader Preset")
                preset_combo_box = ui.ComboBox(self.create_combo_list_model('', [''] + ShaderPresets.get_labels()), name='loadShaderPreset')
                preset_combo_box.model.add_item_changed_fn(lambda model, i: self.load_shader_preset(prim, model.get_current_string()))
