import os
import carb
from pkgutil import iter_modules
import omni.usd
import omni.kit.app

from typing import List
from pxr import Usd
from functools import lru_cache
from pathlib import Path
from omni.kit import ui_test
from omni.kit.ui_test import Vec2
from omni.kit.ui_test import WidgetRef


@lru_cache()
def get_test_data_path(module: str, subpath: str = "") -> str:
    if not subpath:
        ext_path = omni.kit.app.get_app().get_extension_manager().get_extension_path_by_module(module)
        return str(Path(ext_path) / "data" / "tests")
    return str(Path(get_test_data_path(module)) / subpath)


async def wait():
    await omni.kit.app.get_app().next_update_async()
    await omni.kit.app.get_app().next_update_async()


async def wait_stage_loading(usd_context=omni.usd.get_context()):
    while True:
        _, files_loaded, total_files = usd_context.get_stage_loading_status()
        if files_loaded or total_files:
            await omni.kit.app.get_app().next_update_async()
            continue
        break

    await wait()


async def open_stage(path: str, usd_context=omni.usd.get_context()):
    await usd_context.open_stage_async(path)
    await wait_stage_loading(usd_context)


async def select_prims(paths, usd_context=omni.usd.get_context()):
    usd_context.get_selection().set_selected_prim_paths(paths, True)
    await wait()
    # wait for any marterials to load
    await wait_stage_loading()


def get_prims(stage, exclude_list=[]):
    prims = []
    for p in stage.Traverse(Usd.TraverseInstanceProxies(Usd.PrimIsActive and Usd.PrimIsDefined and Usd.PrimIsLoaded)):
        if p not in exclude_list:
            prims.append(p)
    return prims


async def wait_for_window(window_name: str):
    MAX_WAIT = 100

    # Find active window
    for _ in range(MAX_WAIT):
        window_root = ui_test.find(f"{window_name}")
        if window_root and window_root.widget.visible:
            await ui_test.human_delay()
            break
        await ui_test.human_delay(1)

    if not window_root:
        raise Exception("Can't find window {window_name}, wait time exceeded.")


async def handle_assign_material_dialog(index, strength_index=0):
    carb.log_warn("WARNING: 'handle_assign_material_dialog' is being DEPRECATED. Please use the function of the same name from 'omni.kit.material.library.test_helper.MaterialLibraryTestHelper'")

    from pxr import Sdf

    # handle assign dialog
    prims = omni.usd.get_context().get_selection().get_selected_prim_paths()
    if len(prims) == 1:
        shape = Sdf.Path(prims[0]).name
        window_name = f"Bind material to {shape}###context_menu_bind"
    else:
        window_name = f"Bind material to {len(prims)} selected models###context_menu_bind"

    await wait_for_window(window_name)

    # open listbox
    widget = ui_test.find(f"{window_name}//Frame/**/Button[*].identifier=='combo_open_button'")
    await ui_test.emulate_mouse_move_and_click(widget.center, human_delay_speed=4)
    # select material item on listbox
    await wait_for_window("MaterialPropertyPopupWindow")
    widget = ui_test.find(f"MaterialPropertyPopupWindow//Frame/**/TreeView[*]")
    # FIXME - can't use widget.click as open combobox has no readable size and clicks goto stage window
    item_name = widget.model.get_item_children(None)[index].name_model.as_string if index else "None"
    await ui_test.find(f"MaterialPropertyPopupWindow//Frame/**/Label[*].text=='{item_name}'").click(human_delay_speed=4)

    # select strength item on listbox
    widget = ui_test.find(f"{window_name}//Frame/**/ComboBox[*]")
    if widget:
        widget.model.set_value(strength_index)

    # click ok
    widget = ui_test.find(f"{window_name}//Frame/**/Button[*].identifier=='assign_material_ok_button'")
    await ui_test.emulate_mouse_move_and_click(widget.center, human_delay_speed=4)

    # wait for materials to load
    await ui_test.human_delay()
    await wait_stage_loading()


async def handle_create_material_dialog(mdl_path: str, mtl_name: str):
    carb.log_warn("WARNING: 'handle_create_material_dialog' is being DEPRECATED. Please use the function of the same name from 'omni.kit.material.library.test_helper.MaterialLibraryTestHelper'")

    subid_list = []
    def have_subids(id_list):
        nonlocal subid_list
        subid_list = id_list
    await omni.kit.material.library.get_subidentifier_from_mdl(mdl_file=mdl_path, on_complete_fn=have_subids)
    if len(subid_list)> 1:
        # material has subid and dialog is shown
        await wait_for_window("Create Material")
        create_widget = ui_test.find("Create Material//Frame/**/Button[*].identifier=='create_material_ok_button'")
        subid_widget = ui_test.find("Create Material//Frame/**/ComboBox[*].identifier=='create_material_subid_combo'")
        subid_list = subid_widget.model.get_item_list()
        subid_index = 0
        for index, subid in enumerate(subid_list):
            if subid.name == mtl_name:
                subid_index = index
        subid_widget.model.set_current_index(subid_index)
        await ui_test.human_delay()
        create_widget.widget.call_clicked_fn()
        await ui_test.human_delay(4)
        await wait_stage_loading()


