r"""Deprecated accessor of attribute values within a graph context.  Use og.DataView or og.Controller instead.

  _____   ______  _____   _____   ______  _____         _______  ______  _____
 |  __ \ |  ____||  __ \ |  __ \ |  ____|/ ____|    /\ |__   __||  ____||  __ \
 | |  | || |__   | |__) || |__) || |__  | |        /  \   | |   | |__   | |  | |
 | |  | ||  __|  |  ___/ |  _  / |  __| | |       / /\ \  | |   |  __|  | |  | |
 | |__| || |____ | |     | | \ \ | |____| |____  / ____ \ | |   | |____ | |__| |
 |_____/ |______||_|     |_|  \_\|______|\_____|/_/    \_\|_|   |______||_____/

"""
from contextlib import suppress
from typing import Any, Optional, Union

import omni.graph.core as og
import omni.graph.tools as ogt

from ..attribute_values import AttributeDataValueHelper, AttributeValueHelper, WrappedArrayType
from ..lookup_tables import (
    IDX_ARRAY_GET,
    IDX_ARRAY_GET_TENSOR,
    IDX_ARRAY_SET,
    IDX_GET,
    IDX_GET_TENSOR,
    IDX_SET,
    UNSUPPORTED_METHODS,
    type_access_methods,
)
from ..utils import DBG_EVAL, dbg_eval, list_dimensions

# ==============================================================================================================
# Try to import torch for tensor APIs. If not available the non-tensor APIs will still work.
try:
    import torch_wrap
except ImportError:
    torch_wrap = None


