plugin

class BlenderCreator(Creator):
    """Base class for Blender Creator plug-ins."""
    defaults = ['Main']

    settings_category = "blender"
    create_as_asset_group = False

    @staticmethod
    def cache_instance_data(shared_data):
        """Cache instances for Creators shared data.

        Create `blender_cached_instances` key when needed in shared data and
        fill it with all collected instances from the scene under its
        respective creator identifiers.

        If legacy instances are detected in the scene, create
        `blender_cached_legacy_instances` key and fill it with
        all legacy products from this family as a value.  # key or value?

        Args:
            shared_data(Dict[str, Any]): Shared data.

        """
        if not shared_data.get('blender_cached_instances'):
            cache = {}
            cache_legacy = {}
            convert_avalon_instances()
            ayon_instances = bpy.data.collections.get(AYON_INSTANCES)
            ayon_instance_objs = (
                ayon_instances.objects if ayon_instances else []
            )

            # Consider any node tree objects as well
            node_tree_objects = []
            if bpy.context.scene.node_tree:
                node_tree_objects = bpy.context.scene.node_tree.nodes

            for obj_or_col in itertools.chain(
                    ayon_instance_objs,
                    bpy.data.collections,
                    node_tree_objects
            ):
                ayon_prop = get_ayon_property(obj_or_col)
                if not ayon_prop:
                    continue

                if ayon_prop.get('id') not in {
                    AYON_INSTANCE_ID, AVALON_INSTANCE_ID
                }:
                    continue

                creator_id = ayon_prop.get("creator_identifier")
                if "openpype" in creator_id:
                    creator_id = creator_id.replace("openpype", "ayon")
                    ayon_prop["creator_identifier"] = creator_id

                if creator_id:
                    # Creator instance
                    cache.setdefault(creator_id, []).append(obj_or_col)
                else:
                    # Legacy creator instance
                    family = ayon_prop.get('family')
                    cache_legacy.setdefault(family, []).append(obj_or_col)

            shared_data["blender_cached_instances"] = cache
            shared_data["blender_cached_legacy_instances"] = cache_legacy
        return shared_data

    def create(
        self, product_name: str, instance_data: dict, pre_create_data: dict
    ):
        """Override abstract method from Creator.
        Create new instance and store it.

        Args:
            product_name (str): Product name of created instance.
            instance_data (dict): Instance base data.
            pre_create_data (dict): Data based on pre creation attributes.
                Those may affect how creator works.
        """
        # Get Instance Container or create it if it does not exist
        ayon_instances = self._ensure_ayon_instances_collection()

        # Create asset group
        folder_name = instance_data["folderPath"].split("/")[-1]

        name = prepare_scene_name(folder_name, product_name)
        if self.create_as_asset_group:
            # Create instance as empty
            instance_node = bpy.data.objects.new(name=name, object_data=None)
            instance_node.empty_display_type = 'SINGLE_ARROW'
            ayon_instances.objects.link(instance_node)
        else:
            # Create instance collection
            instance_node = bpy.data.collections.new(name=name)
            ayon_instances.children.link(instance_node)

        self.set_instance_data(product_name, instance_data)

        instance = CreatedInstance(
            self.product_type, product_name, instance_data, self,
            transient_data={"instance_node": instance_node}
        )
        self._add_instance_to_context(instance)

        self.imprint(instance_node, instance_data)

        return instance_node

    def collect_instances(self):
        """Override abstract method from BlenderCreator.
        Collect existing instances related to this creator plugin."""

        # Cache instances in shared data
        self.cache_instance_data(self.collection_shared_data)

        # Get cached instances
        cached_instances = self.collection_shared_data.get(
            "blender_cached_instances"
        )
        if not cached_instances:
            return

        # Process only instances that were created by this creator
        for instance_node in cached_instances.get(self.identifier, []):
            instance_data = self.read(instance_node)

            # Create instance object from existing data
            instance = CreatedInstance.from_existing(
                instance_data=instance_data,
                creator=self,
                transient_data={"instance_node": instance_node}
            )

            # Add instance to create context
            self._add_instance_to_context(instance)

    def update_instances(self, update_list):
        """Override abstract method from BlenderCreator.
        Store changes of existing instances so they can be recollected.

        Args:
            update_list(List[UpdateData]): Changed instances
                and their changes, as a list of tuples.
        """

        for created_instance, changes in update_list:
            data = created_instance.data_to_store()
            node = created_instance.transient_data["instance_node"]
            if not node:
                # We can't update if we don't know the node
                self.log.error(
                    f"Unable to update instance {created_instance} "
                    f"without instance node."
                )
                return

            # Rename the instance node in the scene if product
            #   or folder changed.
            # Do not rename the instance if the family is workfile, as the
            # workfile instance is included in the AYON_CONTAINER collection.
            if (
                "productName" in changes.changed_keys
                or "folderPath" in changes.changed_keys
            ) and created_instance.product_type != "workfile":
                folder_name = data["folderPath"].split("/")[-1]
                name = prepare_scene_name(
                    folder_name, data["productName"]
                )
                node.name = name

            self.imprint(node, data)

    def remove_instances(self, instances: List[CreatedInstance]):
        for instance in instances:
            node = instance.transient_data["instance_node"]

            # Remove collection node and its children
            if isinstance(node, bpy.types.Collection):
                # Remove recursively linked child collections and objects
                for child in node.children_recursive:
                    if isinstance(child, bpy.types.Collection):
                        # If only linked to this collection, relink to scene before removal
                        if len(child.users_collection) == 1:
                            if child.name not in bpy.context.scene.collection:
                                bpy.context.scene.collection.children.link(child)

                    elif isinstance(child, bpy.types.Object):
                        if len(child.users_collection) == 1:
                            if child.name not in bpy.context.scene.collection:
                                bpy.context.scene.collection.objects.link(child)
                # Remove directly linked objects
                for obj in node.objects:
                    if len(obj.users_collection) == 1:
                        if obj.name not in bpy.context.scene.collection.objects:
                            bpy.context.scene.collection.objects.link(obj)

                bpy.data.collections.remove(node)

            # Remove object node
            elif isinstance(node, bpy.types.Object):
                bpy.data.objects.remove(node)

            # Remove compositor node
            elif isinstance(node, bpy.types.CompositorNode):
                bpy.context.scene.node_tree.nodes.remove(node)

            self._remove_instance_from_context(instance)

    def set_instance_data(
        self,
        product_name: str,
        instance_data: dict
    ):
        """Fill instance data with required items.

        Args:
            product_name (str): Product name of created instance.
            instance_data (dict): Instance base data.
            instance_node (bpy.types.ID): Instance node in blender scene.
        """
        if not instance_data:
            instance_data = {}

        instance_data.update(
            {
                "id": AYON_INSTANCE_ID,
                "creator_identifier": self.identifier,
                "productName": product_name,
            }
        )

    def _ensure_ayon_instances_collection(self) -> bpy.types.Collection:
        """Create AYON Instances collections that contains created instances.

        Returns:
            bpy.types.Collection: AYON Instances collection
        """
        node = bpy.data.collections.get(AYON_INSTANCES)
        if node:
            # Already exists, return it
            return node

        node = bpy.data.collections.new(AYON_INSTANCES)
        node.color_tag = "COLOR_04"
        node.use_fake_user = True
        bpy.context.scene.collection.children.link(node)
        return node

    def get_pre_create_attr_defs(self):
        return [
            BoolDef("use_selection",
                    label="Use selection",
                    default=True)
        ]

    def imprint(self, node, data: dict):
        """Imprint data to the instance node.

        This can be overridden by a sub-class to store certain data of the
        instance in a different way, e.g. in a custom property or the 'mute'
        state of a node.
        """
        imprint(node, data)

    def read(self, node) -> dict:
        """Read data from the instance node.

        The `data` dictionary is the readily stored
        """
        ayon_property = node.get(AYON_PROPERTY)
        if not ayon_property:
            return {}

        return ayon_property.to_dict()

