# Copyright (c) 2021, 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.
#
import omni.kit.context_menu
from typing import Set
from pxr import Usd, Sdf, UsdShade

material_clipboard = []

# menu functions....
def is_material_selected(objects):
    material_prim = objects["material_prim"]
    if material_prim:
        return True
    return False

def is_multiple_material_selected(objects):
    selected_paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
    mtl_list = []
    for path in selected_paths:
        prim = objects["stage"].GetPrimAtPath(path)
        if not prim:
            return False
        bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
        if bound_material and bound_material not in mtl_list:
            mtl_list.append(bound_material)
    return bool(len(mtl_list)>1)

def is_material_copied(objects):
    return bool(len(material_clipboard) > 0)

def is_bindable_prim_selected(objects):
    for prim_path in omni.usd.get_context().get_selection().get_selected_prim_paths():
        prim = objects["stage"].GetPrimAtPath(prim_path)
        if prim and omni.usd.is_prim_material_supported(prim):
            return True
    return False

def is_prim_selected(self, objects: dict):
    """
    checks if any prims are selected
    """
    if not any(item in objects for item in ["prim", "prim_list"]):
        return False
    return True

def is_paste_all(objects):
    return bool(len(omni.usd.get_context().get_selection().get_selected_prim_paths()) > len(material_clipboard))

def goto_material(objects):
    material_prim = objects["material_prim"]
    if material_prim:
        omni.usd.get_context().get_selection().set_prim_path_selected(material_prim.GetPath().pathString, True, True, True, True)

def copy_material(objects: dict, all: bool):
    global material_clipboard

    material_clipboard = []
    for prim in objects["prim_list"]:
        mat, rel = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
        material_path = mat.GetPath().pathString if mat else None
        material_strength = UsdShade.MaterialBindingAPI.GetMaterialBindingStrength(rel) if rel else UsdShade.Tokens.weakerThanDescendants
        material_clipboard.append((material_path, material_strength))

def paste_material(objects: dict):
    selected_paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
    prims = [objects["stage"].GetPrimAtPath(p) for p in selected_paths]

    with omni.kit.undo.group():
        max_paste = len(material_clipboard)
        for index, prim in enumerate(prims):
            if index >=  max_paste:
                break
            material_path, material_strength = material_clipboard[index]
            if material_path:
                omni.kit.commands.execute(
                    "BindMaterial",
                    prim_path=prim.GetPath(),
                    material_path=material_path,
                    strength=material_strength,
                )

def paste_material_all(objects: dict):
    selected_paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
    prims = [objects["stage"].GetPrimAtPath(p) for p in selected_paths]

    with omni.kit.undo.group():
        max_paste = len(material_clipboard)
        for index, prim in enumerate(prims):
            material_path, material_strength = material_clipboard[index % max_paste]
            if material_path:
                omni.kit.commands.execute(
                    "BindMaterial",
                    prim_path=prim.GetPath(),
                    material_path=material_path,
                    strength=material_strength,
                )

def get_paste_name(objects):
    if len(material_clipboard) <2:
        return "Paste"
    return f"Paste ({len(material_clipboard)} Materials)"


def show_context_menu(stage, material_path, bound_prims: Set[Usd.Prim]):
    # get context menu core functionality & check its enabled
    context_menu = omni.kit.context_menu.get_instance()
    if context_menu is None:
        carb.log_error("context_menu is disabled!")
        return None

    # setup objects, this is passed to all functions
    if not stage:
        return
    objects = {"material_prim": stage.GetPrimAtPath(material_path), "prim_list": bound_prims}

    # setup menu
    menu_dict = [
        {"name": "Copy", "enabled_fn": [], "onclick_fn": lambda o: copy_material(o, False)},
        {"name_fn": get_paste_name, "show_fn": [is_bindable_prim_selected, is_material_copied], "onclick_fn": paste_material},
        {"name": "", "show_fn": is_material_selected},
        {"name": "Select Material", "show_fn": is_material_selected, "onclick_fn": goto_material},
    ]

    # show menu
    menu_dict += omni.kit.context_menu.get_menu_dict("MENU_THUMBNAIL", "")
    omni.kit.context_menu.reorder_menu_dict(menu_dict)
    context_menu.show_context_menu("material_thumbnail", objects, menu_dict)
