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.viewport.utility.camera_state import ViewportCameraState
from pathlib import Path

async def edit_test_func_edit_focus(tester, menu_item: MenuItemDescription):
    # 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("focus_prims.usda")))

    # test focus
    stage =  omni.usd.get_context().get_stage()
    prim_pos = ViewportCameraState().position_world
    for prim_path in ['/World/Cone', '/World/Cube', '/World/Sphere', '/World/defaultLight']:
        # select prim
        omni.usd.get_context().get_selection().set_selected_prim_paths([prim_path], True)

        # use menu
        await ui_test.menu_click("Edit/Focus")
        await ui_test.human_delay(50)

        # verify camera has moved
        pos = ViewportCameraState().position_world
        tester.assertTrue(pos != prim_pos)
        prim_pos = pos


async def edit_test_hotkey_func_edit_focus(tester, menu_item: MenuItemDescription):
    # 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("focus_prims.usda")))

    # test focus
    stage =  omni.usd.get_context().get_stage()
    prim_pos = ViewportCameraState().position_world
    for prim_path in ['/World/Cone', '/World/Cube', '/World/Sphere', '/World/defaultLight']:
        # select prim
        omni.usd.get_context().get_selection().set_selected_prim_paths([prim_path], True)

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

        # verify camera has moved
        pos = ViewportCameraState().position_world
        tester.assertTrue(pos != prim_pos)
        prim_pos = pos


async def edit_test_func_edit_toggle_visualization_mode(tester, menu_item: MenuItemDescription):
    # this test is odd due to the way viewport changed grid/display state
    # it is only tested on legacy Viewport as the test is against implememtation details / settings
    from omni.kit.viewport.utility import get_active_viewport
    if not hasattr(get_active_viewport(), "legacy_window"):
        return

    settings = carb.settings.get_settings()
    def get_state():
        return (settings.get("/app/viewport/grid/enabled"),
                settings.get("/persistent/app/viewport/displayOptions"))

    async def run_toggle():
        # use menu item
        await ui_test.menu_click("Edit/Toggle Visualization Mode")
        await ui_test.human_delay(1)

        return get_state()

    # Grab state to restore to
    grid_restore, display_restore = get_state()


    try:
        # Setup state to test against
        grid_state_a, display_state_a = True, 32767
        settings.set("/app/viewport/grid/enabled", grid_state_a)
        settings.set("/persistent/app/viewport/displayOptions", display_state_a)

        # run the menu item
        grid_state_b, display_state_b = await run_toggle()

        # values should not be equal fro mthe toggle
        tester.assertNotEqual(grid_state_a, grid_state_b)
        tester.assertNotEqual(display_state_a, display_state_b)

        # run the menu item
        grid_state_b, display_state_b = await run_toggle()

        # verify state changed back to original
        tester.assertEqual(grid_state_a, grid_state_b)
        tester.assertEqual(display_state_a, display_state_b)
    finally:
        # Restore to incoming state for any other tests
        settings.set("/app/viewport/grid/enabled", grid_restore)
        settings.set("/persistent/app/viewport/displayOptions", display_restore)

async def edit_test_func_edit_capture_screenshot(tester, menu_item: MenuItemDescription):
    warning_count = 0
    def on_log_event(e):
        if e.payload["level"] == carb.logging.LEVEL_WARN:
            nonlocal warning_count
            warning_count += 1

    # log event, count number of warnings outputed - FIXME: can I parse log and get messages?
    log_stream = omni.kit.app.get_app().get_log_event_stream()
    log_sub = log_stream.create_subscription_to_pop(on_log_event, name="edit_test_func_edit_capture_screenshot log event")

    # use menu
    await ui_test.menu_click("Edit/Capture Screenshot")

    # wait for screenshot to finish
    await ui_test.human_delay(100)

    log_stream = None
    log_sub = None

    tester.assertTrue(warning_count == 1)


async def edit_test_hotkey_func_edit_capture_screenshot(tester, menu_item: MenuItemDescription):
    warning_count = 0
    def on_log_event(e):
        if e.payload["level"] == carb.logging.LEVEL_WARN:
            nonlocal warning_count
            warning_count += 1

    # log event, count number of warnings outputed - FIXME: can I parse log and get messages?
    log_stream = omni.kit.app.get_app().get_log_event_stream()
    log_sub = log_stream.create_subscription_to_pop(on_log_event, name="edit_test_func_edit_capture_screenshot log event")

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

    # wait for screenshot to finish
    await ui_test.human_delay(100)

    log_stream = None
    log_sub = None

    tester.assertTrue(warning_count == 1)


async def _edit_test_hotkey_func_edit_toggle_some_visibility(tester, menu_item: MenuItemDescription, setting_key: str):
    settings = carb.settings.get_settings()
    initial_state = settings.get(setting_key) or False

    await ui_test.emulate_keyboard_press(menu_item.hotkey[1])
    await ui_test.human_delay(2)
    tester.assertNotEqual(settings.get(setting_key), initial_state)

    await ui_test.emulate_keyboard_press(menu_item.hotkey[1])
    await ui_test.human_delay(2)
    tester.assertEqual(settings.get(setting_key), initial_state)
