import omni.ext
import omni.kit.menu.utils
import omni.kit.context_menu
import omni.kit.commands
from omni.kit.menu.utils import MenuItemDescription, MenuItemOrder
from pxr import Sdf, Usd, UsdUtils, Gf
from .aov_actions import register_actions, deregister_actions


class AOVMenuExtension(omni.ext.IExt):
    def __init__(self):
        self._aov_menu_list = []

    def on_startup(self, ext_id):
        self._ext_name = omni.ext.get_extension_name(ext_id)
        register_actions(self._ext_name, AOVMenuExtension)

        self._register_menu()
        self._register_context_menu()

    def on_shutdown(self):
        deregister_actions(self._ext_name)

        omni.kit.menu.utils.remove_menu_items(self._aov_menu_list, "Create")
        self._aov_menu_list = None
        self._context_menu = None

    def _register_menu(self):

        self._aov_menu_list = [
            MenuItemDescription(name="AOV", glyph="AOV_dark.svg", appear_after=[MenuItemOrder.LAST], onclick_action=("omni.kit.menu.aov", "create_single_aov"))
        ]
        omni.kit.menu.utils.add_menu_items(self._aov_menu_list, "Create")

    def _register_context_menu(self):
        menu_dict = {
            'glyph': "AOV_dark.svg",
            'name': 'AOV',
            'show_fn': [],
            'onclick_fn': AOVMenuExtension._on_create_single_aov
        }
        self._context_menu = omni.kit.context_menu.add_menu(menu_dict, "CREATE")

    @staticmethod
    def _create_render_prim(prim_type, prim_path):
        import omni.usd
        from omni.kit.widget.layers import LayerUtils

        stage = omni.usd.get_context().get_stage()
        session_layer = stage.GetSessionLayer()
        current_edit_layer = Sdf.Find(LayerUtils.get_edit_target(stage))
        swap_edit_targets = current_edit_layer != session_layer

        rendervar_list = [  ("ldrColor", {"sourceName": "LdrColor"}), 
                            ("hdrColor", {"sourceName": "HdrColor"}), 
                            ("PtDirectIllumation", {"sourceName": "PtDirectIllumation"}), 
                            ("PtGlobalIllumination", {"sourceName": "PtGlobalIllumination"}), 
                            ("PtReflections", {"sourceName": "PtReflections"}), 
                            ("PtRefractions", {"sourceName": "PtRefractions"}), 
                            ("PtSelfIllumination", {"sourceName": "PtSelfIllumination"}), 
                            ("PtBackground", {"sourceName": "PtBackground"}), 
                            ("PtWorldNormal", {"sourceName": "PtWorldNormal"}), 
                            ("PtWorldPos", {"sourceName": "PtWorldPos"}), 
                            ("PtZDepth", {"sourceName": "PtZDepth"}), 
                            ("PtVolumes", {"sourceName": "PtVolumes"}),
                            ("PtMultiMatte0", {"sourceName": "PtMultiMatte0"}),
                            ("PtMultiMatte1", {"sourceName": "PtMultiMatte1"}),
                            ("PtMultiMatte2", {"sourceName": "PtMultiMatte2"}),
                            ("PtMultiMatte3", {"sourceName": "PtMultiMatte3"}),
                            ("PtMultiMatte4", {"sourceName": "PtMultiMatte4"}),
                            ("PtMultiMatte5", {"sourceName": "PtMultiMatte5"}),
                            ("PtMultiMatte6", {"sourceName": "PtMultiMatte6"}),
                            ("PtMultiMatte7", {"sourceName": "PtMultiMatte7"}),
                        ]

        ## create prim. /Render is on session layer but need new prims in currrent layer
        try:
            if not current_edit_layer.GetPrimAtPath('/Render'):
                omni.kit.commands.execute("CreatePrim", prim_path="/Render", prim_type="Scope", attributes={}, select_new_prim=False)
            if not current_edit_layer.GetPrimAtPath('/Render/Vars'):
                omni.kit.commands.execute("CreatePrim", prim_path="/Render/Vars", prim_type="Scope", attributes={}, select_new_prim=False)

            omni.kit.commands.execute("CreatePrim", prim_path=prim_path, prim_type="RenderProduct", attributes={}, select_new_prim=False)

            for var_name, attributes in rendervar_list:
                if not current_edit_layer.GetPrimAtPath(f"/Render/Vars/{var_name}"):
                    omni.kit.commands.execute("CreatePrim", prim_path=f"/Render/Vars/{var_name}", prim_type="RenderVar", attributes=attributes, select_new_prim=False)

            if swap_edit_targets:
                LayerUtils.set_edit_target(stage, session_layer.identifier)

            # set /Render visible & non-deletable in stage window...
            render_prim = stage.GetPrimAtPath("/Render")
            omni.usd.editor.set_hide_in_stage_window(render_prim, False)
            omni.usd.editor.set_no_delete(render_prim, True)

            # set /Render/Vars visible & non-deletable in stage window...
            vars_prim = stage.GetPrimAtPath("/Render/Vars")
            omni.usd.editor.set_hide_in_stage_window(vars_prim, False)
            omni.usd.editor.set_no_delete(vars_prim, True)

        finally:
            if swap_edit_targets:
                LayerUtils.set_edit_target(stage, current_edit_layer.identifier)

        prim = stage.GetPrimAtPath(prim_path)
        if prim:
            for var_name, attributes in rendervar_list:
                omni.kit.commands.execute("AddRelationshipTarget", relationship=prim.GetRelationship('orderedVars'), target=f"/Render/Vars/{var_name}")

            prim.CreateAttribute("resolution", Sdf.ValueTypeNames.Int2, False).Set(Gf.Vec2i(1280, 720))

    @staticmethod
    def _on_create_single_aov(objects=None):
        with omni.kit.undo.group():
            stage = omni.usd.get_context().get_stage()
            prim_path = Sdf.Path(omni.usd.get_stage_next_free_path(stage, "/Render/RenderView", False))
            AOVMenuExtension._create_render_prim("RenderProduct", prim_path)
