import textwrap

from Katana import (
    NodegraphAPI,
    Nodes3DAPI,
)


class CustomVariableSetNode(Nodes3DAPI.Node3D):
    """
    Custom node type that matches the built-in VariableSet node type in Katana.
    """
    def __init__(self):
        """
        Initializes an instance of the class.
        """
        super(CustomVariableSetNode, self).__init__()
        self.addInputPort('i0')
        self.addOutputPort('out')
        self.getParameters().parseXML(
            textwrap.dedent('''
            <group_parameter>
            <string_parameter name='name' value=''/>
            <string_parameter name='value' value=''/>
            </group_parameter>'''))

    def _getOpChain(self, interface):
        """
        Implementation of the :py:obj:`_getOpChain()` function that provides no Ops.
        """
        _ = interface
        # no-op

    def getInputPortAndGraphState(self, outputPort, graphState):
        """
        Overrides :py:obj:`getInputPortAndGraphState()` to extend the Local Graph
        State with one more Graph State Variable, name and value to be
        determined via the **name** and **value** parameters.
        """
        # Stack the Local Graph State, so that potential parameter expressions
        # can access it.
        with NodegraphAPI.StackedLocalGraphState(graphState):
            name = self.getParameter('name').getValue(graphState.getTime())
            value = self.getParameter('value').getValue(graphState.getTime())

        # Add a variable to the Local Graph State.
        newGraphState = graphState
        if name:
            graphStateBuilder = graphState.edit()
            graphStateBuilder.setDynamicEntry('var:' + name, value)
            newGraphState = graphStateBuilder.build()

        # Call through to the default implementation of this function with the
        # potentially modified Graph State.
        return super(CustomVariableSetNode,
                     self).getInputPortAndGraphState(outputPort, newGraphState)


NodegraphAPI.RegisterPythonNodeType('CustomVariableSet', CustomVariableSetNode)
NodegraphAPI.AddNodeFlavor("CustomVariableSet", "3d")