# 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.
#

import weakref
import asyncio
from functools import lru_cache, partial
from typing import Callable, Union, Any

import carb
import omni.client
import omni.client.utils as clientutils
import omni.kit.app
import omni.kit.commands
import omni.kit.ui
import omni.ui as ui
import omni.usd
from omni.kit.window.file_importer import get_file_importer, ImportOptionsDelegate
from pxr import Sdf, Tf, Usd

from .asset_filepicker import show_asset_file_picker, replace_query
from .prim_selection_payload import PrimSelectionPayload
from .usd_property_widget import UsdPropertiesWidget
from .usd_property_widget_builder import UsdPropertiesWidgetBuilder, get_ui_style
from .widgets import ICON_PATH
from .versioning_helper import VersioningHelper

DEFAULT_PRIM_TAG = "<Default Prim>"
REF_LABEL_WIDTH = 80


@lru_cache()
def _get_plus_glyph():
    return omni.kit.ui.get_custom_glyph_code("${glyphs}/menu_context.svg")


def anchor_reference_asset_path_to_layer(ref: Sdf.Reference, intro_layer: Sdf.Layer, anchor_layer: Sdf.Layer):
    asset_path = ref.assetPath
    if asset_path:
        asset_path = intro_layer.ComputeAbsolutePath(asset_path)
        asset_url = clientutils.make_relative_url_if_possible(anchor_layer.identifier, asset_path)

        # make a copy as Reference is immutable
        ref = Sdf.Reference(
            assetPath=asset_url,
            primPath=ref.primPath,
            layerOffset=ref.layerOffset,
            customData=ref.customData,
        )
    return ref


def anchor_payload_asset_path_to_layer(ref: Sdf.Payload, intro_layer: Sdf.Layer, anchor_layer: Sdf.Layer):
    asset_path = ref.assetPath
    if asset_path:
        asset_path = intro_layer.ComputeAbsolutePath(asset_path)
        asset_url = clientutils.make_relative_url_if_possible(anchor_layer.identifier, asset_path)

        # make a copy as Payload is immutable
        ref = Sdf.Payload(
            assetPath=asset_url,
            primPath=ref.primPath,
            layerOffset=ref.layerOffset,
        )
    return ref


# Model that calls edited_fn when end_edit or set_value (not when typing)
class TrackEditingStringModel(ui.SimpleStringModel):
    def __init__(self, value: str = ""):
        super().__init__()
        self._editing = False
        self._edited_fn = []
        self.set_value(value)

    def begin_edit(self):
        super().begin_edit()
        self._editing = True

    def end_edit(self):
        self._editing = False
        super().end_edit()
        self._call_edited_fn()

    def set_value(self, value: str):
        super().set_value(value)
        if not self._editing:
            self._call_edited_fn()

    def add_edited_fn(self, fn):
        self._edited_fn.append(fn)

    def clear_edited_fn(self):
        self._edited_fn.clear()

    def _call_edited_fn(self):
        for fn in self._edited_fn:
            fn(self)


class PayloadReferenceInfo:
    def __init__(self, asset_path_field, prim_path_field, prim_path_field_model, checkpoint_model):
        self.asset_path_field = asset_path_field
        self.prim_path_field = prim_path_field
        self.prim_path_field_model = prim_path_field_model
        self.checkpoint_model = checkpoint_model

    def destroy(self):
        if self.checkpoint_model:
            self.checkpoint_model.destroy()
        self.asset_path_field = None
        self.prim_path_field = None
        self.prim_path_field_model = None
        self.checkpoint_model = None


