# Copyright (c) 2025 The Foundry Visionmongers Ltd. All Rights Reserved.
"""
Module containing an example Parameters Handler plug-in for UsdSuperLayer-derived node types.
"""

import logging

import usg

import NodegraphAPI
from NodesUsdAPI.UsdSuperLayer import ParametersHandlerPlugin, Utils

log = logging.getLogger("ExampleParametersHandlerPlugin")


class UsdSuperLayerExampleParametersHandlerPlugin(
    ParametersHandlerPlugin.ParametersHandlerPluginBase,
):
    """
    An example of a ``UsdSuperLayerParametersHandler`` plug-in that adds a new parameter group called
    ``'image'`` that includes some user-defined parameters that are then translated into USD
    Attributes on the target USD Prim in the UsdSuperLayer-derived nodes.
    """

    @classmethod
    def isPrimSupported(
        cls,
        node: NodegraphAPI.Node,
        primPath: str,
        usdPrim: usg.Prim | None,
        sdfPrim: usg.Prim | None,
    ) -> bool:
        """
        Implements `ParametersHandlerPluginBase.isPrimSupported()`.

        Invoked by UsdSuperLayer-derived nodes when the selected prim in the node changes, the prim
        in the incoming USD Stage changes, or the prim in the content layer changes.

        Plug-ins must implement the function to specify whether a given prim is to be handled by
        this plug-in. If supported, :py:obj:`buildOrUpdateParameters()` will be invoked next after
        a dedicated group parameter has been created for this prim path and plug-in.

        :param node: The UsdSuperLayer-derived node.
        :param primPath: The absolute path to the prim.
        :param usdPrim: The USD Prim in the stage if it exists.
        :param sdfPrim: The USD Prim in the content layer if it exists.
        :return: Whether the prim should be handled by the plug-in, in which case
            :py:obj:`buildOrUpdateParameters()` will be invoked next.
        """
        # All prims are supported and will create parameters defined by this plug-in.
        return True

    @classmethod
    def getGroupParameterLabel(cls) -> str:
        """
        Implements `ParametersHandlerPluginBase.getGroupParameterLabel()`.

        Invoked by UsdSuperLayer-derived nodes.

        Plug-ins can implement the function to define the label that will be used for the group
        parameter widget that holds the parameters that the plug-in creates.

        :return: The label for the group parameter that will hold the parameters that the plug-in
            creates. Common labels can be ``'properties'``, ``'inputs'``, or ``'transform'``.
        """
        return 'image'

    @classmethod
    def buildOrUpdateParameters(
        cls,
        node: NodegraphAPI.Node,
        primPath: str,
        usdPrim: usg.Prim | None,
        sdfPrim: usg.Prim | None,
        groupParam: NodegraphAPI.Parameter,
    ) -> None:
        """
        Implements `ParametersHandlerPluginBase.buildOrUpdateParameters()`.

        Invoked by UsdSuperLayer-derived nodes if :py:obj:`isPrimSupported()` in the plug-in has
        returned ``True``.

        Plug-ins can implement the function to create their own parameters based on the prim in the
        incoming stage.

        :param node: The UsdSuperLayer-derived node.
        :param primPath: The absolute path to the prim.
        :param usdPrim: The USD Prim in the stage if it exists.
        :param sdfPrim: The USD Prim in the content layer if it exists.
        :param groupParam: The group parameter, previously manufactured for the plug-in only, in
            which the plug-in can create its relevant child parameters.
        """
        # Build parameters if not yet created. These parameters include a **stageparameter** child
        # parameter that enables the StageParameterPolicy in the group parameter.
        filepathGroupParam = groupParam.getChild('filepath')
        if filepathGroupParam is None:
            filepathGroupParam = groupParam.createChildGroup('filepath')
            filepathGroupParam.createChildNumber('enable', 0)
            filepathGroupParam.createChildString('default', '')
            filepathGroupParam.createChildString('value', '')
            filepathGroupParam.createChildNumber('stageparameter', 0.0)
            filepathGroupParam.setHintString(
                str({
                    'widget': 'assetIdInput',
                    'fileTypes': 'png|gif|jpg|webp',
                    'help': 'An example parameter where the user can select an image filepath.',
                }),
            )
        scaleGroupParam = groupParam.getChild('scale')
        if scaleGroupParam is None:
            scaleGroupParam = groupParam.createChildGroup('scale')
            scaleGroupParam.createChildNumber('enable', 0)
            scaleGroupParam.createChildNumber('default', 0.0)
            scaleGroupParam.createChildNumber('value', 0.0)
            scaleGroupParam.createChildNumber('stageparameter', 0.0)
            scaleGroupParam.setHintString(
                str({'help': 'Another parameter with a hypothetical image scale.'}),
            )
        modeGroupParam = groupParam.getChild('mode')
        if modeGroupParam is None:
            modeGroupParam = groupParam.createChildGroup('mode')
            modeGroupParam.createChildNumber('enable', 0)
            modeGroupParam.createChildString('default', '')
            modeGroupParam.createChildString('value', '')
            modeGroupParam.createChildNumber('stageparameter', 0.0)
            modeGroupParam.setHintString(
                str({
                    'widget': 'popup',
                    'options': ('', 'RGBA', 'RGB', 'Grayscale'),
                    'help': 'Yet another parameter to select a color mode.',
                }),
            )

        # Update state based on content layer.
        if sdfPrim is not None:
            attr = sdfPrim.getAttr('image:filepath')
            enableState = 0.0
            value = None
            if attr is not None:
                if attr.isBlocked():
                    enableState = -1.0
                else:
                    enableState = 1.0
                    value = attr.get()
            filepathGroupParam.getChild('enable').setValue(enableState, 0.0)
            if value is not None:
                filepathGroupParam.getChild('value').setValue(value, 0.0)

            attr = sdfPrim.getAttr('image:scale')
            enableState = 0.0
            value = None
            if attr is not None:
                if attr.isBlocked():
                    enableState = -1.0
                else:
                    enableState = 1.0
                    value = attr.get()
            scaleGroupParam.getChild('enable').setValue(enableState, 0.0)
            if value is not None:
                scaleGroupParam.getChild('value').setValue(value, 0.0)

            attr = sdfPrim.getAttr('image:mode')
            enableState = 0.0
            value = None
            if attr is not None:
                if attr.isBlocked():
                    enableState = -1.0
                else:
                    enableState = 1.0
                    value = attr.get()
            modeGroupParam.getChild('enable').setValue(enableState, 0.0)
            if value is not None:
                modeGroupParam.getChild('value').setValue(value, 0.0)
        else:
            # Remove the enable state from all the parameters if the prim is not valid, since they
            # will no longer be local.
            filepathGroupParam.getChild('enable').setValue(0, 0.0)
            scaleGroupParam.getChild('enable').setValue(0, 0.0)
            modeGroupParam.getChild('enable').setValue(0, 0.0)

    @classmethod
    def onParameterChanged(
        cls,
        node: NodegraphAPI.Node,
        primPath: str,
        sdfPrim: usg.Prim,
        groupParam: NodegraphAPI.Parameter,
        parameter: NodegraphAPI.Parameter,
        parameterSampleTime: float,
        usdSampleTime: float,
    ) -> bool:
        """
        Implements `ParametersHandlerPluginBase.onParameterChanged()`.

        Invoked by UsdSuperLayer-derived nodes when a parameter created previously by
        :meth:`buildOrUpdateParameters()` changes.

        Plug-ins can implement the function to react to parameter changes.

        :param node: The UsdSuperLayer-derived node.
        :param primPath: The absolute path to the prim.
        :param sdfPrim: The prim in the content layer that the plug-in can modify.
        :param groupParam: The group parameter that holds all the plug-in-defined parameters for the
            target prim.
        :param parameter: The parameter, created previously by :meth:`buildOrUpdateParameters()`,
            that has changed.
        :param parameterSampleTime: The time at which to sample the parameter to retrieve the value.
        :param usdSampleTime: The sample time at which to store the USD property. During an
            interactive session, this will typically be the defaultTimeSample. Whilst exporting or
            rendering, this will be called for all required time samples.
        :return: ``True`` if the ``sdfPrim`` was modified as a result of the parameter change;
            otherwise ``False``.
        """
        # Select the children of the StageParameterPolicy group parameter.
        parameterParts = Utils.getEnableParameterParts(parameter)
        if parameterParts is None:
            return False
        sppGroupParam, _valueParamType, enableState, valueParam, _defaultParam = parameterParts

        paramName = sppGroupParam.getName()
        attributeName = f'image:{paramName}'

        # Incoming value.
        if enableState == 0.0:
            # If no local value is set, the property is cleared
            sdfPrim.removeProperty(attributeName)
        else:
            # If a local value is set, an attribute will be created (if it does not exist yet), and
            # the value will be set.
            attr = sdfPrim.getAttr(attributeName)
            if attr is None:
                # In this example, only two types are used: either string or double.
                if valueParam.getType() == 'string':
                    attributeType = usg.Value.Type.String
                else:
                    attributeType = usg.Value.Type.Double
                attr = sdfPrim.createAttr(attributeName, attributeType)

            # Local value.
            if enableState == 1.0:
                # Here we use the parameter sample time to read the value from the parameter, and
                # the usd sample time to set the value on the attribute.
                attr.set(valueParam.getValue(parameterSampleTime), usdSampleTime)
            # Force default.
            else:
                attr.block()

        # The method must return `True` if a change has been made to the `sdfPrim` prim. In this
        # example, all code paths lead to a change.
        return True

    @classmethod
    def getIncomingLocalUsgValue(
        cls,
        node: NodegraphAPI.Node,
        prim: usg.Prim,
        groupParam: NodegraphAPI.Parameter,
        parameter: NodegraphAPI.Parameter,
    ) -> usg.Value | None:
        """
        Implements `ParametersHandlerPluginBase.getIncomingLocalUsgValue()`.

        Invoked by UsdSuperLayer-derived nodes when a parameter, created previously by
        :meth:`buildOrUpdateParameters()`, requires the local value from the incoming USD Stage.

        Plug-ins must implement this function if the parameters created in
        :meth:`buildOrUpdateParameters()` are managed by the StageParameterPolicy.

        :param node: The UsdSuperLayer-derived node.
        :param prim: The USD Prim from which to source the `usg.Value`.
        :param groupParam: The group parameter that holds all the plug-in-defined parameters for the
            target prim.
        :param parameter: The parameter; a descendant of the `groupParam` group parameter; that is
            associated to the incoming local value request. Plug-in implementations can use it to
            determine how to source or synthesize the value.
        :return: The `usg.Value` retrieved from the given USD Prim based on the parameter.
        """
        paramName = parameter.getName()
        attributeName = f'image:{paramName}'

        # The method needs to query the incoming USD Stage to retrieve the potential incoming value
        # from the target attribute.
        attr = prim.getAttr(attributeName)
        if attr is None or not attr.isValid():
            return None

        return attr.getUsgValue()


PluginRegistry = [
    (
        'UsdSuperLayerParametersHandler',
        1,
        'ExampleParametersHandlerPlugin',
        (
            UsdSuperLayerExampleParametersHandlerPlugin,
            {
                'active': True,
                'priority': 0,
                'appliedNodes': ['all'],
            },
        ),
    ),
]
