"""Load a rig asset in Blender."""

import logging
from pathlib import Path
from pprint import pformat
from typing import Dict, List, Optional

from avalon import api, blender
import bpy
import pype.hosts.blender.plugin as plugin


class BlendRigLoader(plugin.AssetLoader):
    """Load rigs from a .blend file.

    Because they come from a .blend file we can simply link the collection that
    contains the model. There is no further need to 'containerise' it.
    """

    families = ["rig"]
    representations = ["blend"]

    label = "Link Rig"
    icon = "code-fork"
    color = "orange"

    def _remove(self, objects, obj_container):
        for obj in list(objects):
            if obj.type == 'ARMATURE':
                bpy.data.armatures.remove(obj.data)
            elif obj.type == 'MESH':
                bpy.data.meshes.remove(obj.data)
            elif obj.type == 'CURVE':
                bpy.data.curves.remove(obj.data)

        for child in obj_container.children:
            bpy.data.collections.remove(child)

        bpy.data.collections.remove(obj_container)

    def make_local_and_metadata(self, obj, collection_name):
        local_obj = plugin.prepare_data(obj, collection_name)
        plugin.prepare_data(local_obj.data, collection_name)

        if not local_obj.get(blender.pipeline.AVALON_PROPERTY):
            local_obj[blender.pipeline.AVALON_PROPERTY] = dict()

        avalon_info = local_obj[blender.pipeline.AVALON_PROPERTY]
        avalon_info.update({"container_name": collection_name + '_CON'})

        return local_obj

    def _process(
        self, libpath, lib_container, collection_name,
        action, parent_collection
    ):
        relative = bpy.context.preferences.filepaths.use_relative_paths
        with bpy.data.libraries.load(
            libpath, link=True, relative=relative
        ) as (_, data_to):
            data_to.collections = [lib_container]

        parent = parent_collection

        if parent is None:
            parent = bpy.context.scene.collection

        parent.children.link(bpy.data.collections[lib_container])

        rig_container = parent.children[lib_container].make_local()
        rig_container.name = collection_name

        objects = []
        armatures = [
            obj for obj in rig_container.objects
            if obj.type == 'ARMATURE'
        ]

        for child in rig_container.children:
            local_child = plugin.prepare_data(child, collection_name)
            objects.extend(local_child.objects)

        # for obj in bpy.data.objects:
        #     obj.select_set(False)

        constraints = []

        for armature in armatures:
            for bone in armature.pose.bones:
                for constraint in bone.constraints:
                    if hasattr(constraint, 'target'):
                        constraints.append(constraint)

        # Link armatures after other objects.
        # The armature is unparented for all the non-local meshes,
        # when it is made local.
        for obj in objects:
            local_obj = self.make_local_and_metadata(obj, collection_name)

            if obj != local_obj:
                for constraint in constraints:
                    if constraint.target == obj:
                        constraint.target = local_obj

        for armature in armatures:
            local_obj = self.make_local_and_metadata(armature, collection_name)

            if action is not None:
                local_obj.animation_data.action = action

            # Set link the drivers to the local object
            if local_obj.data.animation_data:
                for d in local_obj.data.animation_data.drivers:
                    for v in d.driver.variables:
                        for t in v.targets:
                            t.id = local_obj

        rig_container.pop(blender.pipeline.AVALON_PROPERTY)

        bpy.ops.object.select_all(action='DESELECT')

        return rig_container

    def process_asset(
        self, context: dict, name: str, namespace: Optional[str] = None,
        options: Optional[Dict] = None
    ) -> Optional[List]:
        """
        Arguments:
            name: Use pre-defined name
            namespace: Use pre-defined namespace
            context: Full parenthood of representation to load
            options: Additional settings dictionary
        """
        libpath = self.fname
        asset = context["asset"]["name"]
        subset = context["subset"]["name"]
        lib_container = plugin.asset_name(
            asset, subset
        )
        unique_number = plugin.get_unique_number(
            asset, subset
        )
        namespace = namespace or f"{asset}_{unique_number}"
        collection_name = plugin.asset_name(
            asset, subset, unique_number
        )

        container = bpy.data.collections.new(collection_name)
        blender.pipeline.containerise_existing(
            container,
            name,
            namespace,
            context,
            self.__class__.__name__,
        )

        container_metadata = container.get(
            blender.pipeline.AVALON_PROPERTY)

        container_metadata["libpath"] = libpath
        container_metadata["lib_container"] = lib_container

        obj_container = self._process(
            libpath, lib_container, collection_name, None, None)

        container_metadata["obj_container"] = obj_container
        # Save the list of objects in the metadata container
        container_metadata["objects"] = obj_container.all_objects

        nodes = list(container.objects)
        nodes.append(container)
        self[:] = nodes
        return nodes

    def update(self, container: Dict, representation: Dict):
        """Update the loaded asset.

        This will remove all objects of the current collection, load the new
        ones and add them to the collection.
        If the objects of the collection are used in another collection they
        will not be removed, only unlinked. Normally this should not be the
        case though.

        Warning:
            No nested collections are supported at the moment!
        """
        collection = bpy.data.collections.get(
            container["objectName"]
        )
        libpath = Path(api.get_representation_path(representation))
        extension = libpath.suffix.lower()

        self.log.info(
            "Container: %s\nRepresentation: %s",
            pformat(container, indent=2),
            pformat(representation, indent=2),
        )

        assert collection, (
            f"The asset is not loaded: {container['objectName']}"
        )
        assert not (collection.children), (
            "Nested collections are not supported."
        )
        assert libpath, (
            "No existing library file found for {container['objectName']}"
        )
        assert libpath.is_file(), (
            f"The file doesn't exist: {libpath}"
        )
        assert extension in plugin.VALID_EXTENSIONS, (
            f"Unsupported file: {libpath}"
        )

        collection_metadata = collection.get(
            blender.pipeline.AVALON_PROPERTY)
        collection_libpath = collection_metadata["libpath"]
        lib_container = collection_metadata["lib_container"]

        obj_container = plugin.get_local_collection_with_name(
            collection_metadata["obj_container"].name
        )
        objects = obj_container.all_objects

        container_name = obj_container.name

        normalized_collection_libpath = (
            str(Path(bpy.path.abspath(collection_libpath)).resolve())
        )
        normalized_libpath = (
            str(Path(bpy.path.abspath(str(libpath))).resolve())
        )
        self.log.debug(
            "normalized_collection_libpath:\n  %s\nnormalized_libpath:\n  %s",
            normalized_collection_libpath,
            normalized_libpath,
        )
        if normalized_collection_libpath == normalized_libpath:
            self.log.info("Library already loaded, not updating...")
            return

        # Get the armature of the rig
        armatures = [obj for obj in objects if obj.type == 'ARMATURE']
        assert(len(armatures) == 1)

        action = None
        if armatures[0].animation_data and armatures[0].animation_data.action:
            action = armatures[0].animation_data.action

        parent = plugin.get_parent_collection(obj_container)

        self._remove(objects, obj_container)

        obj_container = self._process(
            str(libpath), lib_container, container_name, action, parent)

        # Save the list of objects in the metadata container
        collection_metadata["obj_container"] = obj_container
        collection_metadata["objects"] = obj_container.all_objects
        collection_metadata["libpath"] = str(libpath)
        collection_metadata["representation"] = str(representation["_id"])

        bpy.ops.object.select_all(action='DESELECT')

    def remove(self, container: Dict) -> bool:
        """Remove an existing container from a Blender scene.

        Arguments:
            container (avalon-core:container-1.0): Container to remove,
                from `host.ls()`.

        Returns:
            bool: Whether the container was deleted.

        Warning:
            No nested collections are supported at the moment!
        """

        collection = bpy.data.collections.get(
            container["objectName"]
        )
        if not collection:
            return False
        assert not (collection.children), (
            "Nested collections are not supported."
        )

        collection_metadata = collection.get(
            blender.pipeline.AVALON_PROPERTY)

        obj_container = plugin.get_local_collection_with_name(
            collection_metadata["obj_container"].name
        )
        objects = obj_container.all_objects

        self._remove(objects, obj_container)

        bpy.data.collections.remove(collection)

        return True