def build_path_field(stage, init_value: str, jump_button: bool, use_payloads: bool, layer: Sdf.Layer = None):
    with ui.HStack(style={"Button": {"margin": 0, "padding": 3, "border_radius": 2, "background_color": 0x00000000}}):

        def assign_value_fn(model, path):
            model.set_value(path)

        def assign_reference_value(stage_weak, model_weak, path: str, assign_value_fn: Callable[[Any, str], None], frame=None):
            stage = stage_weak()
            if not stage:
                return

            model = model_weak()
            if not model:
                return

            edit_layer = stage.GetEditTarget().GetLayer()
            # make the path relative to current edit target layer
            relative_url = clientutils.make_relative_url_if_possible(edit_layer.identifier, path)
            assign_value_fn(model, relative_url)

        def drop_accept(url: str):
            if len(url.split('\n')) > 1:
                carb.log_warn(f"build_path_field multifile drag/drop not supported")
                return False

            # TODO support filtering by file extension
            if "." not in url:
                # TODO dragging from stage view also result in a drop, which is a prim path not an asset path
                # For now just check if dot presents in the url (indicating file extension).
                return False
            return True

        with ui.ZStack():
            model = TrackEditingStringModel(value=init_value)
            string_field = ui.StringField(model=model)
            string_field.set_accept_drop_fn(drop_accept)
            string_field.set_drop_fn(
                lambda event, model_weak=weakref.ref(model), stage_weak=weakref.ref(stage): assign_reference_value(
                    stage_weak, model_weak, event.mime_data, assign_value_fn=assign_value_fn
                )
            )
        ui.Spacer(width=3)

        def on_refresh_layer(model, weak_layer):
            intro_layer = weak_layer()
            url = model.get_value_as_string()
            if url and intro_layer:
                layer = Sdf.Find(intro_layer.ComputeAbsolutePath(url))
                if layer:
                    async def reload(layer):
                        layer.Reload(force=True)

                    asyncio.ensure_future(reload(layer))

        changed = False
        ui.Button(
            style={
                "image_url": str(ICON_PATH.joinpath("refresh.svg")),
                "padding": 1,
                "margin": 0,
                "Button.Image": {
                    "color": 0xFF9E9E9E if not changed else 0xFF248AE3,
                },
                "Button:hovered": {
                    "background_color": 0x00000000
                }
            },
            width=20,
            tooltip="Refresh Payload" if use_payloads else "Refresh Reference",
            clicked_fn=lambda model=model, layer_weak=weakref.ref(layer) if layer else None: on_refresh_layer(model, layer_weak),
        )

        style = {
            "image_url": str(ICON_PATH.joinpath("small_folder.png")),
            "padding": 1,
            "margin": 0,
            "Button:hovered": {
                "background_color": 0x00000000
                }
        }
        ui.Button(
            style=style,
            width=20,
            tooltip="Browse..." if get_file_importer() is not None else "File importer not available",
            clicked_fn=lambda model_weak=weakref.ref(model), stage_weak=weakref.ref(stage), layer_weak=weakref.ref(
                layer
            ) if layer else None: show_asset_file_picker(
                "Select Payload..." if use_payloads else "Select Reference...",
                assign_value_fn,
                model_weak,
                stage_weak,
                layer_weak=layer_weak,
                on_selected_fn=assign_reference_value,
            ),
            enabled=get_file_importer() is not None,
        )
        if jump_button:
            ui.Spacer(width=3)

            # Button to jump to the file in Content Window
            def locate_file(model, weak_layer):
                # omni.kit.window.content_browser is optional dependency
                try:
                    url = model.get_value_as_string()
                    if len(url) == 0:
                        return

                    if weak_layer:
                        weak_layer = weak_layer()
                        if weak_layer:
                            url = weak_layer.ComputeAbsolutePath(url)

                    # Remove the checkpoint and branch so navigate_to works
                    url = replace_query(url, None)

                    import omni.kit.window.content_browser

                    instance = omni.kit.window.content_browser.get_instance()
                    instance.navigate_to(url)
                except Exception as e:
                    carb.log_warn(f"Failed to locate file: {e}")

            style["image_url"] = str(ICON_PATH.joinpath("find.png"))
            ui.Button(
                style=style,
                width=20,
                tooltip="Locate File",
                clicked_fn=lambda model=model, weak_layer=weakref.ref(layer) if layer else None: locate_file(
                    model, weak_layer
                ),
            )

    return string_field


class AddPayloadReferenceOptionsDelegate(ImportOptionsDelegate):
    def __init__(self, prim_path_model):
        super().__init__(
            build_fn=self._build_ui_impl,
            destroy_fn=self._destroy_impl
        )
        self._widget = None
        self._prim_path_model = prim_path_model

    def should_load_payload(self):
        if self._load_payload_checkbox:
            return self._load_payload_checkbox.model.get_value_as_bool()

    def _build_ui_impl(self):
        self._widget = ui.Frame()
        with self._widget:
            with ui.HStack(height=0, spacing=2):
                ui.Label("Prim Path", width=0)
                ui.StringField().model = self._prim_path_model

    def _destroy_impl(self, _):
        self._prim_path_model = None
        self._widget = None


