import os
import carb
import weakref
import omni.ext
import omni.kit.ui
import omni.usd
from typing import Callable
from pathlib import Path
from pxr import Usd, UsdGeom, Sdf, UsdShade
from omni.kit.window.filepicker import FilePickerDialog
from omni.kit.widget.filebrowser import FileBrowserItem

TEST_DATA_PATH = ""
_extension_instance = None


class MaterialSwapExtension(omni.ext.IExt):
    def on_startup(self, ext_id):
        self._usd_context = omni.usd.get_context()
        self._selection = self._usd_context.get_selection()
        self._build_ui()

        manager = omni.kit.app.get_app().get_extension_manager()
        extension_path = manager.get_extension_path(ext_id)
        global TEST_DATA_PATH
        TEST_DATA_PATH = Path(extension_path).joinpath("data").joinpath("tests")

        global _extension_instance
        _extension_instance = self

    def _build_ui(self):

        self._window = omni.kit.ui.Window("Material Swap", 300, 200, menu_path="Window/Material Swap")

        mat_settings = omni.kit.ui.CollapsingFrame("Material Swap Parameters", True)

        field_width = 120

        self._default_material_name = "None Selected"
        self._old_mat = omni.kit.ui.TextBox(self._default_material_name)
        self._old_mat_label = omni.kit.ui.Label("Material to Replace")
        self._old_mat_selected = omni.kit.ui.Button("Use Selected")
        self._old_mat_selected.set_clicked_fn(self._get_selectedold_fn)

        self._old_mat.width = field_width * 3
        self._old_mat_label.width = field_width
        self._old_mat_selected.width = field_width

        old_row = omni.kit.ui.RowLayout()
        old_row.add_child(self._old_mat_label)
        old_row.add_child(self._old_mat)
        old_row.add_child(self._old_mat_selected)

        self._new_mat = omni.kit.ui.TextBox(self._default_material_name)
        self._new_mat_label = omni.kit.ui.Label("Material to Use")
        self._new_mat_browse = omni.kit.ui.Button("Browse")
        self._new_mat_browse.set_clicked_fn(self._on_browse_fn)
        self._new_mat_selected = omni.kit.ui.Button("Use Selected")
        self._new_mat_selected.set_clicked_fn(self._get_selectednew_fn)

        def on_filter_item(dialog: FilePickerDialog, item: FileBrowserItem) -> bool:
            if item and not item.is_folder and dialog.current_filter_option == 0:
                _, ext = os.path.splitext(item.path)
                return ext in [".mdl"]
            return True

        # create filepicker
        dialog = FilePickerDialog(
            "Open File",
            apply_button_label="Open",
            click_apply_handler=lambda filename, dirname, o=weakref.ref(self): MaterialSwapExtension._on_file_open(o, dialog, filename, dirname),
            item_filter_options=["MDL Files (*.mdl)", "All Files (*)"],
            item_filter_fn=lambda item: on_filter_item(dialog, item),
        )
        dialog.hide()
        self._new_mat_fp = dialog

        self._new_mat.width = field_width * 3
        self._new_mat_label.width = field_width
        self._new_mat_browse.width = field_width
        self._new_mat_selected.width = field_width

        new_row = omni.kit.ui.RowLayout()
        new_row.add_child(self._new_mat_label)
        new_row.add_child(self._new_mat)
        new_row.add_child(self._new_mat_selected)
        new_row.add_child(self._new_mat_browse)

        mat_settings.add_child(old_row)
        mat_settings.add_child(new_row)

        # buttons
        actions = omni.kit.ui.CollapsingFrame("Actions", True)
        self._ui_swap_button = omni.kit.ui.Button("Swap")
        self._ui_swap_button.set_clicked_fn(self._on_swap_fn)
        self._ui_swap_button.width = omni.kit.ui.Percent(30)

        actions.add_child(self._ui_swap_button)

        self._window.layout.add_child(mat_settings)
        self._window.layout.add_child(actions)

    def _on_swap_fn(self, widget):
        if self._old_mat.value == self._default_material_name:
            carb.log_error('"Material To Replace" not set')
            return

        if self._new_mat.value == self._default_material_name:
            carb.log_error('"Material to Use" not set')
            return

        stage = self._usd_context.get_stage()
        for prim in stage.Traverse():
            if not prim.GetMetadata("hide_in_stage_window") and omni.usd.is_prim_material_supported(prim):
                mat, rel = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
                mat_path = mat.GetPath()
                if mat_path == self._old_mat.value:
                    path = prim.GetPath()
                    omni.kit.commands.execute(
                        "BindMaterial",
                        prim_path=path,
                        material_path=self._new_mat.value,
                        strength=UsdShade.Tokens.weakerThanDescendants,
                    )

    def _on_browse_fn(self, widget):
        self._new_mat_fp.show()

    @staticmethod
    def _on_file_open(owner_weak, dialog: FilePickerDialog, filename: str, dirname: str):
        owner = owner_weak()
        if not owner:
            return
        file_path = ""
        if dirname:
            file_path = f"{dirname}/{filename}"
        elif filename:
            file_path = filename
        dialog.hide()

        if not "omniverse:" in file_path:
            carb.log_warn("No local file support, please select an Omniverse asset")
            owner._new_mat.value = "No local file support, please select an Omniverse asset"
        else:
            owner._import_file_path = file_path
            base_name = os.path.basename(owner._import_file_path)
            file_name, _ = os.path.splitext(os.path.basename(base_name))
            file_name = owner._make_valid_identifier(file_name)
            owner._add_material(file_path, file_name)

    def _add_material(self, mat_path, mat_name):
        stage = self._usd_context.get_stage()
        d_xform = stage.GetDefaultPrim()
        d_xname = str(d_xform.GetPath())
        mat_prim = stage.DefinePrim("{}/Looks/{}".format(d_xname, mat_name), "Material")
        mat = UsdShade.Material.Get(stage, mat_prim.GetPath())
        if mat:
            shader_prim = stage.DefinePrim("{}/Looks/{}/Shader".format(d_xname, mat_name), "Shader")
            shader = UsdShade.Shader.Get(stage, shader_prim.GetPath())
            if shader:
                mdl_token = "mdl"
                shader_out = shader.CreateOutput("out", Sdf.ValueTypeNames.Token)
                mat.CreateSurfaceOutput(mdl_token).ConnectToSource(shader_out)
                mat.CreateVolumeOutput(mdl_token).ConnectToSource(shader_out)
                mat.CreateDisplacementOutput(mdl_token).ConnectToSource(shader_out)
                shader.GetImplementationSourceAttr().Set(UsdShade.Tokens.sourceAsset)
                shader.SetSourceAsset(mat_path, mdl_token)
                shader.SetSourceAssetSubIdentifier(mat_name, mdl_token)

        self._new_mat.value = str(mat.GetPath())

    def _make_valid_identifier(self, identifier):
        if len(identifier) == 0:
            return "_"

        result = identifier[0]
        if not identifier[0].isalpha():
            result = "_"

        for i in range(len(identifier) - 1):
            if identifier[i + 1].isalnum():
                result += identifier[i + 1]
            else:
                result += "_"

        return result

    def _on_file_selected_fn(self, widget):
        carb.log_warn("file selected fn")

    def _get_selectedold_fn(self, widget):
        paths = self._selection.get_selected_prim_paths()
        base_path = paths[0] if len(paths) > 0 else None
        base_path = Sdf.Path(base_path) if base_path else None

        if base_path:
            mat = self._get_material(base_path)
            self._old_mat.value = str(mat)
        else:
            carb.log_error("No selected prim")

    def _get_selectednew_fn(self, widget):
        paths = self._selection.get_selected_prim_paths()
        base_path = paths[0] if len(paths) > 0 else None
        base_path = Sdf.Path(base_path) if base_path else None

        if base_path:
            mat = self._get_material(base_path)
            self._new_mat.value = str(mat)
        else:
            carb.log_error("No selected prim")

    def _get_material(self, path):
        stage = self._usd_context.get_stage()
        prim = stage.GetPrimAtPath(path)
        p_type = prim.GetTypeName()
        if p_type == "Material":
            return path
        else:
            return None

    def on_shutdown(self):
        if self._new_mat_fp:
            del self._new_mat_fp
        self._new_mat_fp = None

        global _extension_instance
        _extension_instance = None


def get_instance():
    global _extension_instance
    return _extension_instance
