## 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 os
import omni.kit.app
from omni.kit.test.async_unittest import AsyncTestCase
import omni.usd

from omni.kit import ui_test
from pxr import Vt, Gf
from omni.kit.test_suite.helpers import open_stage, get_test_data_path, select_prims, wait_stage_loading, wait_for_window, arrange_windows



class PrimRelationship(AsyncTestCase):

    # Before running each test
    async def setUp(self):
        await arrange_windows("Stage", 224)
        await open_stage(get_test_data_path(__name__, "usd/relationship.usda"))
        await self._open_raw_frame(False)

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

    async def _open_raw_frame(self, state):
        await select_prims(['/World'])
        frame = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Raw USD Properties'")
        frame.widget.collapsed = state

    async def handle_select_targets(self, prim_name):
        from pxr import Sdf

        # handle assign dialog
        window_name = "Select Target(s)"
        await wait_for_window(window_name)

        # select prim
        stage_widget = ui_test.find(f"{window_name}//Frame/**/ScrollingFrame/TreeView[*].visible==True")
        await stage_widget.find(f"**/StringField[*].model.path=='{prim_name}'").click()
        await ui_test.human_delay()

        # click add
        await ui_test.find(f"{window_name}//Frame/**/Button[*].identifier=='add_button'").click()
        await ui_test.human_delay(50)

    async def test_l1_relationship(self):
        def get_buttons():
            buttons={}
            frame = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Raw USD Properties'")
            for widget in frame.find_all("**/Button[*]"):
                buttons[widget.widget.identifier] = widget
            return buttons

        def remove_add_buttons():
            buttons = get_buttons()
            for w in buttons.copy():
                if w.startswith("add_relationship_"):
                    del buttons[w]
            return buttons

        await ui_test.find("Property").focus()

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

        # select relationship prim
        await select_prims(['/World/Sphere_02'])

        # Raw USD Properties as relationship attribute is not part of any schema
        widgets = get_buttons()
        await widgets["remove_relationship_World_Sphere_01"].click()
        await ui_test.human_delay(50)

        # verify only "add_relationship"
        widgets = remove_add_buttons()
        self.assertEqual(widgets, {})

        # "add_relationship"
        await get_buttons()["add_relationship_World_Sphere_02_Attr2"].click()
        await self.handle_select_targets("/World/Sphere_01")

        # "add_relationship"
        await get_buttons()["add_relationship_World_Sphere_02_Attr2"].click()
        await self.handle_select_targets("/World/Cube_01")

        # "add_relationship"
        await get_buttons()["add_relationship_World_Sphere_02_Attr2"].click()
        await self.handle_select_targets("/World/Cube_02")

        # verify
        widgets = remove_add_buttons()
        self.assertEqual(list(widgets.keys()), ['remove_relationship_World_Sphere_01', 'remove_relationship_World_Cube_01', 'remove_relationship_World_Cube_02'])