class AddPayloadReferenceWindow:
    def __init__(self, payload: PrimSelectionPayload, on_payref_added_fn: Callable, use_payloads=False):
        self._payrefs = None
        self._use_payloads = use_payloads

        self._prim_path_model = ui.SimpleStringModel()
        self._stage = payload.get_stage()
        self._payload = payload
        self._on_payref_added_fn = on_payref_added_fn

    def set_payload(self, payload: PrimSelectionPayload):
        self._stage = payload.get_stage()
        self._payload = payload

    def show(self, payrefs: Union[Sdf.Reference, Sdf.Payload]):
        self._payrefs = payrefs
        fallback = None
        if self._stage and not self._stage.GetRootLayer().anonymous:
            # If asset path is empty, open the USD rootlayer folder
            # But only if filepicker didn't already have a folder remembered (thus fallback)
            fallback = self._stage.GetRootLayer().identifier

        self._prim_path_model.set_value(DEFAULT_PRIM_TAG)

        def _on_file_selected(weak_self, filename, dirname, selections=[]):
            weak_self = weak_self()
            if not weak_self:
                return

            if not selections:
                return

            asset_path = f"{dirname or ''}/{filename}"
            edit_layer = weak_self._stage.GetEditTarget().GetLayer()
            # make the path relative to current edit target layer
            asset_url = clientutils.make_relative_url_if_possible(edit_layer.identifier, asset_path)

            prim_path = self._prim_path_model.get_value_as_string()
            if prim_path == DEFAULT_PRIM_TAG:
                prim_path = Sdf.Path()

            payrefs = weak_self._payrefs
            if not payrefs:
                stage = weak_self._stage
                payload_prim_path = weak_self._payload[-1]
                if not stage or not payload_prim_path:
                    carb.log_warn(f"Cannot create payload/reference as stage/prim {payload_prim_path} is invalid")
                    return
                anchor_prim = stage.GetPrimAtPath(payload_prim_path)
                if not anchor_prim:
                    carb.log_warn(f"Cannot create payload/reference as failed to get prim {payload_prim_path}")
                    return

                if weak_self._use_payloads:
                    anchor_prim.GetPayloads().SetPayloads([])
                    payrefs = anchor_prim.GetPayloads()
                    if not payrefs:
                        carb.log_warn("Failed to create payload")
                        return
                else:
                    anchor_prim.GetReferences().SetReferences([])
                    payrefs = anchor_prim.GetReferences()
                    if not payrefs:
                        carb.log_warn("Failed to create reference")
                        return

            command_kwargs = {
                'stage': payrefs.GetPrim().GetStage(),
                'prim_path': payrefs.GetPrim().GetPath()
            }
            if weak_self._use_payloads:
                if prim_path:
                    payref = Sdf.Payload(asset_url, prim_path)
                else:
                    payref = Sdf.Payload(asset_url)
                command = "AddPayload"
                command_kwargs['payload'] = payref
            else:
                if prim_path:
                    payref = Sdf.Reference(asset_url, prim_path)
                else:
                    payref = Sdf.Reference(asset_url)
                command = "AddReference"
                command_kwargs['reference'] = payref

            omni.kit.commands.execute(command, **command_kwargs)

            weak_self._on_payref_added_fn(payrefs.GetPrim())

        file_importer = get_file_importer()
        if file_importer:
            file_importer.show_window(
                title="Select Payload..." if self._use_payloads else "Select Reference...",
                import_button_label="Select",
                import_handler=partial(_on_file_selected, weakref.ref(self)),
            )
        if fallback and not file_importer._dialog.get_current_directory():
            file_importer._dialog.set_current_directory(fallback)
        file_importer.add_import_options_frame(
            "Payload Options" if self._use_payloads else "Reference Options",
            AddPayloadReferenceOptionsDelegate(self._prim_path_model)
        )

    def destroy(self):
        self._payrefs = None
        self._prim_path_model = None


