# -*- coding: utf-8 -*-
"""
:author:
    Jeremy Ernst
:description:
    This module contains the class for saving and loading templates, which store components, their settings, their
    positions, and their joint name overrides. These can then be loaded in whole or in part, and can have their
    settings, positions, or overrides loaded onto existing components as well.
"""

import artv2.tools.system.logger.output_logger as logger
import artv2.utilities.general_utilities as utils
import artv2.utilities.component_utilities as component_utils
import artv2.components.base_components.rig_asset as rig_asset
import artv2.components.root as root
import pymel.core as pm


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
class Template(object):
    """
    The template class saves and loads templates. Templates contain a dictionary of components, their settings, their
    positions, and their joint name overrides.
    """

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def __init__(self, path=None):

        self.logger = logger.OutputLogger(self.__class__.__name__)

        self.path = utils.path_unify(path)
        self.load_settings = True
        self.load_positions = True
        self.load_overrides = True

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def save_template(self):
        """
        Saves a template file to disk.

        example usage:
            .. code-block:: python

                # get a filename to save a template to, then create a template instance and save the template.
                path = QtWidgets.QFileDialog.getSaveFileName(self, "Save Template", template_path, "*.json")
                if os.path.exists(os.path.dirname(path[0])):
                    template_instance = template.Template(utils.path_unify(path[0]))
                    template_instance.save_template()
        """

        self.logger.info("{0}.save_template: Saving template to {1}.".format(self.__class__.__name__, self.path))
        template_data = self._create_template_data()
        utils.save_to_file(self.path, template_data)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def load_template(self, components, aim=False):
        """
        Loads a template file from disk. If the components don't exist (from the passed in list), then those components
        are created. Then, depending on the settings (load settings, positions, overrides), data is applied to those
        components.

        example usage:
            .. code-block:: python

                # create a template instance using an existing path to a template on disk.
                inst = template.Template(r"C:\\Users\\jernst\\Documents\\ARTv2\\resources\\templates\\test.json")
                inst.load_settings=False
                inst.load_template(["clavicle_l", "arm_l", "torso", "master"])

        :param components: A list of the components from the template to load.
        :param aim: Whether or not aim model should be turned on the component.
        """

        self.logger.info("{0}.load_template: Loading template (from {1}) onto the following components: {2}"
                         "".format(self.__class__.__name__, self.path, components))

        template_data = utils.load_from_file(self.path)

        # ensure an asset and master exist in the scene. If not, create them
        assets = component_utils.get_all_rig_asset_nodes()
        if not assets:
            rig_asset.RigAsset()
            root.Root()

        # load components in scene
        self._load_components(template_data, components)

        # load data onto components
        existing_components = component_utils.get_all_component_names()
        for component in components:
            if component in template_data.keys():
                data = template_data.get(component)
                inst = component_utils.get_component_instance(existing_components.get(component))

                if self.load_settings:
                    self._load_settings(inst, data)

                if self.load_positions:
                    mover_data = data.get("POSITIONS")
                    inst.joint_mover.paste_transforms(template=True, data=mover_data)
                    if aim:
                        inst.joint_mover.aim_helper.toggle_aim_mode()

                if self.load_overrides:
                    self._set_overrides(inst, data)
                del inst

        # set parents
        self._set_parents(components, template_data)
        self._set_mirrors(components, template_data)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _set_overrides(self, inst, data):

        self.logger.info("{0}._set_overrides: Setting joint name override data on {1}"
                         "".format(self.__class__.__name__, inst))

        override_data = data.get("OVERRIDES")
        movers = inst.joint_mover.get_movers()
        offset_movers = movers.get("offset")

        for offset in offset_movers:
            if offset.nodeName() in override_data.keys():
                connected_joint = offset.connected_joint.connections(type="transform")[0]
                created_joint = offset.created_joint.get()
                old_name = inst.prefix + connected_joint.otherType.get() + inst.suffix
                if created_joint != override_data.get(offset.nodeName()):
                    inst.rename_joint(override_data.get(offset.nodeName()), old_name, offset)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _set_parents(self, components, template_data):

        self.logger.info("{0}._set_parents: Settings parents for {1}".format(self.__class__.__name__, components))

        existing_components = component_utils.get_all_component_names()
        for component in components:
            data = template_data.get(component)
            inst = component_utils.get_component_instance(existing_components.get(component))

            settings_data = data.get("SETTINGS")
            if not settings_data.get("parent") is None:
                existing_joints = component_utils.get_all_created_joints()
                if settings_data.get("parent") in existing_joints:
                    inst.parent = settings_data.get("parent")
                else:
                    inst.parent = "root"

            del inst

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _set_mirrors(self, components, template_data):

        self.logger.info("{0}._set_mirrors: Setting mirrors for {1}".format(self.__class__.__name__, components))

        existing_components = component_utils.get_all_component_names()

        for component in components:
            data = template_data.get(component)
            inst = component_utils.get_component_instance(existing_components.get(component))

            settings_data = data.get("SETTINGS")
            if settings_data.get("has_mirror"):
                if settings_data.get("mirror_component") is not None:
                    inst.set_mirror(pm.PyNode(settings_data.get("mirror_component")))

            del inst

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _load_settings(self, instance, data):

        self.logger.info("{0}._load_settings: Loading settings onto {1}".format(self.__class__.__name__, instance))

        settings_data = data.get("SETTINGS")
        properties = utils.get_class_properties(instance)

        if "side" in settings_data.keys():
            if settings_data.get("side") == 1:
                instance.set_side("right")

        for key in settings_data:
            if key in properties:
                setattr(instance, key, settings_data.get(key))

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _load_components(self, template_data, components):

        existing_components = component_utils.get_all_component_names()
        for component in components:
            if component in template_data.keys():

                data = template_data.get(component)
                info = data.get("INFO")
                settings_data = data.get("SETTINGS")

                # check if component exists in scene. If not, create it
                if component not in existing_components.keys():
                    self.logger.info("{0}._load_components: {1} does not exist. Creating component."
                                     "".format(self.__class__.__name__, component))

                    mod = __import__(info.get("PATH"), {}, {}, [info.get("MODULE")])
                    module_class = getattr(mod, info.get("CLASS"))
                    module_class(prefix=settings_data.get("prefix"), suffix=settings_data.get("suffix"))

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _create_template_data(self):

        self.logger.info("{0}._create_template_data: Gathering data to build template.".format(self.__class__.__name__))

        template = {}

        # get all component instances
        components = component_utils.return_all_components()

        # for each component, get its settings, mover positions, and joint name overrides
        for component in components:
            inst = component_utils.get_component_instance(component)

            # CLASS INFO
            class_name = component.className.get()
            module_path = component.modulePath.get()
            module_name = module_path.rpartition(".")[2]

            class_info = {"CLASS": class_name, "PATH": module_path, "MODULE": module_name}

            # SETTINGS
            settings = utils.get_class_properties(inst)
            settings["prefix"] = inst.prefix
            settings["suffix"] = inst.suffix
            settings["parent"] = inst.parent
            settings["has_mirror"] = component.has_mirror.get()
            settings["mirror_component"] = None
            if component.has_mirror.get():
                settings["mirror_component"] = component.mirror_component.connections()[0].nodeName()
            if inst.has_sides:
                settings["side"] = component.side.get()

            # MOVER POSITIONS
            mover_positions = inst.joint_mover.copy_transforms(template=True)

            # JOINT NAME OVERRIDES
            movers = inst.joint_mover.get_movers()
            offset_movers = movers.get("offset")
            override_data = {}
            for mover in offset_movers:

                driven_joint = mover.connected_joint.connections(type="transform")[0]
                original_name = inst.prefix + driven_joint.otherType.get() + inst.suffix
                if mover.created_joint.get() != original_name:
                    override_data[mover.nodeName()] = mover.created_joint.get()

            template[component.componentName.get()] = {"SETTINGS": settings,
                                                       "INFO": class_info,
                                                       "POSITIONS": mover_positions,
                                                       "OVERRIDES": override_data}

            del inst
        return template