async def delete_prim_path_children(prim_path: str):
    stage =  omni.usd.get_context().get_stage()
    root_prim = stage.GetPrimAtPath(prim_path)
    purge_list = []
    for prim in Usd.PrimRange(root_prim):
        if prim.GetPath().pathString != root_prim.GetPath().pathString:
            purge_list.append(prim)
    for prim in purge_list:
        stage.RemovePrim(prim.GetPath())

    # wait for refresh after deleting prims
    await ui_test.human_delay()

async def build_sdf_asset_frame_dictonary():
    widget_table = {}
    for frame in ui_test.find_all("Property//Frame/**/CollapsableFrame[*]"):
        if frame.widget.title != "Raw USD Properties":
            for widget in frame.find_all("Property//Frame/**/StringField[*].identifier!=''"):
                if widget.widget.identifier.startswith('sdf_asset_'):
                    if not frame.widget.title in widget_table:
                        widget_table[frame.widget.title] = {}
                    if not widget.widget.identifier in widget_table[frame.widget.title]:
                        widget_table[frame.widget.title][widget.widget.identifier] = 0
                    widget_table[frame.widget.title][widget.widget.identifier] += 1

    return widget_table


def push_window_height(cls, window_name, new_height=None):
    window = ui_test.find(window_name)
    if window:
        if not hasattr(cls, "_original_window_height"):
            cls._original_window_height = {}
        cls._original_window_height[window_name] = window.widget.height
        if new_height is not None:
            window.widget.height = new_height


def pop_window_height(cls, window_name):
    window = ui_test.find(window_name)
    if window:
        if hasattr(cls, "_original_window_height"):
            window.widget.height = cls._original_window_height[window_name]
            del cls._original_window_height[window_name]


async def handle_multiple_descendents_dialog(stage, prim_path: str, target_prim: str):
    root_prim = stage.GetPrimAtPath(prim_path)
    if not root_prim:
        return

    descendents = omni.usd.get_prim_descendents(root_prim)

    # skip if only root_prim
    if descendents == [root_prim]:
        return

    await ui_test.human_delay(10)
    await wait_for_window("Target prim has multiple descendents")
    await ui_test.human_delay(10)

    # need to select target_prim in combo_widget
    combo_widget = ui_test.find("Target prim has multiple descendents//Frame/**/ComboBox[*].identifier=='multi_descendents_combo'")
    combo_list = combo_widget.model.get_item_children(None)
    combo_index = 0
    for index, item in enumerate(combo_list):
        if item.prim.GetPrimPath().pathString == target_prim:
            combo_index = index

    combo_widget.model.set_current_index(combo_index)
    await ui_test.human_delay()

    ok_widget = ui_test.find("Target prim has multiple descendents//Frame/**/Button[*].identifier=='multi_descendents_ok_button'")
    await ok_widget.click()

    await wait_stage_loading()


async def arrange_windows(topleft_window="Stage", topleft_height=421.0, topleft_width=436.0, hide_viewport=False):
    from omni.kit.viewport.utility import get_active_viewport_window
    viewport_window = get_active_viewport_window()
    # omni.ui & legacy viewport synch
    await wait()
    if viewport_window:
        vp_width = int(1436 - topleft_width)
        viewport_window.position_x = 0
        viewport_window.position_y = 0
        viewport_window.width = vp_width
        viewport_window.height = 425
        viewport_window.visible = (not hide_viewport)
        # # XXX: force the legacy API
        # if hasattr(viewport_window, 'legacy_window'):
        #     legacy_window = viewport_window.legacy_window
        #     legacy_window.set_window_pos(0, 0)
        #     legacy_window.set_window_size(vp_width, 425)
        #     viewport.show_hide_window(not hide_viewport)

    import omni.ui as ui
    content_window = ui.Workspace.get_window("Content")
    if content_window:
        content_window.position_x = 0.0
        content_window.position_y = 448.0
        content_window.width = 1436.0 - topleft_width
        content_window.height = 421.0
        await ui_test.human_delay()

    stage_window = ui.Workspace.get_window("Stage")
    if stage_window:
        stage_window.position_x = 1436.0 - topleft_width
        stage_window.position_y = 0.0
        stage_window.width = topleft_width
        stage_window.height = topleft_height
        await ui_test.human_delay()

    layer_window = ui.Workspace.get_window("Layer")
    if layer_window:
        layer_window.position_x = 1436.0 - topleft_width
        layer_window.position_y = 0.0
        layer_window.width = topleft_width
        layer_window.height = topleft_height
        await ui_test.human_delay()

    tl_window = ui.Workspace.get_window(topleft_window)
    if tl_window:
        tl_window.focus()

    property_window = ui.Workspace.get_window("Property")
    if property_window:
        property_window.position_x = 1436.0 - topleft_width
        property_window.position_y = topleft_height + 27.0
        property_window.width = topleft_width
        property_window.height = 846.0 - topleft_height
        await ui_test.human_delay()

    # Wait for the layout to complete
    await wait()
    return viewport_window