class BlenderLoader(LoaderPlugin):
    """A basic AssetLoader for Blender

    This will implement the basic logic for linking/appending assets
    into another Blender scene.

    The `update` method should be implemented by a subclass, because
    it's different for different types (e.g. model, rig, animation,
    etc.).
    """
    settings_category = "blender"

    @staticmethod
    def _get_instance_empty(instance_name: str, nodes: List) -> Optional[bpy.types.Object]:
        """Get the 'instance empty' that holds the collection instance."""
        for node in nodes:
            if not isinstance(node, bpy.types.Object):
                continue
            if (node.type == 'EMPTY' and node.instance_type == 'COLLECTION'
                    and node.instance_collection and node.name == instance_name):
                return node
        return None

    @staticmethod
    def _get_instance_collection(instance_name: str, nodes: List) -> Optional[bpy.types.Collection]:
        """Get the 'instance collection' (container) for this asset."""
        for node in nodes:
            if not isinstance(node, bpy.types.Collection):
                continue
            if node.name == instance_name:
                return node
        return None

    @staticmethod
    def _get_library_from_container(container: bpy.types.Collection) -> bpy.types.Library:
        """Find the library file from the container.

        It traverses the objects from this collection, checks if there is only
        1 library from which the objects come from and returns the library.

        Warning:
            No nested collections are supported at the moment!
        """
        assert not container.children, "Nested collections are not supported."
        assert container.objects, "The collection doesn't contain any objects."
        libraries = set()
        for obj in container.objects:
            assert obj.library, f"'{obj.name}' is not linked."
            libraries.add(obj.library)

        assert len(
            libraries) == 1, "'{container.name}' contains objects from more then 1 library."

        return list(libraries)[0]

    def process_asset(self,
                      context: dict,
                      name: str,
                      namespace: Optional[str] = None,
                      options: Optional[Dict] = None):
        """Must be implemented by a subclass"""
        raise NotImplementedError("Must be implemented by a subclass")

    def load(self,
             context: dict,
             name: Optional[str] = None,
             namespace: Optional[str] = None,
             options: Optional[Dict] = None) -> Optional[bpy.types.Collection]:
        """ Run the loader on Blender main thread"""
        mti = MainThreadItem(self._load, context, name, namespace, options)
        execute_in_main_thread(mti)

    def _load(self,
              context: dict,
              name: Optional[str] = None,
              namespace: Optional[str] = None,
              options: Optional[Dict] = None
    ) -> Optional[bpy.types.Collection]:
        """Load asset via database

        Arguments:
            context: Full parenthood of representation to load
            name: Use pre-defined name
            namespace: Use pre-defined namespace
            options: Additional settings dictionary
        """
        # TODO: make it possible to add the asset several times by
        # just re-using the collection
        filepath = self.filepath_from_context(context)
        assert Path(filepath).exists(), f"{filepath} doesn't exist."

        folder_name = context["folder"]["name"]
        product_name = context["product"]["name"]
        unique_number = get_unique_number(
            folder_name, product_name
        )
        namespace = namespace or f"{folder_name}_{unique_number}"
        name = name or prepare_scene_name(
            folder_name, product_name, unique_number
        )

        nodes = self.process_asset(
            context=context,
            name=name,
            namespace=namespace,
            options=options,
        )

        # Only containerise if anything was loaded by the Loader.
        if not nodes:
            return None

        # Only containerise if it's not already a collection from a .blend file.
        # representation = context["representation"]["name"]
        # if representation != "blend":
        #     from ayon_blender.api.pipeline import containerise
        #     return containerise(
        #         name=name,
        #         namespace=namespace,
        #         nodes=nodes,
        #         context=context,
        #         loader=self.__class__.__name__,
        #     )

        # folder_name = context["folder"]["name"]
        # product_name = context["product"]["name"]
        # instance_name = prepare_scene_name(
        #     folder_name, product_name, unique_number
        # ) + '_CON'

        # return self._get_instance_collection(instance_name, nodes)

    def exec_update(self, container: Dict, context: Dict):
        """Must be implemented by a subclass"""
        raise NotImplementedError("Must be implemented by a subclass")

    def update(self, container: Dict, context: Dict):
        """ Run the update on Blender main thread"""
        mti = MainThreadItem(self.exec_update, container, context)
        execute_in_main_thread(mti)

    def exec_remove(self, container: Dict) -> bool:
        """Must be implemented by a subclass"""
        raise NotImplementedError("Must be implemented by a subclass")

    def remove(self, container: Dict) -> bool:
        """ Run the remove on Blender main thread"""
        mti = MainThreadItem(self.exec_remove, container)
        execute_in_main_thread(mti)

    def switch(self, container, context):
        # Support switch in scene inventory
        self.update(container, context)