"""
Support for handling attributes of type "bundle" - i.e. attributes whose job it is to encapsulate arbitrary
collections of other attributes, including other bundle attributes.
"""
from typing import List

from ..keys import CudaPointerValues, MemoryTypeValues
from ..utils import IndentedOutput, ParseError, ensure_quoted, to_usd_docs
from .AttributeManager import AttributeManager, CppConfiguration
from .naming import INPUT_GROUP, OUTPUT_GROUP, STATE_GROUP


class BundleAttributeManager(AttributeManager):
    """
    Support class for attributes of type attribute bundle.
    This type of attribute is more complex than standard attributes since it has many more features to handle.
    """

    OGN_TYPE = "bundle"

    # Hardcoded type configuration using USD. Dictionary key is the OGN attribute type name (without arrays).
    CPP_CONFIGURATION = {
        # Type information is overridden but the include file is important to specialize
        "bundle": CppConfiguration(None, include_files=["omni/graph/core/ogn/UsdTypes.h"])
    }

    def requires_default(self):
        """Bundles never need default values as nothing other than an empty bundle makes sense"""
        return False

    @staticmethod
    def array_depths_supported() -> List[int]:
        """Bundle arrays are not yet supported in Fabric"""
        return [0]

    def memory_storage(self) -> str:
        """Bundle handles will always be stored on the CPU as that is where Fabric forces them"""
        return MemoryTypeValues.CPU

    def cpp_base_type_name(self):
        """This type name switches based on read status so this has to override the default method"""
        return "ConstBundleHandle" if self.is_read_only() else "BundleHandle"

    def cpp_element_type_name(self) -> str:
        """The configuration is all manual here so override the default method"""
        return self.cpp_base_type_name()

    def cpp_default_initializer(self):
        """The bundle doesn't really have a default possible so initialize it to an invalid handle or empty array."""
        return "nullptr, 0" if self.array_depth > 0 else "BundleHandle::invalidValue()"

    def cpp_includes(self) -> List[str]:
        """Tack on the include implementing the bundle wrappers"""
        includes = super().cpp_includes()
        includes.append("omni/graph/core/ogn/Bundle.h")
        return includes

    def cpp_accessor_on_cpu(self) -> bool:
        """Bundle wrappers provide a type-casting accessor that will always live on the CPU"""
        return True

    def cpp_wrapper_class(self) -> str:
        """Returns the bundle-specific wrapper class name used to access attribute data in the C++ database"""
        template_args = [self.attribute_group, MemoryTypeValues.CPP[self.memory_type]]
        if self.cuda_pointer_type is not None:
            template_args.append(CudaPointerValues.CPP[self.cuda_pointer_type])
        return ("ogn::BundleAttribute", template_args)

    def cpp_set_handle_at_runtime(self) -> bool:
        """Bundle types do not use attribute handles"""
        return False

    def validate_value(self, value):
        """Raises a ParseError if value is not a valid bundle value"""
        raise ParseError("Bundles do not have values")

    def cuda_includes(self) -> List[str]:
        """The bundle data is the same type of data in CUDA as it is in C++"""
        includes = super().cuda_includes()
        includes.append("omni/graph/core/Handle.h")
        return includes

    def cuda_base_type_name(self) -> str:
        """Returns a string with the CUDA base type of the attribute data"""
        return "omni::graph::core::ConstBundleHandle" if self.is_read_only() else "omni::graph::core::BundleHandle"

    def cuda_element_type_name(self) -> str:
        """The configuration is all manual here so override the default method"""
        return self.cuda_base_type_name()

    def fabric_data_variable_name(self) -> str:
        """Returns a string containing the generated name of the Fabric pointer for this attribute."""
        return f"{self.cpp_variable_name()}.m_bundleHandle"

    def fabric_pointer_exists(self) -> str:
        """Return a string that checks for the existence of the Fabric pointer variable value"""
        return f"{self.fabric_data_variable_name()}.isValid()"

    def python_role_name(self) -> str:
        """Returns a string with the Python role name for this attribute"""
        return "og.Database.ROLE_BUNDLE"

    def usd_name(self) -> str:
        """Bundled output and state attributes are represented as virtual prims in USD so they cannot have namespaces"""
        # Output and state bundles are prims so they have to follow the prim naming restrictions (no colons allowed)
        if self.attribute_group == INPUT_GROUP:
            return self.name
        return self.name.replace(":", "_")

    def create_type_name(self) -> str:
        """Bundled attributes have a special name when creating so that they can be instantiated differently"""
        return "bundle"

    # ----------------------------------------------------------------------
    def generate_python_property_code(self, out: IndentedOutput):
        """Emits the generated code implementing a property for this bundle attribute.
        This class overrides the default behaviour because it needs a wrapper class to access the internal
        functionality of the bundle.
        """
        property_name = self.python_property_name()
        out.write()
        out.write("@property")
        if out.indent(f"def {property_name}(self) -> og.BundleContents:"):
            out.write(f'"""Get the bundle wrapper class for the attribute {self.namespace}.{property_name}"""')
            out.write(f"return self.__bundles.{property_name}")
            out.exdent()

        # No setters at all for read only bundles
        if self.is_read_only():
            return
        property_name = self.python_property_name()
        out.write()
        out.write(f"@{property_name}.setter")
        if out.indent(f"def {property_name}(self, bundle: og.BundleContents):"):
            out.write(f'"""Overwrite the bundle attribute {self.namespace}.{property_name} with a new bundle"""')
            if out.indent("if not isinstance(bundle, og.BundleContents):"):
                out.write('carb.log_error("Only bundle attributes can be assigned to another bundle attribute")')
                out.exdent()
            out.write(f"self.__bundles.{property_name}.bundle = bundle")
            out.exdent()

    # ----------------------------------------------------------------------
    def python_imports(self) -> List[str]:
        """Return a list of modules to import in the Python header for proper parsing of this type"""
        return super().python_imports() + ["import carb"]

    # ----------------------------------------------------------------------
    def emit_usd_declaration(self, out: IndentedOutput):
        """Print a declaration for this attribute in USD

        Args:
            out: Output handler where the USD will be emitted
        """
        try:
            usd_name = self.usd_name()
        except ParseError:
            # Attributes without USD representations can be skipped
            return

        docs = to_usd_docs(self.description)

        # Bundle attributes are stored as virtual prims, not actual USD attributes, so they require a
        # different type of declaration. Inputs are relationships, outputs are defined as nested prims.
        if self.attribute_group == INPUT_GROUP:
            if out.indent(f"custom rel {usd_name} ("):
                out.write(docs)
                out.exdent(")")
        else:
            # Prim names cannot have colons in them
            usd_name = usd_name.replace(":", "_")
            if self.attribute_group == OUTPUT_GROUP:
                side = "Output"
            elif self.attribute_group == STATE_GROUP:
                side = "State"
            else:
                side = "Unknown"
            if out.indent(f"def {side} {ensure_quoted(usd_name)} ("):
                out.write(docs)
                out.exdent(")")
            out.write("{")
            out.write("}")
