import asyncio
import carb
import omni.usd
import omni.ui as ui
from omni.kit import ui_test
from omni.kit.menu.utils import MenuItemDescription, MenuLayout


async def edit_test_func_edit_select_select_all(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere', '/Xform', '/Cube'])

    # use menu
    await ui_test.menu_click("Edit/Select/Select All")

    # verify all prims are selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == prim_list)


async def edit_test_hotkey_func_edit_select_select_all(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere', '/Xform', '/Cube'])

    # hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])

    # verify all prims are selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == prim_list)


async def edit_test_func_edit_select_select_none(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere', '/Xform', '/Cube'])

    # select prims
    omni.usd.get_context().get_selection().set_selected_prim_paths(prim_list, True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Select/Select None")

    # verify all prims are selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == [])


async def edit_test_hotkey_func_edit_select_select_none(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere', '/Xform', '/Cube'])

    # select prims
    omni.usd.get_context().get_selection().set_selected_prim_paths(prim_list, True)
    await ui_test.human_delay()

    # hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])

    # verify all prims are selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == [])


async def edit_test_func_edit_select_select_invert(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere', '/Xform', '/Cube'])

    # select prims
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/Xform'], True)
    await ui_test.human_delay()

    # verify prim is selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == ['/Xform'])

    # use menu
    await ui_test.menu_click("Edit/Select/Select Invert")

    # verify all prims are selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == ['/Sphere', '/Cube'])


async def edit_test_hotkey_func_edit_select_select_invert(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere', '/Xform', '/Cube'])

    # select prims
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/Xform'], True)
    await ui_test.human_delay()

    # verify prim is selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == ['/Xform'])

    # hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])

    # verify all prims are selected
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == ['/Sphere', '/Cube'])


async def edit_test_func_edit_select_select_parent(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    # verify no prims
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == [])

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere'])

    # select prim
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/Sphere'], True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Group")
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/Group/Sphere'], True)

    # use menu
    await ui_test.menu_click("Edit/Select/Select Parent")

    # verify
    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == ['/Group'])


async def edit_test_hotkey_func_edit_select_select_parent(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()

    # verify no prims
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == [])

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    await ui_test.human_delay()

    # verify prims were created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere'])

    # select prim
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/Sphere'], True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Group")
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/Group/Sphere'], True)

    # hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])

    tester.assertTrue(omni.usd.get_context().get_selection().get_selected_prim_paths() == ['/Group'])


async def edit_test_func_edit_select_select_leaf(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

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

    # select prim
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/World'], True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Select/Select Leaf")
    await ui_test.human_delay()

    # verify world leafs
    leafs = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertTrue(leafs == ['/World/Cone', '/World/Cube', '/World/Cylinder', '/World/Disk', '/World/Plane', '/World/Scope/Capsule', '/World/Scope/Cylinder', '/World/Sphere', '/World/Torus', '/World/defaultLight'])

    # select prim
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/World/Scope'], True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Select/Select Leaf")
    await ui_test.human_delay()

    # verify scope leafs
    leafs = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertTrue(leafs == ['/World/Scope/Capsule', '/World/Scope/Cylinder'])


async def edit_test_func_edit_select_select_hierarchy(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

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

    # select prim
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/World'], True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Select/Select Hierarchy")
    await ui_test.human_delay()

    # verify world leafs
    hierarchy = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertTrue(hierarchy == ['/World', '/World/Cone', '/World/Cube', '/World/Cylinder', '/World/Disk', '/World/Plane', '/World/Scope', '/World/Scope/Capsule', '/World/Scope/Cylinder', '/World/Sphere', '/World/Torus', '/World/defaultLight'])

    # select prim
    omni.usd.get_context().get_selection().set_selected_prim_paths(['/World/Scope'], True)
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Select/Select Hierarchy")
    await ui_test.human_delay()

    # verify scope leafs
    hierarchy = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertTrue(hierarchy == ['/World/Scope', '/World/Scope/Capsule', '/World/Scope/Cylinder'])


async def edit_test_func_edit_select_select_similar(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

    stage =  omni.usd.get_context().get_stage()
    # NOTE: prim names do not reflect prim types
    expected = {
        "/World": ['/World'],
        "/World/defaultLight": ['/World/defaultLight'],
        "/World/Cone": ['/World/Cone', '/World/Cylinder', '/World/Torus', '/World/Scope/Capsule'],
        "/World/Cube": ['/World/Cube', '/World/Plane'],
        "/World/Cylinder": ['/World/Cone', '/World/Cylinder', '/World/Torus', '/World/Scope/Capsule'],
        "/World/Disk": ['/World/Disk', '/World/Sphere', '/World/Scope/Cylinder'],
        "/World/Plane": ['/World/Cube', '/World/Plane'],
        "/World/Sphere": ['/World/Disk', '/World/Sphere', '/World/Scope/Cylinder'],
        "/World/Torus": ['/World/Cone', '/World/Cylinder', '/World/Torus', '/World/Scope/Capsule'],
        "/World/Scope": ['/World/Scope'],
        "/World/Scope/Capsule": ['/World/Cone', '/World/Cylinder', '/World/Torus', '/World/Scope/Capsule'],
        "/World/Scope/Cylinder": ['/World/Disk', '/World/Sphere', '/World/Scope/Cylinder'],
        }

    for prim_name in tester.get_stage_prims():
        # select prim
        omni.usd.get_context().get_selection().set_selected_prim_paths([prim_name], True)
        await ui_test.human_delay()

        # use menu
        await ui_test.menu_click("Edit/Select/Select Similar")
        await ui_test.human_delay()

        similar = omni.usd.get_context().get_selection().get_selected_prim_paths()
        tester.assertTrue(expected[prim_name] == similar)


async def edit_test_func_edit_select_create_selection_set(tester, menu_item: MenuItemDescription):
    from carb.input import KeyboardInput

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

    ## create prim
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Sphere")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Xform")
    omni.kit.commands.execute("CreatePrimWithDefaultXform", prim_type="Cube")
    await ui_test.human_delay()

    # use menu
    await ui_test.menu_click("Edit/Select/Select All")
    await ui_test.human_delay()
    await ui_test.menu_click("Edit/Select/Create Selection Set")
    await ui_test.human_delay()
    widget = ui_test.find("Selection Set Name//Frame/**/StringField[*]")
    await widget.click(double=True)
    await ui_test.human_delay()
    await ui_test.emulate_keyboard_press(KeyboardInput.DEL)
    await ui_test.human_delay()
    await ui_test.emulate_char_press("All")
    await ui_test.human_delay()
    await ui_test.emulate_keyboard_press(KeyboardInput.ENTER)
    await ui_test.human_delay()

    # verify "All"
    omni.usd.get_context().get_selection().set_selected_prim_paths([], True)
    await ui_test.human_delay()
    await ui_test.menu_click("Edit/Select/Select Set/All")
    await ui_test.human_delay()
    selected_prims = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertTrue(selected_prims == ['/Sphere', '/Xform', '/Cube'])


async def edit_test_func_edit_select_select_by_kind_subcomponent(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

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

    # use menu
    await ui_test.menu_click("Edit/Select/Select by Kind/Subcomponent")
    await ui_test.human_delay()

    # verify world leafs
    selection = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertEqual(selection, ['/World/Cone', '/World/Scope/Capsule'])


async def edit_test_func_edit_select_select_by_kind_component(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

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

    # use menu
    await ui_test.menu_click("Edit/Select/Select by Kind/Component")
    await ui_test.human_delay()

    # verify world leafs
    selection = omni.usd.get_context().get_selection().get_selected_prim_paths()
    # The /OmniverseKit_* prims are added by Kit
    tester.assertEqual(selection, ['/World/Cube', '/World/Disk', '/OmniverseKit_Persp', '/OmniverseKit_Front', '/OmniverseKit_Top', '/OmniverseKit_Right'])


async def edit_test_func_edit_select_select_by_kind_group(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

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

    # use menu
    await ui_test.menu_click("Edit/Select/Select by Kind/Group")
    await ui_test.human_delay()

    # verify world leafs
    selection = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertEqual(selection, ['/World/Cylinder', '/World/Plane', '/World/Sphere'])


async def edit_test_func_edit_select_select_by_kind_assembly(tester, menu_item: MenuItemDescription):
    from pathlib import Path

    # get file path
    extension_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(__name__)
    await omni.usd.get_context().open_stage_async(str(Path(extension_path).joinpath("data").joinpath("tests").joinpath("stages").joinpath("leaf_test.usda")))

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

    # use menu
    await ui_test.menu_click("Edit/Select/Select by Kind/Assembly")
    await ui_test.human_delay()

    # verify world leafs
    selection = omni.usd.get_context().get_selection().get_selected_prim_paths()
    tester.assertEqual(selection, ['/World/Plane'])