# ================================================================================
@ogt.DeprecatedClass("og.ContextHelper is deprecated after version 1.5.0. Use og.Controller instead.")
class ContextHelper:
    """Helper class for managing compute graph contexts

    Attributes:
        _py_context: Context on which to apply the operations
    """

    def __init__(self, py_context=None):
        """Remember the context for future operations.

        Args:
            py_context: Context for the operations - if None then get the current one
        """
        if py_context is None:
            for context in og.get_compute_graph_contexts():
                self._py_context = context
        else:
            self._py_context = py_context

    # ----------------------------------------------------------------------
    @property
    def context(self):
        """Returns the context being used for evaluation"""
        return self._py_context

    # ----------------------------------------------------------------------
    def get_attribute_configuration(self, attr: Union[og.Attribute, og.AttributeData]):
        """Get the array configuration information from the attribute.

        Attributes can be simple, tuples, or arrays of either. The information on what this is will be
        encoded in the attribute type name. This method decodes that type name to find out what type of
        attribute data the attribute will use.

        The method also gives the right answers for both Attribute and AttributeData with some suitable AttributeError
        catches to special case on unimplemented methods.

        Args:
            attr: Attribute whose configuration is being determined

        Returns:
            Tuple of:
                str: Name of the full/resolved data type used by the attribute (e.g. "float[3][]")
                str: Name of the simple data type used by the attribute (e.g. "float")
                bool: True if the data type is a tuple or array (e.g. "float3" or "float[]")
                bool: True if the data type is a matrix and should be flattened (e.g. "matrixd[3]" or "framed[4][]")
                List: List of lookup methods for this type

        Raises:
            TypeError: If the attribute type is not yet supported
        """
        # The actual data in extended types is the resolved type name; the attribute type is always token
        ogn_type = attr.get_resolved_type() if isinstance(attr, og.Attribute) else attr.get_type()
        type_name = ogn_type.get_ogn_type_name()

        # The root name is important for lookup
        root_type_name = ogn_type.get_base_type_name()

        if ogn_type.base_type == og.BaseDataType.UNKNOWN and (
            attr.get_extended_type()
            in (og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_ANY, og.ExtendedAttributeType.EXTENDED_ATTR_TYPE_UNION)
        ):
            raise TypeError(f"Attribute '{attr.get_name()}' is not resolved, and so has no concrete type")

        is_array_type = ogn_type.array_depth > 0 and ogn_type.role not in [og.AttributeRole.TEXT, og.AttributeRole.PATH]
        is_matrix_type = (
            type_name.startswith("matrix") or type_name.startswith("frame") or type_name.startswith("transform")
        )
        lookup_methods = UNSUPPORTED_METHODS

        # Gather nodes automatically add one level of array to their attributes
        is_gather_node = False
        with suppress(AttributeError):
            is_gather_node = attr.get_node().get_type_name() == "Gather"
            if is_gather_node:
                if is_array_type:
                    raise TypeError("Array types on Gather nodes are not yet supported in Python")
                is_array_type = True

        # Arrays of arrays are not yet supported in flatcache so they cannot be supported in Python
        if ogn_type.array_depth > 1:
            raise TypeError("Nested array types are not yet supported in Python")

        try:
            lookup_methods = type_access_methods(ogn_type)[0]
        except (KeyError, TypeError) as error:
            raise TypeError(f"Root type {ogn_type.get_ogn_type_name()} is not yet supported in Python") from error

        return (type_name, root_type_name, is_array_type, is_matrix_type, lookup_methods)

    # ----------------------------------------------------------------------
    def get(
        self, attr_object: Union[og.Attribute, og.AttributeData], on_gpu: bool = False, update_immediately: bool = False
    ):
        """
        Get the value of an attribute or attribute data. Can be used for either type.

        This method is intended for simple value retrievals. For array values use get_array().

        Args:
            attr_object: attribute data whose value is to be retrieved
            on_gpu: Is the value stored on the GPU?
            update_immediately: Should the stage update before getting the value or wait for the next regular update?

        Returns:
            Value of the attribute or attribute data

        Raises:
            TypeError: If an unsupported object type was passed in
        """
        if isinstance(attr_object, og.AttributeData):
            helper = AttributeDataValueHelper(attr_object)
            return helper.get(on_gpu)
        if isinstance(attr_object, og.Attribute):
            helper = AttributeValueHelper(attr_object)
            return helper.get(on_gpu)

        raise TypeError(f"Object type {type(attr_object)} cannot be set with this method")

    # ----------------------------------------------------------------------
    def get_array(
        self,
        attr_object: Union[og.Attribute, og.AttributeData],
        on_gpu: bool = False,
        update_immediately: bool = False,
        get_for_write: bool = False,
        reserved_element_count: int = 0,
        return_type: Optional[WrappedArrayType] = None,
    ) -> Any:
        """
        Get the value of an attribute or attribute data. Can be used for either type.

        Args:
            attr_object: attribute whose value is to be retrieved
            on_gpu: Is the value stored on the GPU?
            update_immediately: Should the stage update before getting the value or wait for the next regular update?
            get_for_write: If True then get the data in writable form (which may do other things under the covers)
            reserved_element_count: For writable array types set the element count to this before retrieving the data.
                This guarantees a buffer of this many elements is available in the returned data.

        Returns:
            Array value of the attribute or attribute data

        Raises:
            TypeError: If an unsupported object type was passed in
        """
        # Ask the attribute update its value.  If we are using a push graph, this
        # won't do anything, but in a dirty and pull graph, this generates the "pull" that causes the graph
        # to evaluate the dirtied attribute.
        with suppress(AttributeError):
            attr_object.update_object_value(update_immediately)

        helper = AttributeValueHelper(attr_object)
        if get_for_write is not None:
            if reserved_element_count is None:
                raise ValueError("Tried to get an array for write without setting the element count")
            return helper.get_array(get_for_write, reserved_element_count, on_gpu, return_type)

        return helper.get(on_gpu)

    # ----------------------------------------------------------------------
    def get_attr_value(
        self,
        attr: og.Attribute,
        isGPU: bool = False,  # noqa: N803
        isTensor: bool = False,  # noqa: N803
        getDefault: bool = False,  # noqa: N803
        updateImmediately=False,  # noqa: N803
        getForWrite=False,  # noqa: N803
        writeElemCount=0,  # noqa: N803
    ):
        """
        Get the value of a node attribute in the context managed by this class.

        Args:
            attr: attribute whose value is to be retrieved
            isGPU: Is the attribute stored on the GPU?
            isTensor: Is the attribute value a tensor type? For some types of data this doesn't mean anything
                      but is silently accepted anyway.
            isDefault: Whether or not we want to retrieve the default value
            updateImmediately: Should the stage update immediately or wait for the next regular update?
            getForWrite: Is the value going to be written to after calling?
            writeElemCount: If the value to get is a writable array then set it at this size

        Returns:
            Value of the attribute

        Raises:
            AttributeError: If you try to access array data with "isTensor=True" and torch was not imported
            TypeError: Raised if the data type of the attribute is not yet supported
        """
        (type_name, root_type_name, is_array_type, _, lookup_methods) = self.get_attribute_configuration(attr)

        # Verify that getting arrays for write has a valid element count. 0 is okay, None indicates uninitialized
        if is_array_type and getForWrite and (writeElemCount is None):
            raise ValueError(f"Attribute {attr.get_name()} requires a size to be set before getting values")

        # Ask the attribute update its value.  If we are using a push graph, this
        # won't do anything, but in a dirty and pull graph, this generates the "pull" that causes the graph
        # to evaluate the dirtied attribute.
        with suppress(AttributeError):
            attr.update_attribute_value(updateImmediately)

        can_be_tensor = (isTensor or isGPU) and is_array_type and torch_wrap is not None

        _ = DBG_EVAL and dbg_eval(
            f"Getting value of type {type_name} on {attr.get_name()}. GPU = {isGPU}, Tensor = {isTensor}"
        )
        try:
            _ = DBG_EVAL and dbg_eval("    --> Has Get/Set methods")
            if is_array_type:
                if lookup_methods[IDX_ARRAY_GET] is not None:
                    if can_be_tensor and (lookup_methods[IDX_ARRAY_GET_TENSOR] is not None):
                        _ = DBG_EVAL and dbg_eval(
                            f"    --> Returning array tensor data with {lookup_methods[IDX_ARRAY_GET_TENSOR]}"
                        )
                        tensor = lookup_methods[IDX_ARRAY_GET_TENSOR](
                            self._py_context, attr, isGPU, getForWrite, writeElemCount
                        )
                        return torch_wrap.wrap_tensor(tensor) if tensor is not None else None
                    _ = DBG_EVAL and dbg_eval(f"    --> Returning array data with {lookup_methods[IDX_ARRAY_GET]}")
                    return lookup_methods[IDX_ARRAY_GET](
                        self._py_context, attr, getDefault, getForWrite, writeElemCount
                    )
                raise TypeError(f"Getting array data of type {root_type_name} is not yet supported")

            if can_be_tensor and lookup_methods[IDX_GET_TENSOR] is not None:
                _ = DBG_EVAL and dbg_eval(f"    --> Returning tensor data with {lookup_methods[IDX_GET_TENSOR]}")
                tensor = lookup_methods[IDX_GET_TENSOR](self._py_context, attr, isGPU, getForWrite, writeElemCount)
                return torch_wrap.wrap_tensor(tensor) if tensor is not None else None

            _ = DBG_EVAL and dbg_eval(f"    --> Returning normal data with {lookup_methods[IDX_GET]}")
            return lookup_methods[IDX_GET](self._py_context, attr, getDefault, getForWrite, writeElemCount)

        except TypeError as error:
            raise TypeError(f"Trying to get unsupported attribute type {type_name}") from error

        except KeyError as error:
            raise TypeError(f"Trying to get unknown attribute type {type_name}") from error

    # ----------------------------------------------------------------------
    def set_attr_value(self, new_value, attr: og.Attribute):
        """
        Set the value of a attribute in the context managed by this class
            new_value: New value to be set on the attribute
            attr: Attribute whose value is to be set

        Raises:
            TypeError: Raised if the data type of the attribute is not yet supported
                       or if the parameters were passed in the wrong order.
        """
        if isinstance(new_value, og.Attribute):
            raise TypeError("set_attr_value is called with value first, attribute second")
        (type_name, root_type_name, is_array_type, is_matrix_type, lookup_methods) = self.get_attribute_configuration(
            attr
        )
        # Flatten the matrix data out, if it wasn't already flattened
        # 2x2 matrix can come in as [[a, b], [c, d]] or [a, b, c, d] but will always go to the ABI as the latter
        if is_matrix_type:
            value_dimensions = list_dimensions(new_value)
            if is_array_type:
                if value_dimensions == 3:
                    value_to_set = [[item for sublist in matrix for item in sublist] for matrix in new_value]
                else:
                    value_to_set = new_value
            else:
                if value_dimensions == 2:
                    value_to_set = [item for sublist in new_value for item in sublist]
                else:
                    value_to_set = new_value
        else:
            value_to_set = new_value
        try:
            _ = DBG_EVAL and dbg_eval(f"Set attribute {attr.get_name()} of type {type_name} to {new_value}")
            _ = DBG_EVAL and dbg_eval(f"    as {root_type_name}, array={is_array_type}, matrix={is_matrix_type}")
            if is_array_type:
                if lookup_methods[IDX_ARRAY_GET] and lookup_methods[IDX_ARRAY_SET]:
                    _ = DBG_EVAL and dbg_eval(f"    --> Setting array node data using {lookup_methods[IDX_ARRAY_SET]}")
                    # When setting an entire array all at once you need to first ensure the memory is allocated
                    # by calling the get method.
                    lookup_methods[IDX_ARRAY_GET](self._py_context, attr, False, True, len(value_to_set))
                    lookup_methods[IDX_ARRAY_SET](self._py_context, value_to_set, attr)
                else:
                    raise TypeError(f"Setting array data of type {type_name} is not yet supported")
            else:
                _ = DBG_EVAL and dbg_eval(f"    --> Setting simple data with {lookup_methods[IDX_SET]}")
                lookup_methods[IDX_SET](self._py_context, value_to_set, attr)
        except KeyError as error:
            raise TypeError(f"Trying to set unsupported attribute type {type_name}") from error

    # ----------------------------------------------------------------------
    def get_elem_count(self, attr):
        """
        Get the number of elements on the attribute in the context managed by this class
            attr: Attribute whose element values are to be counted

        Returns:
            Number of elements currently on the attribute
        """
        return self._py_context.get_elem_count(attr)
