# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

from omni.kit.window.property.templates import SimplePropertyWidget, LABEL_WIDTH, LABEL_HEIGHT
from omni.kit.material.library.listbox_widget import MaterialListBoxWidget
from omni.kit.material.library.thumbnail_loader import ThumbnailLoader
from omni.kit.material.library.search_widget import SearchWidget
from .material_utils import Constant, get_binding_from_prims
from .context_menu import show_context_menu
from typing import List, Set
from functools import partial
from pxr import Usd, Tf, Sdf, UsdShade

import copy
import os
import weakref
import asyncio
import weakref
import carb
import omni.ui as ui
import omni.usd
import omni.kit.material.library
import omni.kit.context_menu


class UsdBindingAttributeWidget(SimplePropertyWidget):
    def __init__(self, extension_path, add_context_menu=False):
        super().__init__(title="Materials on selected models", collapsed=False)
        self._strengths = {
            "Weaker than Descendants": UsdShade.Tokens.weakerThanDescendants,
            "Stronger than Descendants": UsdShade.Tokens.strongerThanDescendants,
        }
        self._extension_path = extension_path
        self._thumbnail_loader = ThumbnailLoader()

        BUTTON_BACKGROUND_COLOR = 0xFF323434
        GENERAL_BACKGROUND_COLOR = 0xFF1F2124
        self._theme = carb.settings.get_settings().get_as_string("/persistent/app/window/uiStyle") or "NvidiaDark"
        if self._theme != "NvidiaDark":
            BUTTON_BACKGROUND_COLOR = 0xFF545454
            GENERAL_BACKGROUND_COLOR = 0xFF545454

        self._listbox_widget = None
        self._search_widget = SearchWidget(theme=self._theme, icon_path=None)

        self._style = {
            "Image::material": {"margin": 12},
            "Button::material_mixed": {"margin": 10},
            "Button.Image::material_mixed": {"image_url": f"{self._extension_path}/data/icons/material_mixed@3x.png"},
            "Button::material_solo": {"margin": 10},
            "Button.Image::material_solo": {"image_url": f"{self._extension_path}/data/icons/material@3x.png"},
            "Field::prims" : {"font_size": Constant.FONT_SIZE},
            "Field::prims_mixed" : {"font_size": Constant.FONT_SIZE, "color": Constant.MIXED_COLOR},
            "Label::prim" : {"font_size": Constant.FONT_SIZE},
            "ComboBox": {"font_size": Constant.FONT_SIZE},
            "Label::combo_mixed": {
                "font_size": Constant.FONT_SIZE,
                "color": Constant.MIXED_COLOR,
                "margin_width": 5,
            },
         }

        self._listener = None
        self._pending_frame_rebuild_task = None

        self._paste_all_menu = None
        if add_context_menu:
            # paste all on collapsible frame header
            from .context_menu import is_bindable_prim_selected, is_material_copied, is_paste_all, paste_material_all

            menu_dict = {"name": "Paste To All", "enabled_fn": [is_bindable_prim_selected, is_material_copied, is_paste_all], "onclick_fn": paste_material_all}
            self._paste_all_menu = omni.kit.context_menu.add_menu(menu_dict, "group_context_menu.Materials on selected models", "omni.kit.window.property")

    def clean(self):
        self._thumbnail_loader = None
        if self._listbox_widget:
            self._listbox_widget.clean()
        self._listbox_widget = None
        self._search_widget.clean()
        self._paste_all_menu = None
        super().clean()

    def _materials_changed(self):
        self.request_rebuild()

    def _get_prim(self, prim_path):
        if prim_path:
            stage = self._payload.get_stage()
            if stage:
                return stage.GetPrimAtPath(prim_path)
        return None

    def on_new_payload(self, payload):
        """
        See PropertyWidget.on_new_payload
        """
        if not super().on_new_payload(payload):
            return False

        if not payload:
            return False

        if len(payload) == 0:
            return False

        for prim_path in payload:
            prim = self._get_prim(prim_path)
            if prim and not omni.usd.is_prim_material_supported(prim):
                return False
            if not prim and not Usd.CollectionAPI.IsCollectionAPIPath(prim_path):
                return False

        return True

    def _get_index(self, items: List[str], value: str, default_value: int = 0):
        index = default_value
        try:
            index = items.index(value)
        except:
            pass
        return index

    def _get_theme_name(self, darkname, lightname):
        if self._theme=="NvidiaDark":
            return darkname
        return lightname

    def _show_material_popup(self, name_field: ui.StringField, material_index: int, bind_material_fn: callable):
        if self._listbox_widget:
            self._listbox_widget.clean()
            del self._listbox_widget
            self._listbox_widget = None

        self._listbox_widget = MaterialListBoxWidget(icon_path=None, index=material_index, on_click_fn=bind_material_fn, theme=self._theme)
        self._listbox_widget.set_parent(name_field)
        self._listbox_widget.set_selection_on_loading_complete(name_field.model.get_value_as_string() if material_index >= 0 else None)
        self._listbox_widget.build_ui()

    def _build_material_popup(self, material_list: List[str], material_index: int, thumbnail_image: ui.Button, bound_prim_list: Set[Usd.Prim], bind_material_fn: callable, on_goto_fn: callable, on_dragdrop_fn: callable):
        def drop_accept(url):
            if len(url.split('\n')) > 1:
                carb.log_warn(f"build_material_popup multifile drag/drop not supported")
                return False

            if url.startswith("material::"):
                return True

            url_ext = os.path.splitext(url)[1]
            return url_ext.lower() in [".mdl"]

        def dropped_mtl(event, bind_material_fn):
            url = event.mime_data
            subid = None
            # remove omni.kit.browser.material "material::" prefix
            if url.startswith("material::"):
                url=url[10:]
                parts = url.split(".mdl@")
                if len(parts) == 2:
                    subid = parts[1]
                    url = url.replace(f".mdl@{subid}", ".mdl")

            bind_material_fn(url=url, subid=subid)

        name_field, listbox_button, goto_button = self._search_widget.build_ui_popup(search_size=LABEL_HEIGHT,
                                                                                     popup_text=material_list[material_index] if material_index >= 0 else "Mixed",
                                                                                     index=material_index,
                                                                                     update_fn=bind_material_fn)

        name_field.set_mouse_pressed_fn(lambda x, y, b, m, f=name_field: self._show_material_popup(f, material_index, bind_material_fn))
        name_field.set_accept_drop_fn(drop_accept)
        name_field.set_drop_fn(partial(dropped_mtl, bind_material_fn=on_dragdrop_fn))

        listbox_button.set_mouse_pressed_fn(lambda x, y, b, m, f=name_field: self._show_material_popup(f, material_index, bind_material_fn))

        if material_index > 0:
            goto_button.set_mouse_pressed_fn(lambda x, y, b, m, f=name_field: on_goto_fn(f.model))

        if self._collapsable_frame:
            def mouse_clicked(b, f):
                if b == 0:
                   on_goto_fn(f.model)
                elif b == 1:
                   show_context_menu(self._payload.get_stage(), f.model.get_value_as_string(), bound_prim_list)

            thumbnail_image.set_mouse_pressed_fn(lambda x, y, b, m, f=name_field: mouse_clicked(b, f))

        thumbnail_image.set_accept_drop_fn(drop_accept)
        thumbnail_image.set_drop_fn(partial(dropped_mtl, bind_material_fn=on_dragdrop_fn))

    def _build_combo(self, material_list: List[str], material_index: int, on_fn: callable):
        def set_visible(widget, visible):
            widget.visible = visible

        with ui.ZStack():
            combo = ui.ComboBox(material_index, *material_list)
            combo.model.add_item_changed_fn(on_fn)
            if material_index == -1:
                placeholder = ui.Label(
                    "Mixed",
                    height=LABEL_HEIGHT,
                    name="combo_mixed",
                )
                self._combo_subscription.append(
                    combo.model.get_item_value_model().subscribe_value_changed_fn(
                        lambda m, w=placeholder: set_visible(w, m.as_int < 0)
                    )
                )
            return combo

    def _build_binding_info(
        self,
        inherited: bool,
        style_name: str,
        material_list: List[str],
        material_path: str,
        strength_value: str,
        bound_prim_list: Set[str],
        on_material_fn: callable,
        on_strength_fn: callable,
        on_goto_fn: callable,
        on_dragdrop_fn: callable,
    ):
        prim_style = "prims"

        ## fixup Constant.SDF_PATH_INVALID vs "None"
        material_list = copy.copy(material_list)
        if material_list[0] == Constant.SDF_PATH_INVALID:
            material_list[0] = "None"
        if material_path == Constant.SDF_PATH_INVALID:
            material_path = "None"

        type_label = 'Prim'
        if len(bound_prim_list) == 1:

            if isinstance(bound_prim_list[0], Usd.Prim):
                bound_prims = bound_prim_list[0].GetPath().pathString
            elif isinstance(bound_prim_list[0], Usd.CollectionAPI):
                bound_prims = bound_prim_list[0].GetCollectionPath().pathString
                type_label = 'Collection'
            else:
                bound_prims = bound_prim_list[0]
        elif len(bound_prim_list) > 1:
            bound_prims = Constant.MIXED
            prim_style = "prims_mixed"
        else:
            bound_prims = "None"

        if not self._first_row:
            ui.Separator(height=10)
        self._first_row = False

        index = self._get_index(material_list, material_path, -1)

        with ui.HStack(style=self._style):
            with ui.ZStack(width=0):
                image_button = ui.Button(
                    "",
                    width=Constant.ICON_SIZE,
                    height=Constant.ICON_SIZE,
                    fill_policy=ui.FillPolicy.PRESERVE_ASPECT_FIT,
                    name=style_name,
                    identifier="preview_drop_target"
                )

                if index != -1:
                    material_prim = self._get_prim(material_path)
                    if material_prim:
                        def on_image_progress(button: ui.Button, image: ui.Image, progress: float):
                            if progress > 0.999:
                                button.image_url = image.source_url
                                image.visible = False

                        thumbnail_image = ui.Image(
                            width=Constant.ICON_SIZE,
                            height=Constant.ICON_SIZE,
                            fill_policy=ui.FillPolicy.PRESERVE_ASPECT_FIT,
                            name="material"
                        )
                        self._thumbnail_loader.load(material_prim, thumbnail_image)
                        # callback when image is loading
                        thumbnail_image.set_progress_changed_fn(lambda p, b=image_button, i=thumbnail_image: on_image_progress(b, i, p))

            with ui.VStack(spacing=10):
                with ui.HStack():
                    ui.Label(
                        type_label,
                        width=Constant.BOUND_LABEL_WIDTH,
                        height=LABEL_HEIGHT,
                        name="prim",
                    )
                    ui.StringField(name=prim_style, height=LABEL_HEIGHT, enabled=False).model.set_value(
                        bound_prims
                    )
                if index != -1 and inherited:
                    inherited_list = copy.copy(material_list)
                    inherited_list[index] = f"{inherited_list[index]} (inherited)"
                    self._build_material_popup(inherited_list, index, image_button, bound_prim_list, on_material_fn, on_goto_fn, on_dragdrop_fn)
                else:
                    self._build_material_popup(material_list, index, image_button, bound_prim_list, on_material_fn, on_goto_fn, on_dragdrop_fn)

                with ui.HStack():
                    ui.Label(
                        "Strength",
                        width=Constant.BOUND_LABEL_WIDTH,
                        height=LABEL_HEIGHT,
                        name="prim",
                    )
                    index = self._get_index(list(self._strengths.values()), strength_value, -1)
                    self._build_combo(list(self._strengths.keys()), index, on_strength_fn)

    def build_items(self):
        binding = get_binding_from_prims(self._payload.get_stage(), self._payload)
        if binding is None:
            return
        strengths = self._strengths
        material_list = [Constant.SDF_PATH_INVALID] + [mtl for mtl in binding["material"]]

        stage = self._payload.get_stage()
        if not stage or len(self._payload) == 0:
            return
        last_prim = self._get_prim(self._payload[-1])

        if not self._listener:
            self._listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, self._on_usd_changed, stage)

        def bind_material_to_prims(bind_material_path, items):
            prim_path_list = []
            strength_list = []
            for prim_path, material_name, strength in items:
                if str(material_name) != str(bind_material_path):
                    prim_path_list.append(prim_path)
                    strength_list.append(strength)


            if str(bind_material_path) == Constant.SDF_PATH_INVALID:
                bind_material_path = None

            omni.kit.commands.execute(
                "BindMaterial",
                material_path=bind_material_path,
                prim_path=prim_path_list,
                strength=strength_list,
            )

        def on_strength_changed(model, item, relationships):
            index = model.get_item_value_model().as_int
            if index < 0:
                carb.log_error(f"on_strength_changed with invalid index {index}")
                return
            bind_strength = list(strengths.values())[index]
            omni.kit.undo.begin_group()
            for relationship in relationships:
                omni.kit.commands.execute("SetMaterialStrength", rel=relationship, strength=bind_strength)
            omni.kit.undo.end_group()

            if self._collapsable_frame:
                self._collapsable_frame.rebuild()

        def on_material_changed(model, item, items):
            if isinstance(model, omni.ui._ui.AbstractItemModel):
                index = model.get_item_value_model().as_int
                if index < 0:
                    carb.log_error(f"on_material_changed with invalid index {index}")
                    return
                bind_material_path = material_list[index] if material_list[index] != Constant.SDF_PATH_INVALID else None
            elif isinstance(model, omni.ui._ui.SimpleStringModel):
                bind_material_path = model.get_value_as_string()
                if bind_material_path == "None":
                    bind_material_path = None
            else:
                carb.log_error(f"on_material_changed model {type(model)} unsupported")
                return
            bind_material_to_prims(bind_material_path, items)
            if self._collapsable_frame:
                self._collapsable_frame.rebuild()

        def on_dragdrop(url, items, subid=""):
            if len(url.split('\n')) > 1:
                carb.log_warn(f"build_binding_info multifile drag/drop not supported")
                return

            def on_create_func(mtl_prim):
                bind_material_to_prims(mtl_prim.GetPath(), items)
                if self._collapsable_frame:
                    self._collapsable_frame.rebuild()

            def reload_frame(prim):
                self._materials_changed()

            stage = self._payload.get_stage()
            mtl_name, _ = os.path.splitext(os.path.basename(url))

            def loaded_mdl_subids(mtl_list, mdl_path, filename, subid):
                if not mtl_list:
                    return

                if subid:
                    omni.kit.material.library.create_mdl_material(
                        stage=stage,
                        mtl_url=mdl_path,
                        mtl_name=subid,
                        on_create_fn=on_create_func,
                        mtl_real_name=subid)
                    return

                if len(mtl_list) > 1:
                    prim_paths = [prim_path for prim_path, material_name, strength in items]
                    omni.kit.material.library.custom_material_dialog(mdl_path=mdl_path, bind_prim_paths=prim_paths, on_complete_fn=reload_frame)
                    return

                omni.kit.material.library.create_mdl_material(
                    stage=stage,
                    mtl_url=mdl_path,
                    mtl_name=mtl_name,
                    on_create_fn=on_create_func,
                    mtl_real_name=mtl_list[0].name if len(mtl_list) > 0 else "")

            # get subids from material
            asyncio.ensure_future(omni.kit.material.library.get_subidentifier_from_mdl(mdl_file=url, on_complete_fn=lambda l, p=url, n=mtl_name: loaded_mdl_subids(mtl_list=l, mdl_path=p, filename=n, subid=subid), show_alert=True))

        def on_material_goto(model, items):
            prim_list =  sorted(list(set([item[1] for item in items if item[1] != Constant.SDF_PATH_INVALID])))
            if prim_list:
                omni.usd.get_context().get_selection().set_selected_prim_paths(prim_list, True)
            else:
                carb.log_warn(f"on_material_goto failed. No materials in list")

        self._combo_subscription = []
        self._first_row = True

        data = binding["material"]

        # show mixed material for all selected prims
        if len(data) > 1:
            self._build_binding_info(
                inherited=False,
                style_name="material_mixed",
                material_list=material_list,
                material_path=Constant.MIXED,
                strength_value=binding["strength"],
                bound_prim_list=list(binding["bound"]),
                on_material_fn=partial(on_material_changed, items=list(binding["bound_info"])),
                on_strength_fn=partial(on_strength_changed, relationships=list(binding["relationship"])),
                on_goto_fn=partial(on_material_goto, items=list(binding["bound_info"] | binding["inherited_info"])),
                on_dragdrop_fn=partial(on_dragdrop, items=list(binding["bound_info"])),
            )

        for material_path in data:
            if data[material_path]["bound"]:
                self._build_binding_info(
                    inherited=False,
                    style_name="material_solo",
                    material_list=material_list,
                    material_path=material_path,
                    strength_value=data[material_path]["strength"],
                    bound_prim_list=list(data[material_path]["bound"]),
                    on_material_fn=partial(on_material_changed, items=list(data[material_path]["bound_info"])),
                    on_strength_fn=partial(
                        on_strength_changed, relationships=list(data[material_path]["relationship"])
                    ),
                    on_goto_fn=partial(on_material_goto, items=list(data[material_path]["bound_info"])),
                    on_dragdrop_fn=partial(on_dragdrop, items=list(data[material_path]["bound_info"])),
                )

        for material_path in data:
            if data[material_path]["inherited"]:
                self._build_binding_info(
                    inherited=True,
                    style_name="material_solo",
                    material_list=material_list,
                    material_path=material_path,
                    strength_value=data[material_path]["strength"],
                    bound_prim_list=list(data[material_path]["inherited"]),
                    on_material_fn=partial(on_material_changed, items=list(data[material_path]["inherited_info"])),
                    on_strength_fn=partial(
                        on_strength_changed, relationships=list(data[material_path]["relationship"])
                    ),
                    on_goto_fn=partial(on_material_goto, items=list(data[material_path]["inherited_info"])),
                    on_dragdrop_fn=partial(on_dragdrop, items=list(data[material_path]["inherited_info"])),
                )

    def reset(self):
        if self._listener:
            self._listener.Revoke()
            self._listener = None

        if self._pending_frame_rebuild_task:
            self._pending_frame_rebuild_task.cancel()
            self._pending_frame_rebuild_task = None

    def _on_usd_changed(self, notice, stage):
        if self._pending_frame_rebuild_task:
            return
        items = set(notice.GetResyncedPaths()+notice.GetChangedInfoOnlyPaths())
        if any(path.elementString == ".material:binding" for path in items):
            async def rebuild():
                if self._collapsable_frame:
                    self._collapsable_frame.rebuild()
                self._pending_frame_rebuild_task = None

            self._pending_frame_rebuild_task = asyncio.ensure_future(rebuild())
