Alias Default hook

# Copyright (c) 2021 Autodesk, Inc.
#
# CONFIDENTIAL AND PROPRIETARY
#
# This work is provided "AS IS" and subject to the Shotgun Pipeline Toolkit
# Source Code License included in this distribution package. See LICENSE.
# By accessing, using, copying or modifying this work you indicate your
# agreement to the Shotgun Pipeline Toolkit Source Code License. All rights
# not expressly granted therein are reserved by Autodesk, Inc.

import os

import sgtk

import alias_api

HookBaseClass = sgtk.get_hook_baseclass()


class BreakdownSceneOperations(HookBaseClass):
    """
    Breakdown operations for Alias.

    This implementation handles detection of Alias references.
    """

    def __init__(self, *args, **kwargs):
        """Class constructor."""

        super(BreakdownSceneOperations, self).__init__(*args, **kwargs)

        # Keep track of the scene change callbacks that are registered, so that they can be
        # disconnected at a later time.
        self.__alias_event_callbacks = []

    def scan_scene(self):
        """
        The scan scene method is executed once at startup and its purpose is
        to analyze the current scene and return a list of references that are
        to be potentially operated on.

        The return data structure is a list of dictionaries. Each scene reference
        that is returned should be represented by a dictionary with three keys:

        - "node_name": The name of the 'node' that is to be operated on. Most DCCs have
          a concept of a node, path or some other way to address a particular
          object in the scene.
        - "node_type": The object type that this is. This is later passed to the
          update method so that it knows how to handle the object.
        - "path": Path on disk to the referenced object.
        - "extra_data": Optional key to pass some extra data to the update method
          in case we'd like to access them when updating the nodes.

        Toolkit will scan the list of items, see if any of the objects matches
        a published file and try to determine if there is a more recent version
        available. Any such versions are then displayed in the UI as out of date.
        """

        # Get all referencs currently in Alias
        alias_refs = alias_api.get_references()

        # The references result to return. This will be a list of references formatted in a
        # dictionary with the required key-values:
        #   'node_name': str, 'node_type': str, 'path': str
        refs = []

        # The set of file paths already added to the result. The result should contain items
        # with unique paths.
        file_paths = set()

        # For each Alias reference, add an item to the result for both its source path and path,
        # if they both exist. The next step to gather the file items will query for the
        # Published Files for each of these reference objects, which will decide if the source
        # path or path value should be used for this reference.
        for r in alias_refs:
            if r.source_path and r.source_path not in file_paths:
                refs.append(
                    {
                        "node_name": r.name,
                        "node_type": "reference",
                        "path": r.source_path.replace("/", os.path.sep),
                    }
                )
                file_paths.add(r.source_path)

            if r.path and r.path not in file_paths:
                refs.append(
                    {
                        "node_name": r.name,
                        "node_type": "reference",
                        "path": r.path.replace("/", os.path.sep),
                    }
                )
                file_paths.add(r.path)

        return refs

    def update(self, item):
        """
        Perform replacements given a number of scene items passed from the app.

        Once a selection has been performed in the main UI and the user clicks
        the update button, this method is called.

        :param item: Dictionary on the same form as was generated by the scan_scene hook above.
                     The path key now holds the path that the node should be updated *to* rather than the current path.
        :type item: dict

        :return: False to indicate that a file model update is not necessary (Alias events
            will trigger the update).
        :rtype: bool
        """

        node_type = item["node_type"]
        path = item["path"]
        extra_data = item["extra_data"]
        sg_data = item["sg_data"]
        if node_type == "reference":
            return self.update_reference(path, extra_data, sg_data)
        return False

    def update_reference(self, path, extra_data, sg_data):
        """
        Update the Alias reference from the given data.

        :param path: The new file path to set the Alias reference to.
        :type path: str
        :param extra_data: Additional data containing the existing Alias reference file path,
            which will be updated to the new file path. The existing path is used to look up
            the reference to update.
        :type extra_data: dict (required key-values: 'old_path': str)
        :param sg_data: Additional Flow Production Tracking specific data required to look up reference
            templates to perform the reference update.
        :type sg_data: dict (required key-values: 'project': dict, 'task': dict)

        :return: True if the reference was updated, False otherwise.
        :rtype: bool
        """

        old_path = extra_data["old_path"]
        _, ext = os.path.splitext(path)

        # if the new path is not a path to a wref file, we need to handle the conversion
        if ext != ".wref":
            tk = self.parent.engine.get_tk_from_project(sg_data["project"])
            source_template = tk.template_from_path(path)
            reference_template = self.parent.engine.get_reference_template(tk, sg_data)

            if source_template and reference_template:

                template_fields = source_template.get_fields(path)
                template_fields["alias.extension"] = ext[1:]
                reference_path = reference_template.apply_fields(template_fields)

                # do the same for the old path in order to get the associated reference path
                template_fields = source_template.get_fields(old_path)
                template_fields["alias.extension"] = ext[1:]
                old_path = reference_template.apply_fields(template_fields)

                if os.path.exists(reference_path):
                    self.logger.debug("File already converted!")
                    path = reference_path

                else:
                    self.logger.debug("Translating file to wref...")

                    # get the Alias Translations framework to translate the file to wref before importing it
                    framework = self.load_framework(
                        "tk-framework-aliastranslations_v0.x.x"
                    )
                    if not framework:
                        self.logger.error(
                            "Couldn't load tk-framework-aliastranslations. Skipping reference update for file {}.".format(
                                path
                            )
                        )
                        return False

                    tk_framework_aliastranslations = framework.import_module(
                        "tk_framework_aliastranslations"
                    )
                    translator = tk_framework_aliastranslations.Translator(
                        path, reference_path
                    )
                    translator.execute()
                    path = reference_path
            else:
                self.logger.error(
                    "Couldn't convert file to wref, missing templates. Skipping file {}...".format(
                        path
                    )
                )
                return False

        if not old_path or not os.path.exists(old_path):
            self.logger.info(
                "Couldn't find old reference path. Skipping file {}".format(path)
            )
            return False

        # get the reference by its uuid if possible, otherwise use its name to find the right instance
        alias_api.update_reference(old_path, path)
        return True

    def register_scene_change_callback(self, scene_change_callback):
        """
        Register the callback such that it is executed on a scene change event.

        This hook method is useful to reload the breakdown data when the data in the scene has
        changed.

        For Alias, the callback is registered with the AliasEngine event watcher to be
        triggered on a PostRetrieve event (e.g. when a file is opened).

        :param scene_change_callback: The callback to register and execute on scene chagnes.
        :type scene_change_callback: function
        """

        # Define the list of Alias event to that will trigger the scene change callback.
        events = [
            alias_api.AlMessageType.PostRetrieve,
            alias_api.AlMessageType.StageActive,
        ]

        # Alias event messages that are only available in version >= 2023.0
        if hasattr(alias_api.AlMessageType, "ReferenceFileAdded"):
            events.append(alias_api.AlMessageType.ReferenceFileAdded)
            # Only add the deleted event if the added event is available - otherwise
            # add/remove references doesn't work completely
            events.append(alias_api.AlMessageType.ReferenceFileDeleted)

        # Create the scene change callback to register with the Alias event watcher.
        scene_change_cb = (
            lambda result, cb=scene_change_callback: self.__handle_event_callback(
                result, cb
            )
        )

        # Keep track of the Alias event callbacks that will be registered, so that they can
        # properly be unregistered on shut down.
        self.__alias_event_callbacks = [(scene_change_cb, events)]

        # Register the event callbacks to the engine's event watcher
        for callback, events in self.__alias_event_callbacks:
            self.parent.engine.event_watcher.register_alias_callback(callback, events)

    def unregister_scene_change_callback(self):
        """Unregister the scene change callbacks by disconnecting any signals."""

        event_watcher = self.parent.engine.event_watcher
        if not event_watcher:
            # Engine already shutdown and removed event callbacks
            return

        # Unregister the event callbacks from the engine's event watcher
        for callback, events in self.__alias_event_callbacks:
            event_watcher.unregister_alias_callback(callback, events)

    def __handle_event_callback(self, event_result, scene_change_callback):
        """
        Intermediate callback handler for Alias events.

        Process the result returned by the Alias event that triggered the callback, to call
        the scene callback function with the appropriate parameters.

        :param event_result: The object returned by the Alias event.
        :type event_result: alias_api.MessageResult
        :param scene_change_callback: The callback to execute.
        :type scene_change_callback: function
        """

        if event_result.message_type == alias_api.AlMessageType.ReferenceFileDeleted:
            # Remove the reference from the model by its path.
            scene_change_callback(
                event_type="remove",
                data=event_result.reference_file_1_path,
            )

        elif (
            hasattr(alias_api.AlMessageType, "ReferenceFileAdded")
            and event_result.message_type == alias_api.AlMessageType.ReferenceFileAdded
        ):
            # Add the new reference to the model
            file_item_data = {
                "node_name": event_result.reference_file_1_name,
                "node_type": "reference",
                "path": event_result.reference_file_1_path,
            }
            scene_change_callback(event_type="add", data=file_item_data)

        else:
            # Requset a full model reload.
            scene_change_callback()