## Copyright (c) 2022, 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.test
import os
import sys
import unittest
import omni.kit.app
from omni.kit.test.async_unittest import AsyncTestCase
import omni.usd
from pathlib import Path

from omni.kit import ui_test
from pxr import Gf, UsdGeom
from omni.kit.test_suite.helpers import (
    open_stage,
    get_test_data_path,
    select_prims,
    wait_stage_loading,
    arrange_windows
    )


def include_op_code(xform : UsdGeom.Xformable, op_name):
        xform_ops = xform.GetOrderedXformOps()
        for op in xform_ops:
            print(op.GetOpName())
            if op.GetOpName() == op_name:
                return True
        return False

class TransformContextMenu(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows("Stage", 64)

        extension_root_folder = Path(omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__))
        self._usd_path = extension_root_folder.joinpath("data/test_data/test_map")

    # After running each test
    async def tearDown(self):
        await wait_stage_loading()

    @unittest.skipIf(sys.platform.startswith("linux"), "Pyperclip fails on some TeamCity agents")
    async def test_transform_context_menu(self):
        await open_stage(get_test_data_path(__name__, "usd/bound_shapes.usda"))
        await wait_stage_loading()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        # get prim attributes
        torus_translate = stage.GetPrimAtPath("/World/Torus").GetAttribute('xformOp:translate')
        cone_translate = stage.GetPrimAtPath("/World/Cone").GetAttribute('xformOp:translate')
        torus_rotate = stage.GetPrimAtPath("/World/Torus").GetAttribute('xformOp:rotateXYZ')
        cone_rotate = stage.GetPrimAtPath("/World/Cone").GetAttribute('xformOp:rotateXYZ')
        torus_scale = stage.GetPrimAtPath("/World/Torus").GetAttribute('xformOp:scale')
        cone_scale = stage.GetPrimAtPath("/World/Cone").GetAttribute('xformOp:scale')

        # verify transforms different
        self.assertNotEqual(torus_translate.Get(), cone_translate.Get())
        self.assertNotEqual(torus_rotate.Get(), cone_rotate.Get())
        self.assertNotEqual(torus_scale.Get(), cone_scale.Get())

        # select torus
        await select_prims(["/World/Torus"])
        await ui_test.human_delay()

        # right click on transform header
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Transform'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        # context menu copy
        await ui_test.select_context_menu("Copy All Property Values in Transform", offset=ui_test.Vec2(10, 10))

        # select cone
        await select_prims(["/World/Cone"])
        await ui_test.human_delay()

        # right click on Transform header
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Transform'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        # context menu paste
        await ui_test.select_context_menu("Paste All Property Values to Transform", offset=ui_test.Vec2(10, 10))
        # verify transforms same
        self.assertEqual(torus_translate.Get(), cone_translate.Get())
        self.assertEqual(torus_rotate.Get(), cone_rotate.Get())
        self.assertEqual(torus_scale.Get(), cone_scale.Get())

        # select torus
        await select_prims(["/World/Torus"])
        await ui_test.human_delay()

        # right click on Transform header
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Transform'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        # context menu reset
        await ui_test.select_context_menu("Reset All Property Values in Transform", offset=ui_test.Vec2(10, 10))
        # verify transforms reset
        self.assertEqual(torus_translate.Get(), Gf.Vec3d(0, 0, 0))
        self.assertEqual(torus_rotate.Get(), Gf.Vec3d(0, 0, 0))
        self.assertEqual(torus_scale.Get(), Gf.Vec3d(1, 1, 1))

    async def open_property_map(self):
        usd_context = omni.usd.get_context()
        test_file_path = self._usd_path.joinpath("test_transform_property.usda").absolute()
        await usd_context.open_stage_async(str(test_file_path))
        await wait_stage_loading()

    async def test_translate_context_menu(self):

        await self.open_property_map()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        #test disable/delete
        #right click on translate
        #select cone
        await select_prims(["/World/Cone"])
        await ui_test.human_delay()

        prim = stage.GetObjectAtPath("/World/Cone")
        xform = UsdGeom.Xformable(prim)

        #test translate disable/enable/delete
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Translate'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)

        # context menu reset
        await ui_test.select_context_menu("Disable", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        self.assertTrue(not include_op_code(xform, "xformOp:translate"))

        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Translate'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)

        # context menu reset
        await ui_test.select_context_menu("Enable", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        self.assertTrue(include_op_code(xform, "xformOp:translate"))

        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Translate'")
        #now test delete
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        await ui_test.select_context_menu("Delete", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        attr = prim.GetAttribute("xformOp:translate")
        self.assertTrue(not attr)

    async def test_pivot_context_menu(self):
        await self.open_property_map()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        #test disable/delete
        #right click on translate
        #select cone
        await select_prims(["/World/Cone"])
        await ui_test.human_delay()

        prim = stage.GetObjectAtPath("/World/Cone")
        xform = UsdGeom.Xformable(prim)

        #test translate:pivot disable/enable/delete
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Translate:pivot'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)

        # context menu reset
        await ui_test.select_context_menu("Disable", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        self.assertTrue(not include_op_code(xform, "xformOp:translate:pivot"))

        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Translate:pivot'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)

        # context menu resety
        await ui_test.select_context_menu("Enable", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        self.assertTrue(include_op_code(xform, "xformOp:translate:pivot"))

        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Translate:pivot'")
        #now test delete
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        await ui_test.select_context_menu("Delete", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        attr = prim.GetAttribute("xformOp:translate:pivot")
        self.assertTrue(not attr)

    async def test_scale_context_menu(self):

        await self.open_property_map()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        #test disable/delete
        #right click on translate
        #select cone
        await select_prims(["/World/Cone"])
        await ui_test.human_delay()

        prim = stage.GetObjectAtPath("/World/Cone")
        xform = UsdGeom.Xformable(prim)

        #test translate:pivot disable/enable/delete
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Scale'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)

        # context menu reset
        await ui_test.select_context_menu("Disable", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        self.assertTrue(not include_op_code(xform, "xformOp:scale"))

        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Scale'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)

        # context menu resety
        await ui_test.select_context_menu("Enable", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        self.assertTrue(include_op_code(xform, "xformOp:scale"))

        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Scale'")
        #now test delete
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        await ui_test.select_context_menu("Delete", offset=ui_test.Vec2(10, 10))
        await ui_test.human_delay()
        #make sure transslate is gone
        attr = prim.GetAttribute("xformOp:scale")
        self.assertTrue(not attr)

    async def test_rotate_context_menu(self):
        await self.open_property_map()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        #test disable/delete
        #right click on translate
        #select cone
        await select_prims(["/World/Cone"])
        await ui_test.human_delay()

        prim = stage.GetObjectAtPath("/World/Cone")
        xform = UsdGeom.Xformable(prim)

        #test rotate rotation order change
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Rotate'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10))

        #find ZYX label
        widget = ui_test.find("RotationOrder//Frame/**/Label[*].text=='ZYX'")

        #switch and make sure it's changed
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10))
        self.assertTrue(not include_op_code(xform, "xformOp:rotateXYZ"))
        self.assertTrue(include_op_code(xform, "xformOp:rotateZYX"))

        #chgange it to YZX
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*]/**/Label[*].text=='Rotate'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10))

        #find YZX label
        widget = ui_test.find("RotationOrder//Frame/**/Label[*].text=='YZX'")

        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10))
        self.assertTrue(not include_op_code(xform, "xformOp:rotateZYX"))
        self.assertTrue(include_op_code(xform, "xformOp:rotateYZX"))