# Copyright (c) 2021, 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 omni.kit.commands
import omni.kit.viewport_legacy
import omni.usd

from pxr import Sdf, UsdGeom, Usd
from typing import Union


def get_viewport_window_from_name(viewport_name: str):
    vp = omni.kit.viewport_legacy.get_viewport_interface()
    instance = vp.get_instance(viewport_name)
    if instance:
        return vp.get_viewport_window(instance)
    return None


class SetActiveViewportCameraCommand(omni.kit.commands.Command):
    """
    Sets Viewport's actively bound camera to given camera at give path.

    Args:
        new_active_cam_path (Union[str, Sdf.Path): new camera path to bind to viewport.
        viewport_name (str): name of the viewport to set active camera (for multi-viewport).
    """

    def __init__(self, new_active_cam_path: Union[str, Sdf.Path], viewport_name: str = ""):
        self._new_active_cam_path = str(new_active_cam_path)
        self._viewport_window = get_viewport_window_from_name(viewport_name)
        if not viewport_name and not self._viewport_window:
            self._viewport_window = omni.kit.viewport_legacy.get_default_viewport_window()
        self._prev_active_camera = None

    def do(self):
        if not self._viewport_window:
            return

        self._prev_active_camera = self._viewport_window.get_active_camera()
        self._viewport_window.set_active_camera(self._new_active_cam_path)

    def undo(self):
        if not self._viewport_window:
            return

        self._viewport_window.set_active_camera(self._prev_active_camera)


class DuplicateFromActiveViewportCameraCommand(omni.kit.commands.Command):
    """
    Duplicates Viewport's actively bound camera and bind active camera to the duplicated one.

    Args:
        viewport_name (str): name of the viewport to set active camera (for multi-viewport).
    """

    def __init__(self, viewport_name: str = ""):
        self._viewport_name = viewport_name
        self._viewport_window = get_viewport_window_from_name(viewport_name)
        if not viewport_name and not self._viewport_window:
            self._viewport_window = omni.kit.viewport_legacy.get_default_viewport_window()
        self._usd_context_name = ""
        self._usd_context = omni.usd.get_context("")  # TODO get associated usd_context from Viewport itself
        self._prev_active_camera = None

    def do(self):
        if not self._viewport_window:
            return

        self._prev_active_camera = self._viewport_window.get_active_camera()
        if self._prev_active_camera:
            stage = self._usd_context.get_stage()
            target_path = omni.usd.get_stage_next_free_path(stage, "/Camera", True)
            old_prim = stage.GetPrimAtPath(self._prev_active_camera)
            omni.kit.commands.execute("CreatePrim", prim_path=target_path, prim_type="Camera")
            new_prim = stage.GetPrimAtPath(target_path)
            if old_prim and new_prim:
                # copy attributes
                timeline = omni.timeline.get_timeline_interface()
                timecode = timeline.get_current_time() * stage.GetTimeCodesPerSecond()
                for attr in old_prim.GetAttributes():
                    value = attr.Get(timecode)
                    if value is not None:
                        new_prim.CreateAttribute(attr.GetName(), attr.GetTypeName()).Set(value)
                
                default_prim = stage.GetDefaultPrim()
                # OM-35156: It's possible that transform of default prim is not identity.
                # And it tries to copy a camera that's outside of the default prim tree.
                # It needs to be transformed to the default prim space. 
                if default_prim and default_prim.IsA(UsdGeom.Xformable):
                    default_prim_world_mtx = omni.usd.get_world_transform_matrix(default_prim, timecode)
                    old_camera_prim_world_mtx = omni.usd.get_world_transform_matrix(old_prim, timecode)
                    new_camera_prim_local_mtx = old_camera_prim_world_mtx * default_prim_world_mtx.GetInverse() 
                    omni.kit.commands.execute(
                        "TransformPrim", path=new_prim.GetPath(),
                        new_transform_matrix=new_camera_prim_local_mtx
                    )

                omni.usd.editor.set_no_delete(new_prim, False)
                omni.kit.commands.execute(
                    "SetActiveViewportCamera", new_active_cam_path=target_path, viewport_name=self._viewport_name
                )

    def undo(self):
        pass


omni.kit.commands.register_all_commands_in_module(__name__)
