import os
import carb
import carb.settings
import omni.ext
import omni.kit.app
import omni.usd
import asyncio
import glob
import omni.kit.actions.core
from inspect import signature
from functools import partial
from pxr import Sdf, UsdGeom, Usd, Gf
from contextlib import suppress
from .stage_templates_menu import get_action_name

_extension_instance = None
_extension_path = None
_template_list = []
_clear_dirty_task = None

def _try_unregister_page(page: str):
    with suppress(Exception):
        import omni.kit.window.preferences

        omni.kit.window.preferences.unregister_page(page)


# must be initalized before templates
class NewStageExtension(omni.ext.IExt):
    def on_startup(self, ext_id):
        global _extension_path
        global clear_dirty_task
        _extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path(ext_id)
        _clear_dirty_task = None

        register_template("empty", self.new_stage_empty, 0)
        omni.kit.stage_templates.templates.load_templates()
        self.load_user_templates()

        # as omni.kit.stage_templates loads before the UI, it cannot depend on omni.kit.window.preferences or other extensions.
        self._preferences_page = None
        self._menu_button1 = None
        self._menu_button2 = None
        self._hooks = []

        manager = omni.kit.app.get_app().get_extension_manager()
        self._hooks.append(
            manager.subscribe_to_extension_enable(
                on_enable_fn=lambda _: self._register_page(),
                on_disable_fn=lambda _: self._unregister_page(),
                ext_name="omni.kit.window.preferences",
                hook_name="omni.kit.stage_templates omni.kit.window.preferences listener",
            )
        )

        self._stage_template_menu = None
        self._hooks.append(
            manager.subscribe_to_extension_enable(
                on_enable_fn=lambda _: self._register_menu(),
                on_disable_fn=lambda _: self._unregister_menu(),
                ext_name="omni.kit.menu.utils",
                hook_name="omni.kit.stage_templates omni.kit.menu.utils listener",
            )
        )

        self._hooks.append(
            manager.subscribe_to_extension_enable(
                on_enable_fn=lambda _: self._register_property_menu(),
                on_disable_fn=lambda _: self._unregister_property_menu(),
                ext_name="omni.kit.property.usd",
                hook_name="omni.kit.stage_templates omni.kit.property.usd listener",
            )
        )

        global _extension_instance
        _extension_instance = self

    def on_shutdown(self):
        global _extension_instance
        global _template_list
        global _clear_dirty_task

        if _clear_dirty_task:
            _clear_dirty_task.cancel()
            _clear_dirty_task = None

        unregister_template("empty")
        self._unregister_page()
        self._unregister_menu()
        self._unregister_property_menu()
        self._hooks = None
        _extension_instance = None
        _template_list = None
        for user_template in self._user_scripts:
            del user_template
        omni.kit.stage_templates.templates.unload_templates()
        self._user_scripts = None

    def _register_page(self):
        try:
            from omni.kit.window.preferences import register_page
            from .stage_templates_page import StageTemplatesPreferences

            self._preferences_page = register_page(StageTemplatesPreferences())
        except ModuleNotFoundError:
            pass

    def _unregister_page(self):
        if self._preferences_page:
            try:
                import omni.kit.window.preferences

                omni.kit.window.preferences.unregister_page(self._preferences_page)
                self._preferences_page = None
            except ModuleNotFoundError:
                pass

    def _register_menu(self):
        try:
            from .stage_templates_menu import StageTemplateMenu

            self._stage_template_menu = StageTemplateMenu()
            self._stage_template_menu.on_startup()
        except ModuleNotFoundError:
            pass

    def _unregister_menu(self):
        if self._stage_template_menu:
            try:
                self._stage_template_menu.on_shutdown()
                self._stage_template_menu = None
            except ModuleNotFoundError:
                pass

    def _has_axis(self, objects, axis):
        if not "stage" in objects:
            return False
        stage = objects["stage"]
        if stage:
            return UsdGeom.GetStageUpAxis(stage) != axis
        else:
            carb.log_error("_click_set_up_axis stage not found")
        return False

    def _click_set_up_axis(self, payload, axis):
        stage = payload.get_stage()
        if stage:
            rootLayer = stage.GetRootLayer()
            if rootLayer:
                rootLayer.SetPermissionToEdit(True)
                with Usd.EditContext(stage, rootLayer):
                    UsdGeom.SetStageUpAxis(stage, axis)
                    from omni.kit.property.usd import PrimPathWidget

                    PrimPathWidget.rebuild()
            else:
                carb.log_error("_click_set_up_axis rootLayer not found")
        else:
            carb.log_error("_click_set_up_axis stage not found")

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

        context_menu = omni.kit.context_menu.get_instance()
        if context_menu is None:
            self._menu_button1 = None
            self._menu_button2 = None
            carb.log_error("context_menu is disabled!")
            return None

        self._menu_button1 = PrimPathWidget.add_button_menu_entry(
            "Stage/Set up axis +Y",
            show_fn=partial(self._has_axis, axis=UsdGeom.Tokens.y),
            onclick_fn=partial(self._click_set_up_axis, axis=UsdGeom.Tokens.y),
            add_to_context_menu=False,
        )
        self._menu_button2 = PrimPathWidget.add_button_menu_entry(
            "Stage/Set up axis +Z",
            show_fn=partial(self._has_axis, axis=UsdGeom.Tokens.z),
            onclick_fn=partial(self._click_set_up_axis, axis=UsdGeom.Tokens.z),
            add_to_context_menu=False,
        )

    def _unregister_property_menu(self):
        if self._menu_button1 or self._menu_button2:
            try:
                from omni.kit.property.usd import PrimPathWidget

                if self._menu_button1:
                    PrimPathWidget.remove_button_menu_entry(self._menu_button1)
                    self._menu_button1 = None

                if self._menu_button2:
                    PrimPathWidget.remove_button_menu_entry(self._menu_button2)
                    self._menu_button2 = None
            except ModuleNotFoundError:
                pass

    def new_stage_empty(self, rootname):
        pass

    def load_user_templates(self):
        settings = carb.settings.get_settings()
        template_paths = settings.get("/persistent/app/newStage/templatePath")

        # Create template directories
        original_umask = os.umask(0)
        for path in template_paths:
            path = carb.tokens.get_tokens_interface().resolve(path)
            if not os.path.isdir(path):
                try:
                    os.makedirs(path)
                except Exception:
                    carb.log_error(f"Failed to create directory {path}")
        os.umask(original_umask)

        # Load template scripts
        self._user_scripts = []
        for path in template_paths:
            for full_path in glob.glob(f"{path}/*.py"):
                try:
                    with open(os.path.normpath(full_path)) as f:
                        user_script = f.read()
                        carb.log_verbose(f"loaded new_stage template {full_path}")
                        exec(user_script)
                        self._user_scripts.append(user_script)
                except Exception as e:
                    carb.log_error(f"error loading {full_path}: {e}")


