"""Support for simplified access to data on nodes of type omni.graph.test.TestDataModel

Helper node that allows to test that core features of datamodel are working as expected (CoW, DataStealing, ...)
"""

import carb
import numpy

import carb
import omni.graph.core as og
import omni.graph.core._omni_graph_core as _og
import omni.graph.tools.ogn as ogn



class OgnTestDataModelDatabase(og.Database):
    """Helper class providing simplified access to data on nodes of type omni.graph.test.TestDataModel

    Class Members:
        node: Node being evaluated

    Attribute Value Properties:
        Inputs:
            inputs.arrayShouldMatch
            inputs.attrib
            inputs.bundleArraysThatShouldDiffer
            inputs.bundleShouldMatch
            inputs.mutArray
            inputs.mutBundle
            inputs.mutateArray
            inputs.refArray
            inputs.refBundle
        Outputs:
            outputs.array
            outputs.bundle
    """

    # Imprint the generator and target ABI versions in the file for JIT generation
    GENERATOR_VERSION = (1, 41, 3)
    TARGET_VERSION = (2, 139, 12)

    # This is an internal object that provides per-class storage of a per-node data dictionary
    PER_NODE_DATA = {}

    # This is an internal object that describes unchanging attributes in a generic way
    # The values in this list are in no particular order, as a per-attribute tuple
    #     Name, Type, ExtendedTypeIndex, UiName, Description, Metadata,
    #     Is_Required, DefaultValue, Is_Deprecated, DeprecationMsg
    # You should not need to access any of this data directly, use the defined database interfaces
    INTERFACE = og.Database._get_interface([
        ('inputs:arrayShouldMatch', 'bool', 0, 'Array should match', 'Whether or not the input arrays should be the same one one', {ogn.MetadataKeys.DEFAULT: 'true'}, True, True, False, ''),
        ('inputs:attrib', 'token', 0, 'Attrib to mutate', 'Attribute to mutate in the bundle', {}, True, "", False, ''),
        ('inputs:bundleArraysThatShouldDiffer', 'int', 0, 'Number of != arrays in bundles', 'The number of arrays attribute in the input bundles that should differs', {ogn.MetadataKeys.DEFAULT: '0'}, True, 0, False, ''),
        ('inputs:bundleShouldMatch', 'bool', 0, 'Bundles should match', 'Whether or not the input bundles should be the same one', {ogn.MetadataKeys.DEFAULT: 'true'}, True, True, False, ''),
        ('inputs:mutArray', 'point3f[]', 0, 'In array', 'Array meant to be mutated', {}, True, [], False, ''),
        ('inputs:mutBundle', 'bundle', 0, 'In bundle', 'Bundle meant to be mutated (or not)', {}, True, None, False, ''),
        ('inputs:mutateArray', 'bool', 0, 'Mutate array', 'Whether or not to mutate the array or just passthrough', {ogn.MetadataKeys.DEFAULT: 'false'}, True, False, False, ''),
        ('inputs:refArray', 'point3f[]', 0, 'Ref array', 'A reference array used as a point of comparaison', {}, True, [], False, ''),
        ('inputs:refBundle', 'bundle', 0, 'Ref bundle', 'Reference Bundle used as a point of comparaison', {}, True, None, False, ''),
        ('outputs:array', 'point3f[]', 0, 'Output array', 'The outputed array', {}, True, None, False, ''),
        ('outputs:bundle', 'bundle', 0, 'Output bundle', 'The outputed bundle', {}, True, None, False, ''),
    ])

    @classmethod
    def _populate_role_data(cls):
        """Populate a role structure with the non-default roles on this node type"""
        role_data = super()._populate_role_data()
        role_data.inputs.mutArray = og.AttributeRole.POSITION
        role_data.inputs.mutBundle = og.AttributeRole.BUNDLE
        role_data.inputs.refArray = og.AttributeRole.POSITION
        role_data.inputs.refBundle = og.AttributeRole.BUNDLE
        role_data.outputs.array = og.AttributeRole.POSITION
        role_data.outputs.bundle = og.AttributeRole.BUNDLE
        return role_data

    class ValuesForInputs(og.DynamicAttributeAccess):
        LOCAL_PROPERTY_NAMES = { }
        """Helper class that creates natural hierarchical access to input attributes"""
        def __init__(self, node: og.Node, attributes, dynamic_attributes: og.DynamicAttributeInterface):
            """Initialize simplified access for the attribute data"""
            context = node.get_graph().get_default_graph_context()
            super().__init__(context, node, attributes, dynamic_attributes)
            self.__bundles = og.BundleContainer(context, node, attributes, [], read_only=True, gpu_ptr_kinds={})
            self._batchedReadAttributes = []
            self._batchedReadValues = []

        @property
        def arrayShouldMatch(self):
            data_view = og.AttributeValueHelper(self._attributes.arrayShouldMatch)
            return data_view.get()

        @arrayShouldMatch.setter
        def arrayShouldMatch(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.arrayShouldMatch)
            data_view = og.AttributeValueHelper(self._attributes.arrayShouldMatch)
            data_view.set(value)

        @property
        def attrib(self):
            data_view = og.AttributeValueHelper(self._attributes.attrib)
            return data_view.get()

        @attrib.setter
        def attrib(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.attrib)
            data_view = og.AttributeValueHelper(self._attributes.attrib)
            data_view.set(value)

        @property
        def bundleArraysThatShouldDiffer(self):
            data_view = og.AttributeValueHelper(self._attributes.bundleArraysThatShouldDiffer)
            return data_view.get()

        @bundleArraysThatShouldDiffer.setter
        def bundleArraysThatShouldDiffer(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.bundleArraysThatShouldDiffer)
            data_view = og.AttributeValueHelper(self._attributes.bundleArraysThatShouldDiffer)
            data_view.set(value)

        @property
        def bundleShouldMatch(self):
            data_view = og.AttributeValueHelper(self._attributes.bundleShouldMatch)
            return data_view.get()

        @bundleShouldMatch.setter
        def bundleShouldMatch(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.bundleShouldMatch)
            data_view = og.AttributeValueHelper(self._attributes.bundleShouldMatch)
            data_view.set(value)

        @property
        def mutArray(self):
            data_view = og.AttributeValueHelper(self._attributes.mutArray)
            return data_view.get()

        @mutArray.setter
        def mutArray(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.mutArray)
            data_view = og.AttributeValueHelper(self._attributes.mutArray)
            data_view.set(value)
            self.mutArray_size = data_view.get_array_size()

        @property
        def mutBundle(self) -> og.BundleContents:
            """Get the bundle wrapper class for the attribute inputs.mutBundle"""
            return self.__bundles.mutBundle

        @property
        def mutateArray(self):
            data_view = og.AttributeValueHelper(self._attributes.mutateArray)
            return data_view.get()

        @mutateArray.setter
        def mutateArray(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.mutateArray)
            data_view = og.AttributeValueHelper(self._attributes.mutateArray)
            data_view.set(value)

        @property
        def refArray(self):
            data_view = og.AttributeValueHelper(self._attributes.refArray)
            return data_view.get()

        @refArray.setter
        def refArray(self, value):
            if self._setting_locked:
                raise og.ReadOnlyError(self._attributes.refArray)
            data_view = og.AttributeValueHelper(self._attributes.refArray)
            data_view.set(value)
            self.refArray_size = data_view.get_array_size()

        @property
        def refBundle(self) -> og.BundleContents:
            """Get the bundle wrapper class for the attribute inputs.refBundle"""
            return self.__bundles.refBundle

        def _prefetch(self):
            readAttributes = self._batchedReadAttributes
            newValues = _og._prefetch_input_attributes_data(readAttributes)
            if len(readAttributes) == len(newValues):
                self._batchedReadValues = newValues

    class ValuesForOutputs(og.DynamicAttributeAccess):
        LOCAL_PROPERTY_NAMES = { }
        """Helper class that creates natural hierarchical access to output attributes"""
        def __init__(self, node: og.Node, attributes, dynamic_attributes: og.DynamicAttributeInterface):
            """Initialize simplified access for the attribute data"""
            context = node.get_graph().get_default_graph_context()
            super().__init__(context, node, attributes, dynamic_attributes)
            self.__bundles = og.BundleContainer(context, node, attributes, [], read_only=False, gpu_ptr_kinds={})
            self.array_size = None
            self._batchedWriteValues = { }

        @property
        def array(self):
            data_view = og.AttributeValueHelper(self._attributes.array)
            return data_view.get(reserved_element_count=self.array_size)

        @array.setter
        def array(self, value):
            data_view = og.AttributeValueHelper(self._attributes.array)
            data_view.set(value)
            self.array_size = data_view.get_array_size()

        @property
        def bundle(self) -> og.BundleContents:
            """Get the bundle wrapper class for the attribute outputs.bundle"""
            return self.__bundles.bundle

        @bundle.setter
        def bundle(self, bundle: og.BundleContents):
            """Overwrite the bundle attribute outputs.bundle with a new bundle"""
            if not isinstance(bundle, og.BundleContents):
                carb.log_error("Only bundle attributes can be assigned to another bundle attribute")
            self.__bundles.bundle.bundle = bundle

        def _commit(self):
            _og._commit_output_attributes_data(self._batchedWriteValues)
            self._batchedWriteValues = { }

    class ValuesForState(og.DynamicAttributeAccess):
        """Helper class that creates natural hierarchical access to state attributes"""
        def __init__(self, node: og.Node, attributes, dynamic_attributes: og.DynamicAttributeInterface):
            """Initialize simplified access for the attribute data"""
            context = node.get_graph().get_default_graph_context()
            super().__init__(context, node, attributes, dynamic_attributes)

    def __init__(self, node):
        super().__init__(node)
        dynamic_attributes = self.dynamic_attribute_data(node, og.AttributePortType.ATTRIBUTE_PORT_TYPE_INPUT)
        self.inputs = OgnTestDataModelDatabase.ValuesForInputs(node, self.attributes.inputs, dynamic_attributes)
        dynamic_attributes = self.dynamic_attribute_data(node, og.AttributePortType.ATTRIBUTE_PORT_TYPE_OUTPUT)
        self.outputs = OgnTestDataModelDatabase.ValuesForOutputs(node, self.attributes.outputs, dynamic_attributes)
        dynamic_attributes = self.dynamic_attribute_data(node, og.AttributePortType.ATTRIBUTE_PORT_TYPE_STATE)
        self.state = OgnTestDataModelDatabase.ValuesForState(node, self.attributes.state, dynamic_attributes)
