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 file_test_func_file_new_from_stage_template_empty(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()
    layer_name =  stage.GetRootLayer().identifier if stage else "None"
    await tester.reset_stage_event(omni.usd.StageEventType.OPENED)

    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("New From Stage Template").click()

    # select empty and wait for stage open
    await menu_widget.find_menu("Empty").click()
    await tester.wait_for_stage_event()

    # verify Empty stage
    stage =  omni.usd.get_context().get_stage()
    tester.assertFalse(stage.GetRootLayer().identifier == layer_name)
    prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll()]
    tester.assertTrue(prim_list == ['/World'])


async def file_test_func_file_new_from_stage_template_sunlight(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()
    layer_name =  stage.GetRootLayer().identifier if stage else "None"
    await tester.reset_stage_event(omni.usd.StageEventType.OPENED)

    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("New From Stage Template").click()

    # select Sunlight and wait for stage open
    await menu_widget.find_menu("Sunlight").click()
    await tester.wait_for_stage_event()

    # verify Sunlight stage
    stage =  omni.usd.get_context().get_stage()
    tester.assertFalse(stage.GetRootLayer().identifier == layer_name)
    prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll()]
    tester.assertTrue(prim_list == ['/World', '/Environment', '/Environment/defaultLight'])

async def file_test_func_file_new_from_stage_template_default_stage(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()
    layer_name =  stage.GetRootLayer().identifier if stage else "None"
    await tester.reset_stage_event(omni.usd.StageEventType.OPENED)

    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("New From Stage Template").click()

    # select Default Stage and wait for stage open
    await menu_widget.find_menu("Default Stage").click()
    await tester.wait_for_stage_event()

    # verify Default Stage
    stage =  omni.usd.get_context().get_stage()
    tester.assertFalse(stage.GetRootLayer().identifier == layer_name)
    prim_list = [prim.GetPath().pathString for prim in stage.TraverseAll()]
    tester.assertEqual(set(prim_list), set([
        '/World', '/Environment', '/Environment/Sky',
        '/Environment/DistantLight', '/Environment/Looks',
        '/Environment/Looks/Grid', '/Environment/Looks/Grid/Shader',
        '/Environment/ground', '/Environment/groundCollider']))
