# Copyright (c) 2023 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 typing import (
    Optional,
    Sequence,
)

import omni.graph.core as og
from omni.kit.property.usd.custom_layout_helper import (
    CustomLayoutFrame,
    CustomLayoutGroup,
    CustomLayoutProperty,
)
from omni.kit.property.usd.usd_property_widget import UsdPropertyUiEntry


PROPS = (
    "inputs:enabled",
    "inputs:substepCount",
    "inputs:gravity",
    "inputs:globalScale",
    "inputs:contactElasticStiffness",
    "inputs:contactFrictionStiffness",
    "inputs:contactFrictionCoeff",
    "inputs:contactDampingStiffness",
    "inputs:particlesQueryRange",
    "inputs:particlesContactAdhesion",
    "inputs:particlesContactCohesion",
    "inputs:colliderContactDistance",
    "inputs:colliderContactQueryRange",
    "inputs:groundEnabled",
    "inputs:groundAltitude",
)


def find_prop(
    props: Sequence[UsdPropertyUiEntry],
    name: str,
) -> Optional[UsdPropertyUiEntry]:
    try:
        return next(x for x in props if x.prop_name == name)
    except StopIteration:
        return None


class CustomLayout:
    def __init__(self, compute_node_widget):
        self.enable = True
        self.compute_node_widget = compute_node_widget
        self.node_prim_path = self.compute_node_widget._payload[-1]
        self.node = og.Controller.node(self.node_prim_path)

    def apply(self, props) -> Sequence[UsdPropertyUiEntry]:
        user_props = tuple(find_prop(props, x) for x in PROPS)

        frame = CustomLayoutFrame(hide_extra=True)
        with frame:
            with CustomLayoutGroup("Inputs"):
                for prop in user_props:
                    if prop is not None:
                        CustomLayoutProperty(
                            prop.prop_name,
                            display_name=prop.metadata["customData"]["uiName"],
                        )

        return frame.apply(props)
