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_undo(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()
    menu_widget = ui_test.get_menubar()

    # 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 prim created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere'])

    # click undo
    await menu_widget.find_menu("Edit").click()
    await menu_widget.find_menu("Undo").click()
    await ui_test.human_delay()

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


async def edit_test_hotkey_func_edit_undo(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 prim created
    prim_list = tester.get_stage_prims()
    tester.assertTrue(prim_list == ['/Sphere'])

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

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


async def edit_test_func_edit_redo(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()
    menu_widget = ui_test.get_menubar()

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

    # do/undo so redo will work
    await edit_test_func_edit_undo(tester, None)

    # click redo
    await menu_widget.find_menu("Edit").click()
    await menu_widget.find_menu("Redo").click()
    await ui_test.human_delay()

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


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

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

    # do/undo so redo will work
    await edit_test_func_edit_undo(tester, None)

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

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


async def edit_test_func_edit_repeat(tester, menu_item: MenuItemDescription):
    menu_widget = ui_test.get_menubar()

    # 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 prim created
    prim_list = tester.get_stage_prims()
    tester.assertListEqual(prim_list, ['/Sphere'])

    # click repeat
    await menu_widget.find_menu("Edit").click()
    await menu_widget.find_menu("Repeat").click()
    await ui_test.human_delay()

    # verify prim was repeated
    prim_list = tester.get_stage_prims()
    tester.assertListEqual(prim_list, ['/Sphere', '/Sphere_01'])


async def edit_test_hotkey_func_edit_repeat(tester, menu_item: MenuItemDescription):
    # 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 prim created
    prim_list = tester.get_stage_prims()
    tester.assertListEqual(prim_list, ['/Sphere'])

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

    # verify prim was repeated
    prim_list = tester.get_stage_prims()
    tester.assertListEqual(prim_list, ['/Sphere', '/Sphere_01'])
