import os
import carb
import omni.ext
from functools import partial
from pathlib import Path
from pxr import Sdf, Usd, UsdGeom, UsdUI
from typing import Any, Callable
from omni.kit.property.usd.prim_selection_payload import PrimSelectionPayload

_extension_instance = None
TEST_DATA_PATH = ""

def get_instance():
    global _extension_instance
    return _extension_instance

class GeometryPropertyExtension(omni.ext.IExt):
    def __init__(self):
        self._registered = False
        self._button_menu_entry = []
        self._visual_property_widget = None
        super().__init__()

    def on_startup(self, ext_id):
        global _extension_instance
        _extension_instance = self

        self._register_widget()

        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")

        # +add menu item(s)
        from omni.kit.property.usd import PrimPathWidget

        context_menu = omni.kit.context_menu.get_instance()
        if context_menu is None:
            carb.log_error("context_menu is disabled!") # pragma: no cover
            return None                                 # pragma: no cover

        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Instanceable",
                show_fn=context_menu.is_prim_selected,
                onclick_fn=self._click_toggle_instanceable,
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Wireframe Mode",
                name_fn=partial(self._get_primvar_state, prim_name="wireframe", text_name=" Wireframe Mode"),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="wireframe"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Do Not Cast Shadows",
                name_fn=partial(
                    self._get_primvar_state, prim_name="doNotCastShadows", text_name=" Do Not Cast Shadows"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="doNotCastShadows"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Enable Shadow Terminator Fix",
                name_fn=partial(
                    self._get_primvar_state, prim_name="enableShadowTerminatorFix", text_name=" Enable Shadow Terminator Fix"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="enableShadowTerminatorFix"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Enable Fast Refraction Shadow",
                name_fn=partial(
                    self._get_primvar_state, prim_name="enableFastRefractionShadow", text_name=" Enable Fast Refraction Shadow"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="enableFastRefractionShadow"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Disable RT SSS Transmission",
                name_fn=partial(
                    self._get_primvar_state, prim_name="disableRtSssTransmission", text_name=" Disable RT SSS Transmission"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="disableRtSssTransmission"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Multimatted ID:",
                name_fn=partial(
                    self._get_primvar_state, prim_name="multimatte_id", text_name=" ID for multimatte"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="multimatte_id"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Enable Holdout Object",
                name_fn=partial(
                    self._get_primvar_state, prim_name="holdoutObject", text_name=" Enable Holdout Object"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="holdoutObject"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Invisible To Secondary Rays",
                name_fn=partial(
                    self._get_primvar_state, prim_name="invisibleToSecondaryRays", text_name=" Invisible To Secondary Rays"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="invisibleToSecondaryRays"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Is Procedural Volume",
                name_fn=partial(
                    self._get_primvar_state, prim_name="isVolume", text_name=" Is Volume"
                ),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="isVolume"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Matte Object",
                name_fn=partial(self._get_primvar_state, prim_name="isMatteObject", text_name=" Matte Object"),
                show_fn=partial(context_menu.prim_is_type, type=UsdGeom.Boundable),
                onclick_fn=partial(self._click_set_primvar, prim_name="isMatteObject"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Hide From Camera",
                name_fn=partial(
                    self._get_primvar_state, prim_name="hideForCamera", text_name=" Hide From Camera"
                ),
                show_fn=[partial(context_menu.prim_is_type, type=UsdGeom.Boundable)],
                onclick_fn=partial(self._click_set_primvar, prim_name="hideForCamera"),
            )
        )
        self._button_menu_entry.append(
            PrimPathWidget.add_button_menu_entry(
                "Rendering/Toggle Is Light",
                name_fn=partial(self._get_primvar_state, prim_name="isLight", text_name=" Is Light"),
                show_fn=[partial(context_menu.prim_is_type, type=UsdGeom.Boundable)],
                onclick_fn=partial(self._click_set_primvar, prim_name="isLight"),
            )
        )

    def on_shutdown(self): # pragma: no cover
        if self._registered:
            self._unregister_widget()

        # release menu item(s)
        from omni.kit.property.usd import PrimPathWidget

        for item in self._button_menu_entry:
            PrimPathWidget.remove_button_menu_entry(item)

        global _extension_instance
        _extension_instance = None

    def register_custom_visual_attribute(self,
                                         attribute_name: str,
                                         display_name: str,
                                         type_name: str,
                                         default_value: Any,
                                         predicate: Callable[[Any], bool] = None):
        """
        Add custom attribute with placeholder.
        """
        if self._visual_property_widget:
            self._visual_property_widget.add_custom_attribute(
                attribute_name,
                display_name,
                type_name,
                default_value,
                predicate
            )

    def deregister_custom_visual_attribute(self, attribute_name: str):
        if self._visual_property_widget:
            self._visual_property_widget.remove_custom_attribute(attribute_name)

    def _register_widget(self):
        import omni.kit.window.property as p
        from .prim_kind_widget import PrimKindWidget
        from .prim_geometry_widget import GeometrySchemaAttributesWidget, ImageableSchemaAttributesWidget

        w = p.get_window()
        if w:

            w.register_widget(
                "prim",
                "geometry",
                GeometrySchemaAttributesWidget(
                    "Geometry",
                    UsdGeom.Xformable,
                    [
                        UsdGeom.BasisCurves,
                        UsdGeom.Capsule,
                        UsdGeom.Cone,
                        UsdGeom.Cube,
                        UsdGeom.Cylinder,
                        UsdGeom.HermiteCurves,
                        UsdGeom.Mesh,
                        UsdGeom.NurbsCurves,
                        UsdGeom.NurbsPatch,
                        UsdGeom.PointInstancer,
                        UsdGeom.Points,
                        UsdGeom.Subset,
                        UsdGeom.Sphere,
                        UsdGeom.Xform,
                        UsdGeom.Gprim,
                        UsdGeom.PointBased,
                        UsdGeom.Boundable,
                        UsdGeom.Curves,
                        UsdGeom.Imageable,
                        UsdGeom.PointBased,
                        UsdGeom.Subset,
                        UsdGeom.ModelAPI,
                        UsdGeom.MotionAPI,
                        UsdGeom.PrimvarsAPI,
                        UsdGeom.XformCommonAPI,
                        UsdGeom.ModelAPI,
                        UsdUI.Backdrop,
                        UsdUI.NodeGraphNodeAPI,
                        UsdUI.SceneGraphPrimAPI,
                    ],
                    [
                        "proceduralMesh:parameterCheck",
                        "outputs:parameterCheck",
                        "refinementEnableOverride",
                        "refinementLevel",
                        "primvars:doNotCastShadows",
                        "primvars:enableShadowTerminatorFix",
                        "primvars:enableFastRefractionShadow",
                        "primvars:disableRtSssTransmission",
                        "primvars:holdoutObject",
                        "primvars:invisibleToSecondaryRays",
                        "primvars:isMatteObject",
                        "primvars:isVolume",
                        "primvars:multimatte_id",
                        "primvars:numSplits",
                        "primvars:endcaps",
                        UsdGeom.Tokens.proxyPrim,
                    ],
                    [
                        "primvars:displayColor",
                        "primvars:displayOpacity",
                        "doubleSided",
                        "purpose",
                        "visibility",
                        "xformOpOrder",
                    ],
                ),
            )

            self._visual_property_widget = ImageableSchemaAttributesWidget(
                "Visual",
                UsdGeom.Imageable,
                [],
                ["primvars:displayColor", "primvars:displayOpacity", "doubleSided", "singleSided"],
                []
            )

            w.register_widget(
                "prim",
                "geometry_imageable",
                self._visual_property_widget,
            )

            w.register_widget("prim", "kind", PrimKindWidget())
            self._registered = True

    def _unregister_widget(self): # pragma: no cover
        import omni.kit.window.property as p

        w = p.get_window()
        if w:
            w.unregister_widget("prim", "geometry")
            w.unregister_widget("prim", "geometry_imageable")
            w.unregister_widget("prim", "kind")
            self._registered = False

    def _click_set_primvar(self, payload: PrimSelectionPayload, prim_name: str):
        stage = payload.get_stage()
        if not stage:
            return

        omni.kit.commands.execute("TogglePrimVarCommand", prim_path=payload.get_paths(), prim_name=prim_name)

    def _get_primvar_state(self, objects: dict, prim_name: str, text_prefix: str = "", text_name: str = "") -> str:
        if not "stage" in objects or not "prim_list" in objects or not objects["stage"]:
            return None
        stage = objects["stage"]
        primvar_state = []
        for path in objects["prim_list"]:
            prim = stage.GetPrimAtPath(path) if isinstance(path, Sdf.Path) else path
            if prim:
                primvars_api = UsdGeom.PrimvarsAPI(prim)
                is_primvar = primvars_api.GetPrimvar(prim_name)
                if is_primvar:
                    primvar_state.append(is_primvar.Get())
                else:
                    primvar_state.append(False)

        if primvar_state == [False] * len(primvar_state):
            return f"{text_prefix}Set{text_name}"
        elif primvar_state == [True] * len(primvar_state):
            return f"{text_prefix}Clear{text_name}"
        return f"{text_prefix}Toggle{text_name}"

    def _click_toggle_instanceable(self, payload: PrimSelectionPayload):
        omni.kit.commands.execute("ToggleInstanceableCommand", prim_path=payload.get_paths())
