import os
import carb
import omni.ext
from pathlib import Path
from pxr import Sdf, UsdShade, UsdUI
from omni.kit.material.library.treeview_model import MaterialListModel, MaterialListDelegate

TEST_DATA_PATH = ""

class MaterialPropertyExtension(omni.ext.IExt):
    def __init__(self):
        self._registered = False

    def on_startup(self, ext_id):
        self._extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path(ext_id)
        self._register_widget()
        self._hooks = []

        manager = omni.kit.app.get_app().get_extension_manager()
        extension_path = manager.get_extension_path(ext_id)
        global TEST_DATA_PATH
        TEST_DATA_PATH = Path(extension_path).joinpath("data").joinpath("tests")

    def on_shutdown(self):
        self._hooks = None
        if self._registered:
            self._unregister_widget()

    def _register_widget(self):
        import omni.kit.window.property as p
        from omni.kit.window.property.property_scheme_delegate import PropertySchemeDelegate
        from .usd_attribute_widget import UsdMaterialAttributeWidget
        from .usd_binding_widget import UsdBindingAttributeWidget

        w = p.get_window()
        if w:

            w.register_widget("prim", "material_binding", UsdBindingAttributeWidget(self._extension_path, add_context_menu=True))
            w.register_widget(
                "prim",
                "material",
                UsdMaterialAttributeWidget(
                    schema=UsdShade.Material,
                    title="Material and Shader",
                    include_names=[],
                    exclude_names=[])
            )
            w.register_widget(
                "prim",
                "shader",
                UsdMaterialAttributeWidget(
                    schema=UsdShade.Shader,
                    title="Shader",
                    include_names=["info:mdl:sourceAsset", "info:mdl:sourceAsset:subIdentifier"],
                    exclude_names=[],
                ),
            )
            # NOTE: UsdShade.Material is also a UsdShade.NodeGraph. So NodeGraphs ignore Material's
            w.register_widget(
                "prim",
                "nodegraph",
                UsdMaterialAttributeWidget(
                    schema=UsdShade.NodeGraph,
                    schema_ignore=UsdShade.Material,
                    title="Node Graph",
                    include_names=[
                        "ui:nodegraph:node:displayColor",
                        "ui:nodegraph:node:expansionState",
                        "ui:nodegraph:node:icon",
                        "ui:nodegraph:node:pos",
                        "ui:nodegraph:node:stackingOrder",
                        "ui:nodegraph:node:size"],
                    exclude_names=[],
                ),
            )
            w.register_widget(
                "prim",
                "backdrop",
                UsdMaterialAttributeWidget(
                    schema=UsdUI.Backdrop,
                    title="Backdrop",
                    include_names=[
                        "ui:nodegraph:node:displayColor",
                        "ui:nodegraph:node:expansionState",
                        "ui:nodegraph:node:icon",
                        "ui:nodegraph:node:pos",
                        "ui:nodegraph:node:stackingOrder",
                        "ui:nodegraph:node:size"],
                    exclude_names=[],
                ),
            )
            self._registered = True

    def _unregister_widget(self):
        import omni.kit.window.property as p

        w = p.get_window()
        if w:
            w.unregister_widget("prim", "material")
            w.unregister_widget("prim", "shader")
            w.unregister_widget("prim", "nodegraph")
            w.unregister_widget("prim", "backdrop")
            w.unregister_widget("prim", "material_binding")
            self._registered = False
