import os
import shutil
import tempfile
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
from omni.kit.window.file_importer.test_helper import FileImporterTestHelper
from omni.kit.window.file_exporter.test_helper import FileExporterTestHelper


async def file_test_func_file_new(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").click()

    await tester.wait_for_stage_event()

    # verify its a new layer
    tester.assertFalse(omni.usd.get_context().get_stage().GetRootLayer().identifier == layer_name)


async def file_test_func_file_open(tester, menu_item: MenuItemDescription):
    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("Open").click()
    await ui_test.human_delay(5)

    async with FileImporterTestHelper() as file_import_helper:
        dir_url = carb.tokens.acquire_tokens_interface().resolve("${kit}/exts/omni.kit.menu.file/data/tests")
        await file_import_helper.click_apply_async(filename_url=f"{dir_url}/4Lights.usda")
        await tester.wait_for_stage_event()

    # verify name of stage
    stage = omni.usd.get_context().get_stage()
    tester.assertTrue(stage.GetRootLayer().identifier.replace("\\", "/").endswith("data/tests/4Lights.usda"))


async def file_test_func_file_reopen(tester, menu_item: MenuItemDescription):
    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("Open Recent").click()
    await menu_widget.find_menu(tester._light_path, ignore_case=True).click()
    await tester.wait_for_stage_event()

    # verify its a unmodified stage
    tester.assertFalse(omni.kit.undo.can_undo())

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

    # verify its a modified stage
    tester.assertTrue(omni.kit.undo.can_undo())

    await tester.reset_stage_event(omni.usd.StageEventType.OPENED)
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Open Recent").click()
    await menu_widget.find_menu("Current stage:", contains_path=True).click()
    await tester.wait_for_stage_event()

    # verify its a unmodified stage
    tester.assertFalse(omni.kit.undo.can_undo())


async def file_test_func_file_re_open_with_new_edit_layer(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await file_test_func_file_open(tester, {})

    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("Re-open with New Edit Layer").click()
    await ui_test.human_delay(50)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights_layer.usda")
        await tester.wait_for_stage_event()

    # verify its not saved to tmpdir
    layer = omni.usd.get_context().get_stage().GetRootLayer()
    expected_name = f"{tmpdir}/4Lights_layer.usd".replace("\\", "/")

    tester.assertTrue(layer.anonymous)
    tester.assertTrue(len(layer.subLayerPaths) == 2)
    tester.assertTrue(layer.subLayerPaths[0] == expected_name)
    tester.assertTrue(layer.subLayerPaths[1].endswith("exts/omni.kit.menu.file/data/tests/4Lights.usda"))

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_func_file_save(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Save").click()
    await ui_test.human_delay(50)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usd")
        await tester.wait_for_stage_event()

    # verify its saved to tmpdir
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    expected_size = os.path.getsize(expected_name)
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)

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

    # save again, it should not show dialog....
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Save").click()
    await ui_test.human_delay()

    # handle dialog....
    save_widget = ui_test.find("Select Files to Save##file.py//Frame/**/Button[*].text=='Save Selected'")
    await save_widget.click()
    await ui_test.human_delay()
    await tester.wait_for_stage_event()

    # verify file was saved
    tester.assertFalse(os.path.getsize(expected_name) == expected_size)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_func_file_save_with_options(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Save With Options").click()
    await ui_test.human_delay(5)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usda")
        await tester.wait_for_stage_event()

    # verify its saved to tmpdir
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    expected_size = os.path.getsize(expected_name)
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)

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

    # save again, it should not show dialog....
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Save").click()
    await ui_test.human_delay()

    # handle dialog....
    save_widget = ui_test.find("Select Files to Save##file.py//Frame/**/Button[*].text=='Save Selected'")
    await save_widget.click()
    await ui_test.human_delay()
    await tester.wait_for_stage_event()

    # verify file was saved
    tester.assertFalse(os.path.getsize(expected_name) == expected_size)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_func_file_save_as(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await file_test_func_file_open(tester, {})

    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Save As...").click()
    await ui_test.human_delay(5)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usda")
        await tester.wait_for_stage_event()

    # verify its saved to tmpdir
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_func_file_save_flattened_as(tester, menu_item: MenuItemDescription):
    # save flattened doesn't send any stage events as it exports not saves
    tmpdir = tempfile.mkdtemp()
    await file_test_func_file_open(tester, {})

    stage =  omni.usd.get_context().get_stage()
    layer_name =  stage.GetRootLayer().identifier.replace("\\", "/") if stage else "None"

    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Save Flattened As...").click()
    await ui_test.human_delay(5)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usda")
        await ui_test.human_delay(20)

    # verify its exported to tmpdir & hasn't changed stage path
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == layer_name)
    tester.assertFalse(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_func_file_open_recent_exts_omni_kit_menu_file_data_tests_4lights_usda(tester, menu_item: MenuItemDescription):
    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("Open Recent").click()

    await menu_widget.find_menu(tester._light_path, ignore_case=True).click()
    await tester.wait_for_stage_event()

    # verify its 4lights
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/").endswith("exts/omni.kit.menu.file/data/tests/4Lights.usda"))


async def file_test_func_file_exit(tester, menu_item: MenuItemDescription):
    # not sure how to test this without kit exiting...
    pass


async def file_test_func_file_exit_fast__experimental(tester, menu_item: MenuItemDescription):
    # not sure how to test this without kit exiting...
    pass


async def file_test_hotkey_func_file_new(tester, menu_item: MenuItemDescription):
    stage =  omni.usd.get_context().get_stage()
    layer_name =  stage.GetRootLayer().identifier if stage else "None"

    # emulate hotkey
    await tester.reset_stage_event(omni.usd.StageEventType.OPENED)
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])
    await tester.wait_for_stage_event()

    # verify its a new layer
    tester.assertFalse(omni.usd.get_context().get_stage().GetRootLayer().identifier == layer_name)


async def file_test_hotkey_func_file_open(tester, menu_item: MenuItemDescription):
    # emulate hotkey
    await tester.reset_stage_event(omni.usd.StageEventType.OPENED)
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])
    await ui_test.human_delay(5)

    async with FileImporterTestHelper() as file_import_helper:
        dir_url = carb.tokens.acquire_tokens_interface().resolve("${kit}/exts/omni.kit.menu.file/data/tests")
        await file_import_helper.click_apply_async(filename_url=f"{dir_url}/4Lights.usda")
        await tester.wait_for_stage_event()

    # verify its 4lights
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/").endswith("exts/omni.kit.menu.file/data/tests/4Lights.usda"))