def register_template(name, new_stage_fn, group=0, rebuild=True):
    """Register new_stage Template

    Args:
        param1 (str): template name
        param2 (callable): function to create template
        param3 (int): group number. User by menu to split into groups with seperators

    Returns:
        bool: True for success, False when template already exists.
    """
    # check if element exists
    global _template_list

    exists = get_stage_template(name)
    if exists:
        carb.log_warn(f"template {name} already exists")
        return False

    try:
        exists = _template_list[group]
    except IndexError:
        _template_list.insert(group, {})

    _template_list[group][name] = (name, new_stage_fn)

    omni.kit.actions.core.get_action_registry().register_action(
        "omni.kit.stage.templates",
        f"create_new_stage_{get_action_name(name)}",
        lambda t=name: omni.kit.window.file.new(t),
        display_name=f"Create New Stage {name}",
        description=f"Create New Stage {name}",
        tag="Create Stage Template",
    )

    if rebuild:
        rebuild_stage_template_menu()

    return True


def unregister_template(name, rebuild: bool=True):
    """Remove registered new_stage Template

    Args:
        param1 (str): template name

    Returns:
        nothing
    """
    global _template_list
    if _template_list is not None:
        for templates in _template_list:
            if name in templates:
                del templates[name]

    if rebuild:
        rebuild_stage_template_menu()

    omni.kit.actions.core.get_action_registry().deregister_action("omni.kit.stage.templates", f"create_new_stage_{get_action_name(name)}")


def rebuild_stage_template_menu() -> None:
    if _extension_instance and _extension_instance._stage_template_menu:
        _extension_instance._stage_template_menu._rebuild_sub_menu()


def get_stage_template_list():
    """Get list of loaded new_stage templates

    Args:
        none

    Returns:
        list: list of groups of new_stage template names & create function pointers
    """
    global _template_list
    if not _template_list or len(_template_list) == 0:
        return None
    return _template_list


def get_stage_template(name):
    """Get named new_stage template & create function pointer

    Args:
        param1 (str): template name

    Returns:
        tuple: new_stage template name & create function pointer
    """
    global _template_list
    if not _template_list:
        return None

    for templates in _template_list:
        if name in templates:
            return templates[name]
    return None


def get_default_template():
    """Get name of default new_stage template. Used when new_stage is called without template name specified

    Args:
        param1 (str): template name

    Returns:
        str: new_stage template name
    """
    settings = carb.settings.get_settings()
    return settings.get("/persistent/app/newStage/defaultTemplate")