class PayloadReferenceWidget(UsdPropertiesWidget):
    def __init__(self, use_payloads=False):
        super().__init__(title="Payloads" if use_payloads else "References", collapsed=False, multi_edit=False)
        self._ref_list_op = None
        self._payrefs = None
        self._ref_dict = {}
        self._add_ref_window = [None, None]
        self._ref_and_layers = []
        self._checkpoint_combobox = None
        self._use_payloads = use_payloads
        self._payload_loaded_cb = None

        self._stage_event_sub = (
            omni.usd.get_context()
            .get_stage_event_stream()
            .create_subscription_to_pop(self._on_stage_event, name="PayloadReferenceWidget stage update")
        )

        # +add menu item(s)
        from .prim_path_widget import PrimPathWidget

        PrimPathWidget.add_button_menu_entry(
            "Payload" if use_payloads else "Reference", show_fn=self._prim_is_selected, onclick_fn=lambda payload, u=self._use_payloads: self._on_add_payload_reference(payload, u)
        )

    def _undo_redo_on_change(self, cmds):
        async def update_value(cmds):
            if cmds and "SetPayLoadLoadSelectedPrimsCommand" in cmds:
                last_prim = self._get_prim(self._payload[-1])
                if last_prim:
                    self._payload_loaded_cb.model.set_value(last_prim.IsLoaded())
        if cmds and self._payload_loaded_cb:
            asyncio.ensure_future(update_value(cmds))

    def _on_stage_event(self, event):
        if event.type == int(omni.usd.StageEventType.CLOSING) or event.type == int(omni.usd.StageEventType.OPENED):
            for window in self._add_ref_window:
                if window:
                    window.destroy()
            self._add_ref_window = [None, None]

    def _on_add_payload_reference(self, payload: PrimSelectionPayload, use_payloads: bool):
        ref_window_index = 0 if self._use_payloads else 1
        if not self._add_ref_window[ref_window_index]:
            self._add_ref_window[ref_window_index] = AddPayloadReferenceWindow(payload, self._on_payload_reference_added, use_payloads=use_payloads)
        else:
            self._add_ref_window[ref_window_index].set_payload(payload)
        self._add_ref_window[ref_window_index].show(self._payrefs)

    def _on_payload_reference_added(self, prim: Usd.Prim):
        property_window = omni.kit.window.property.get_window()
        if property_window:
            property_window.request_rebuild()

    def _prim_is_selected(self, objects: dict):
        if not "prim_list" in objects or not "stage" in objects:
            return False
        stage = objects["stage"]
        if not stage:
            return False
        return len(objects["prim_list"]) == 1

    def clean(self):
        self.reset()
        for window in self._add_ref_window:
            if window:
                window.destroy()
        self._add_ref_window = [None, None]
        self._stage_event_sub = None
        super().clean()

    def reset(self):
        if self._listener:
            self._listener.Revoke()
            self._listener = None
        omni.kit.undo.unsubscribe_on_change(self._undo_redo_on_change)
        self._ref_list_op = None
        self._payrefs = None
        for ref, info in self._ref_dict.items():
            info.destroy()
        self._ref_dict.clear()
        self._ref_and_layers.clear()
        if self._checkpoint_combobox:
            self._checkpoint_combobox.destroy()
            self._checkpoint_combobox = None
        super().reset()

    def on_new_payload(self, payload):
        """
        See PropertyWidget.on_new_payload
        """

        if not super().on_new_payload(payload):
            return False

        if not self._payload or len(self._payload) != 1:  # single edit for now
            return False

        anchor_prim = None
        for prim_path in self._payload:
            prim = self._get_prim(prim_path)
            if not prim:
                return False
            anchor_prim = prim

        # only show if prim has payloads/references
        if anchor_prim:
            if self._use_payloads:
                ref_and_layers = omni.usd.get_composed_payloads_from_prim(anchor_prim)
                return len(ref_and_layers) > 0

            ref_and_layers = omni.usd.get_composed_references_from_prim(anchor_prim)
            return len(ref_and_layers) > 0

        return False

    def build_impl(self):
        """
        See PropertyWidget.build_impl
        """
        if not self._use_payloads:
            super().build_impl()
        else:
            def on_checkbox_toggle(layer_name, model):
                async def set_loaded(layer_name, state):
                    omni.kit.commands.execute("SetPayLoadLoadSelectedPrimsCommand", selected_paths=[layer_name], value=state)

                asyncio.ensure_future(set_loaded(layer_name, model.get_value_as_bool()))

            last_prim = self._get_prim(self._payload[-1])
            if not last_prim:
                return

            self._button_frame = ui.Frame()
            with self._button_frame:
                with ui.ZStack():
                    super().build_impl()
                    with ui.HStack():
                        ui.Spacer(width=ui.Fraction(0.5))
                        with ui.VStack(width=0, content_clipping=True):
                            ui.Spacer(height=5)
                            self._payload_loaded_cb = ui.CheckBox()
                            self._payload_loaded_cb.model.set_value(last_prim.IsLoaded())
                            self._payload_loaded_cb.model.add_value_changed_fn(
                                partial(on_checkbox_toggle, last_prim.GetPath().pathString))
                        ui.Spacer(width=5)

    def build_items(self):
        self.reset()

        if len(self._payload) == 0:
            return

        last_prim = self._get_prim(self._payload[-1])
        if not last_prim:
            return

        stage = last_prim.GetStage()
        if not stage:
            return

        self._listener = Tf.Notice.Register(Usd.Notice.ObjectsChanged, self._on_usd_changed, stage)
        omni.kit.undo.subscribe_on_change(self._undo_redo_on_change)

        if self._use_payloads:
            self._payrefs = last_prim.GetPayloads()
            self._ref_and_layers = omni.usd.get_composed_payloads_from_prim(last_prim, False)
        else:
            self._payrefs = last_prim.GetReferences()
            self._ref_and_layers = omni.usd.get_composed_references_from_prim(last_prim, False)

        ref_window_index = 0 if self._use_payloads else 1
        self._add_ref_window[ref_window_index] = AddPayloadReferenceWindow(self._payload, self._on_payload_reference_added, use_payloads=self._use_payloads)
        with ui.VStack(height=0, spacing=5, name="frame_v_stack"):
            if self._payrefs:
                for (ref, layer) in self._ref_and_layers:
                    self._build_payload_reference(ref, layer)

            def on_add_payload_reference(weak_self):
                weak_self = weak_self()
                if not weak_self:
                    return

                weak_self._add_ref_window[ref_window_index].show(weak_self._payrefs)

            ui.Button(
                f"{_get_plus_glyph()} Add Payload" if self._use_payloads else f"{_get_plus_glyph()} Add Reference",
                clicked_fn=lambda weak_self=weakref.ref(self): on_add_payload_reference(weak_self),
            )

    def _build_payload_reference(self, payref: Sdf.Reference, intro_layer: Sdf.Layer):
        with ui.Frame():
            with ui.ZStack():
                ui.Rectangle(name="backdrop")
                with ui.VStack(
                    name="ref_group", spacing=5, style={"VStack::ref_group": {"margin_width": 2, "margin_height": 2}}
                ):
                    with ui.HStack(spacing=5):
                        ui.Spacer(width=5)
                        with ui.VStack(spacing=5):
                            asset_path_field = self._build_asset_path_ui(payref, intro_layer)
                            prim_path_field, prim_path_field_model = self._build_prim_path_ui(payref, intro_layer)
                            checkpoint_model = self._build_checkpoint_ui(payref, intro_layer)
                        self._build_remove_payload_reference_button(payref, intro_layer)

                self._ref_dict[payref] = PayloadReferenceInfo(
                    asset_path_field, prim_path_field, prim_path_field_model, checkpoint_model
                )

    def _build_asset_path_ui(self, payref: Union[Sdf.Reference, Sdf.Payload], intro_layer: Sdf.Layer):
        with ui.HStack():
            UsdPropertiesWidgetBuilder._create_label("Asset Path", additional_label_kwargs={"width": REF_LABEL_WIDTH})
            asset_path_field = build_path_field(
                self._payload.get_stage(),
                payref.assetPath,
                True,
                self._use_payloads,
                intro_layer,
            )

            # If payload/reference is not existed
            if payref.assetPath:
                absolute_path = intro_layer.ComputeAbsolutePath(payref.assetPath)
                status, _ = omni.client.stat(absolute_path)
                if status != omni.client.Result.OK:
                    asset_path_field.set_style({"color": 0xFF6F72FF})

            asset_path_field.model.add_edited_fn(
                lambda model, stage=self._payrefs.GetPrim().GetStage(), prim_path=self._payrefs.GetPrim().GetPath(), payref=payref: self._on_payload_reference_edited(
                    model, stage, prim_path, payref, intro_layer
                )
            )
        return asset_path_field

    def _build_prim_path_ui(self, payref: Union[Sdf.Reference,Sdf.Payload], intro_layer: Sdf.Layer):
        with ui.HStack():
            UsdPropertiesWidgetBuilder._create_label("Prim Path", additional_label_kwargs={"width": REF_LABEL_WIDTH})
            prim_path = payref.primPath
            if not prim_path:
                prim_path = DEFAULT_PRIM_TAG
            prim_path_field_model = TrackEditingStringModel(str(prim_path))
            prim_path_field = ui.StringField(model=prim_path_field_model)
            prim_path_field.model.add_edited_fn(
                lambda model, stage=self._payrefs.GetPrim().GetStage(), prim_path=self._payrefs.GetPrim().GetPath(), payref=payref: self._on_payload_reference_edited(
                    model, stage, prim_path, payref, intro_layer
                )
            )
        return prim_path_field, prim_path_field_model

    def _build_checkpoint_ui(self, payref: Union[Sdf.Reference,Sdf.Payload], intro_layer: Sdf.Layer):
        if VersioningHelper.is_versioning_enabled():
            try:
                # Use checkpoint widget in the drop down menu for more detailed information
                from omni.kit.widget.versioning.checkpoint_combobox import CheckpointCombobox

                stack = ui.HStack()
                with stack:
                    UsdPropertiesWidgetBuilder._create_label(
                        "Checkpoint", additional_label_kwargs={"width": REF_LABEL_WIDTH}
                    )

                    absolute_asset_path = payref.assetPath
                    if len(absolute_asset_path):
                        absolute_asset_path = omni.client.combine_urls(intro_layer.identifier, absolute_asset_path)

                    def on_selection_changed(
                        selection,
                        stage: Usd.Stage,
                        prim_path: Sdf.Path,
                        payref: Union[Sdf.Reference, Sdf.Payload],
                        intro_layer: Sdf.Layer,
                    ):
                        self._on_payload_reference_checkpoint_edited(selection, stage, prim_path, payref)

                    self._checkpoint_combobox = CheckpointCombobox(
                        absolute_asset_path,
                        lambda selection, stage=self._payrefs.GetPrim().GetStage(), prim_path=self._payrefs.GetPrim().GetPath(), payref=payref: on_selection_changed(
                            selection, stage, prim_path, payref, intro_layer
                        ),
                    )

                    # reset button
                    def reset_func(payref, stage, prim_path):
                        client_url = omni.client.break_url(payref.assetPath)
                        on_selection_changed(None, stage, prim_path, payref, intro_layer)

                    checkpoint = ""
                    client_url = omni.client.break_url(payref.assetPath)
                    if client_url.query:
                        _, checkpoint = omni.client.get_branch_and_checkpoint_from_query(client_url.query)
                    ui.Spacer(width=4)
                    ui.Image(
                        f"{ICON_PATH}/Default value.svg" if checkpoint == "" else f"{ICON_PATH}/Changed value.svg",
                        mouse_pressed_fn=lambda x, y, b, a, s=self._payrefs.GetPrim().GetStage(), r=payref, p=self._payrefs.GetPrim().GetPath(): reset_func(r, s, p),
                        width=12,
                        height=18,
                        tooltip="Reset Checkpoint" if checkpoint else "",
                    )

                    def on_have_server_info(server:str, support_checkpoint: bool, stack: ui.HStack):
                        if not support_checkpoint:
                            stack.visible = False

                    VersioningHelper.check_server_checkpoint_support(
                        VersioningHelper.extract_server_from_url(absolute_asset_path),
                        lambda s, c, t=stack: on_have_server_info(s,c, t)
                    )

                return None
            except ImportError as e:
                # If the widget is not available, create a simple combo box instead
                carb.log_warn(f"Checkpoint widget in Payload/Reference is not availble due to: {e}")

    def _build_remove_payload_reference_button(self, payref, intro_layer):
        def on_remove_payload_reference(ref, layer_weak, payrefs):
            if not ref or not payrefs:
                return

            stage = payrefs.GetPrim().GetStage()
            edit_target_layer = stage.GetEditTarget().GetLayer()
            intro_layer = layer_weak() if layer_weak else None

            if self._use_payloads:
                # When removing a payload on a different layer, the deleted assetPath should be relative to edit target layer, not introducing layer
                if intro_layer and intro_layer != edit_target_layer:
                    ref = anchor_payload_asset_path_to_layer(ref, intro_layer, edit_target_layer)

                omni.kit.commands.execute(
                    "RemovePayload",
                    stage=stage,
                    prim_path=payrefs.GetPrim().GetPath(),
                    payload=ref,
                ),
            else:
                # When removing a reference on a different layer, the deleted assetPath should be relative to edit target layer, not introducing layer
                if intro_layer and intro_layer != edit_target_layer:
                    ref = anchor_reference_asset_path_to_layer(ref, intro_layer, edit_target_layer)

                omni.kit.commands.execute(
                    "RemoveReference",
                    stage=stage,
                    prim_path=payrefs.GetPrim().GetPath(),
                    reference=ref,
                ),

        style = {"image_url": str(ICON_PATH.joinpath("remove.svg")), "margin": 0, "padding": 0}
        ui.Button(
            style=style,
            clicked_fn=lambda ref=payref, layer_weak=weakref.ref(
                intro_layer
            ) if intro_layer else None, refs=self._payrefs: on_remove_payload_reference(ref, layer_weak, refs),
            width=16,
        )

    def _on_payload_reference_edited(
        self, model_or_item, stage: Usd.Stage, prim_path: Sdf.Path, payref: Union[Sdf.Reference, Sdf.Payload], intro_layer: Sdf.Layer
    ):
        ref_prim_path = self._ref_dict[payref].prim_path_field.model.get_value_as_string()
        ref_prim_path = Sdf.Path(ref_prim_path) if ref_prim_path and ref_prim_path != DEFAULT_PRIM_TAG else Sdf.Path()

        new_asset_path = self._ref_dict[payref].asset_path_field.model.get_value_as_string()

        try:
            from omni.kit.widget.versioning.checkpoints_model import CheckpointItem

            if isinstance(model_or_item, CheckpointItem):
                new_asset_path = replace_query(new_asset_path, model_or_item.get_relative_path())
            elif model_or_item is None:
                new_asset_path = replace_query(new_asset_path, None)
        except Exception as e:
            pass

        if self._use_payloads:
            new_ref = Sdf.Payload(assetPath=new_asset_path.replace("\\", "/"), primPath=ref_prim_path)

            edit_target_layer = stage.GetEditTarget().GetLayer()
            # When replacing a payload on a different layer, the replaced assetPath should be relative to edit target layer, not introducing layer
            if intro_layer != edit_target_layer:
                payref = anchor_payload_asset_path_to_layer(payref, intro_layer, edit_target_layer)

            if payref == new_ref:
                return False

            omni.kit.commands.execute(
                "ReplacePayload",
                stage=stage,
                prim_path=prim_path,
                old_payload=payref,
                new_payload=new_ref,
            )

            return True

        new_ref = Sdf.Reference(assetPath=new_asset_path.replace("\\", "/"), primPath=ref_prim_path)

        edit_target_layer = stage.GetEditTarget().GetLayer()
        # When replacing a reference on a different layer, the replaced assetPath should be relative to edit target layer, not introducing layer
        if intro_layer != edit_target_layer:
            payref = anchor_reference_asset_path_to_layer(payref, intro_layer, edit_target_layer)

        if payref == new_ref:
            return False

        omni.kit.commands.execute(
            "ReplaceReference",
            stage=stage,
            prim_path=prim_path,
            old_reference=payref,
            new_reference=new_ref,
        )

        return True

    def _on_payload_reference_checkpoint_edited(
        self, model_or_item, stage: Usd.Stage, prim_path: Sdf.Path, payref: Union[Sdf.Reference, Sdf.Payload]
    ):
        new_asset_path = self._ref_dict[payref].asset_path_field.model.get_value_as_string()

        try:
            from omni.kit.widget.versioning.checkpoints_model import CheckpointItem

            if isinstance(model_or_item, CheckpointItem):
                new_asset_path = replace_query(new_asset_path, model_or_item.get_relative_path())
            elif model_or_item is None:
                new_asset_path = replace_query(new_asset_path, None)
        except Exception as e:
            pass

        if self._use_payloads:
            new_ref = Sdf.Payload(assetPath=new_asset_path.replace("\\", "/"), primPath=payref.primPath, layerOffset=payref.layerOffset)
            if payref != new_ref:
                omni.kit.commands.execute(
                    "ReplacePayload",
                    stage=stage,
                    prim_path=prim_path,
                    old_payload=payref,
                    new_payload=new_ref,
                )
            return

        new_ref = Sdf.Reference(assetPath=new_asset_path.replace("\\", "/"), primPath=payref.primPath, layerOffset=payref.layerOffset)
        if payref != new_ref:
            omni.kit.commands.execute(
                "ReplaceReference",
                stage=stage,
                prim_path=prim_path,
                old_reference=payref,
                new_reference=new_ref,
            )