async def file_test_hotkey_func_file_save(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    # emulate hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])
    await ui_test.human_delay(5)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usd")
        await tester.wait_for_stage_event()

    # verify its saved to tmpdir
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)

    expected_size = os.path.getsize(expected_name)

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

    # save again, it should not show dialog....
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    await ui_test.emulate_keyboard_press(carb.input.KeyboardInput.S, carb.input.KEYBOARD_MODIFIER_FLAG_CONTROL)

    # handle dialog....
    save_widget = ui_test.find("Select Files to Save##file.py//Frame/**/Button[*].text=='Save Selected'")
    if save_widget:
        await save_widget.click()
        await ui_test.human_delay()
    await tester.wait_for_stage_event()

    # verify file was saved
    tester.assertFalse(os.path.getsize(expected_name) == expected_size)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_hotkey_func_file_save_with_options(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    # emulate hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])
    await ui_test.human_delay(5)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usd")
        await tester.wait_for_stage_event()

    # verify its saved to tmpdir
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    expected_size = os.path.getsize(expected_name)

    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)


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

    # save again, it should not show dialog....
    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    # emulate hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])

    # handle dialog....
    save_widget = ui_test.find("Select Files to Save##file.py//Frame/**/Button[*].text=='Save Selected'")
    await save_widget.click()
    await ui_test.human_delay()
    await tester.wait_for_stage_event()

    # verify file was saved
    tester.assertFalse(os.path.getsize(expected_name) == expected_size)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_hotkey_func_file_save_as(tester, menu_item: MenuItemDescription):
    tmpdir = tempfile.mkdtemp()
    await file_test_func_file_open(tester, {})

    await tester.reset_stage_event(omni.usd.StageEventType.SAVED)
    # emulate hotkey
    await ui_test.emulate_keyboard_press(menu_item.hotkey[1], menu_item.hotkey[0])
    await ui_test.human_delay(5)

    async with FileExporterTestHelper() as file_export_helper:
        await file_export_helper.click_apply_async(filename_url=f"{tmpdir}/4Lights.usda")
        await tester.wait_for_stage_event()

    # verify its saved to tmpdir
    expected_name = f"{tmpdir}/4Lights.usd".replace("\\", "/")
    tester.assertTrue(omni.usd.get_context().get_stage().GetRootLayer().identifier.replace("\\", "/") == expected_name)

    # cleanup
    shutil.rmtree(tmpdir, ignore_errors=True)


async def file_test_func_file_share(tester, menu_item: MenuItemDescription):
    # File->Share will only work on omniverse:// files. Currently there is no way to mock this scenario.
    menu_widget = ui_test.get_menubar()
    await menu_widget.find_menu("File").click()
    await menu_widget.find_menu("Share").click()
    pass