def new_stage_finalize(create_result, error_message, usd_context, template=None, on_new_stage_fn=None):
    global _clear_dirty_task

    if _clear_dirty_task:
        _clear_dirty_task.cancel()
        _clear_dirty_task = None

    if create_result:
        stage = usd_context.get_stage()

        # already finilized
        if stage.HasDefaultPrim():
            return

        with Usd.EditContext(stage, stage.GetRootLayer()):
            settings = carb.settings.get_settings()
            default_prim_name = settings.get("/persistent/app/stage/defaultPrimName")
            up_axis = settings.get("/persistent/app/stage/upAxis")
            time_codes_per_second = settings.get_as_float("/persistent/app/stage/timeCodesPerSecond")
            time_code_range = settings.get("/persistent/app/stage/timeCodeRange")
            if time_code_range is None:
                time_code_range = [0, 100]

            rootname = f"/{default_prim_name}"

            # Set up axis
            if up_axis == "Y" or up_axis == "y":
                UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.y)
            else:
                UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)

            # Set timecodes per second
            stage.SetTimeCodesPerSecond(time_codes_per_second)

            # Start and end time code
            if time_code_range:
                stage.SetStartTimeCode(time_code_range[0])
                stage.SetEndTimeCode(time_code_range[1])

            # Create defaultPrim
            default_prim = UsdGeom.Xform.Define(stage, Sdf.Path(rootname)).GetPrim()
            if not default_prim:
                carb.log_error("Failed to create defaultPrim at {rootname}")
                return
            stage.SetDefaultPrim(default_prim)

            if template is None:
                template = get_default_template()

            # Run script
            item = get_stage_template(template)
            if item and item[1]:
                try:
                    create_fn = item[1]
                    sig = signature(create_fn)
                    if len(sig.parameters) == 1:
                        create_fn(rootname)
                    elif len(sig.parameters) == 2:
                        create_fn(rootname, usd_context.get_name())
                    else:
                        carb.log_error(f"template {template} has incorrect parameter count")

                except Exception as error:
                    carb.log_error(f"exception in {template} {error}")

    omni.kit.undo.clear_stack()
    usd_context.set_pending_edit(False)

    # Clear the stage dirty state again, as bound materials can set it
    async def clear_dirty(usd_context):
        import omni.kit.app

        await omni.kit.app.get_app().next_update_async()
        await omni.kit.app.get_app().next_update_async()
        usd_context.set_pending_edit(False)
        _clear_dirty_task = None
        if on_new_stage_fn:
            on_new_stage_fn(create_result, error_message)

    _clear_dirty_task = asyncio.ensure_future(clear_dirty(usd_context))


def new_stage(on_new_stage_fn=None, template="empty", usd_context=None):
    """Execute new_stage

    Args:
        param2 (str): template name, if not specified default name is used
        param3 (omni.usd.UsdContext): usd_context, the usd_context to create new stage

    Returns:
        nothing
    """

    if usd_context is None:
        usd_context = omni.usd.get_context()

    if on_new_stage_fn is not None:
        carb.log_warn(
            "omni.kit.stage_templates.new_stage(callback, ...) is deprecated. \
            Use `omni.kit.stage_templates.new_stage_with_callback` instead."
        )
        new_stage_with_callback(on_new_stage_fn, template, usd_context)
    else:
        new_stage_finalize(usd_context.new_stage(), "", usd_context, template=template, on_new_stage_fn=None)


def new_stage_with_callback(on_new_stage_fn=None, template="empty", usd_context=None):
    """Execute new_stage

    Args:
        param1 (callable): callback when new_stage is created
        param2 (str): template name, if not specified default name is used
        param3 (omni.usd.UsdContext): usd_context, the usd_context to create new stage

    Returns:
        nothing
    """

    if usd_context is None:
        usd_context = omni.usd.get_context()

    usd_context.new_stage_with_callback(
        partial(new_stage_finalize, usd_context=usd_context, template=template, on_new_stage_fn=on_new_stage_fn)
    )


async def new_stage_async(template="empty", usd_context=None):
    """Execute new_stage asynchronously

    Args:
        param1 (str): template name, if not specified default name is used
        param2 (omni.usd.UsdContext): usd_context, the usd_context to create new stage

    Returns:
        awaitable object until stage is created
    """

    if usd_context is None:
        usd_context = omni.usd.get_context()

    f = asyncio.Future()

    def on_new_stage(result, err_msg):
        if not f.done():
            f.set_result((result, err_msg))

    new_stage_with_callback(on_new_stage, template, usd_context)
    return await f


def get_extension_path(sub_directory):
    global _extension_path
    path = _extension_path
    if sub_directory:
        path = os.path.normpath(os.path.join(path, sub_directory))
    return path


def set_transform_helper(
    prim_path,
    translate=Gf.Vec3d(0, 0, 0),
    euler=Gf.Vec3d(0, 0, 0),
    scale=Gf.Vec3d(1, 1, 1),
):
    rotation = (
        Gf.Rotation(Gf.Vec3d.ZAxis(), euler[2])
        * Gf.Rotation(Gf.Vec3d.YAxis(), euler[1])
        * Gf.Rotation(Gf.Vec3d.XAxis(), euler[0])
    )
    xform = Gf.Matrix4d().SetScale(scale) * Gf.Matrix4d().SetRotate(rotation) * Gf.Matrix4d().SetTranslate(translate)
    omni.kit.commands.execute(
        "TransformPrim",
        path=prim_path,
        new_transform_matrix=xform,
    )
