plugin

class LayoutLoader(Loader):
    """Load Layout from a JSON file"""

    product_base_types = {"layout"}
    product_types = product_base_types
    representations = {"json"}

    label = "Load Layout"
    icon = "code-fork"
    color = "orange"
    loaded_layout_dir = "{folder[path]}/{product[name]}"
    loaded_layout_name = "{folder[name]}_{product[name]}_{version[version]}"
    remove_loaded_assets = False

    @staticmethod
    def _get_fbx_loader(loaders, family):
        name = ""
        if family in ['rig', 'skeletalMesh']:
            name = "SkeletalMeshFBXLoader"
        elif family in ['model', 'staticMesh']:
            name = "StaticMeshFBXLoader"
        elif family == 'camera':
            name = "CameraLoader"

        if not name:
            return None

        for loader in loaders:
            if loader.__name__ == name:
                return loader

        return None

    @staticmethod
    def _get_abc_loader(loaders, family):
        name = ""
        if family in ['rig', 'skeletalMesh']:
            name = "SkeletalMeshAlembicLoader"
        elif family in ['model', 'staticMesh']:
            name = "StaticMeshAlembicLoader"
        elif family in ["animation"]:
            name = "AnimationAlembicLoader"
        if not name:
            return None

        for loader in loaders:
            if loader.__name__ == name:
                return loader

        return None

    def _transform_from_basis(self, transform, basis, unreal_import=False):
        """Transform a transform from a basis to a new basis."""
        # Get the basis matrix
        basis_matrix = unreal.Matrix(
            basis[0],
            basis[1],
            basis[2],
            basis[3]
        )
        transform_matrix = unreal.Matrix(
            transform[0],
            transform[1],
            transform[2],
            transform[3]
        )

        new_transform = None
        if unreal_import:
            new_transform = transform_matrix * basis_matrix
        else:
            new_transform = (
                basis_matrix.get_inverse() * transform_matrix * basis_matrix)

        return new_transform.transform()

    def _get_repre_entities_by_version_id(self, project_name, data, repre_extension, force_loaded=False):
        version_ids = {
            element.get("version")
            for element in data
            if element.get("representation")
        }
        version_ids.discard(None)
        output = collections.defaultdict(list)
        if not version_ids:
            return output
        # Extract extensions from data with backward compatibility for "ma"
        extensions = {
            element.get("extension", "ma")
            for element in data
            if element.get("representation")
        }

        # Update extensions based on the force_loaded flag
        updated_extensions = set()

        for ext in extensions:
            if not force_loaded or repre_extension == "json":
                if ext == "ma":
                    updated_extensions.update({"fbx", "abc"})
                else:
                    updated_extensions.add(ext)
            else:
                updated_extensions.update({repre_extension})

        repre_entities = ayon_api.get_representations(
            project_name,
            representation_names=updated_extensions,
            version_ids=version_ids,
            fields={"id", "versionId", "name"}
        )
        for repre_entity in repre_entities:
            version_id = repre_entity["versionId"]
            output[version_id].append(repre_entity)
        return output

    def imprint(
        self,
        context: dict[str, Any],
        folder_path: str,
        folder_name: str,
        loaded_assets: list[str],
        asset_dir: str,
        asset_name: str,
        container_name: str,
        project_name: str,
        hierarchy_dir: Optional[str] = None,
    ) -> None:
        """Imprint the container with the necessary data.

        Args:
            context (dict): The context of the loading process.
            folder_path (str): The path to the folder where the layout is located.
            folder_name (str): The name of the folder
            loaded_assets (list): List of loaded assets.
            asset_dir (str): The asset directory.
            asset_name (str): The asset name.
            container_name (str): The name of the container.
            project_name (str): The name of the project.
            hierarchy_dir (str, optional): The directory of the hierarchy.
                Defaults to None.

        Note:
            This method is re-implemented with different signatures in
            many loader plugins. We should consider refactoring it in the
            future o avoid code duplication.

        """
        data = {
            "schema": "ayon:container-2.0",
            "id": AYON_CONTAINER_ID,
            "asset": folder_name,
            "folder_path": folder_path,
            "namespace": asset_dir,
            "container_name": container_name,
            "asset_name": asset_name,
            "loader": str(self.__class__.__name__),
            "representation": context["representation"]["id"],
            "parent": context["representation"]["versionId"],
            "product_base_type": context["product"]["productBaseType"],
            "family": context["product"]["productBaseType"],
            "loaded_assets": loaded_assets,
            "project_name": project_name
        }
        if hierarchy_dir is not None:
            data["master_directory"] = hierarchy_dir
        imprint(f"{asset_dir}/{container_name}", data)

    def _load_assets(
            self,
            instance_name,
            repre_id,
            product_base_type,
            repr_format):
        all_loaders = discover_loader_plugins()
        loaders = loaders_from_representation(
            all_loaders, repre_id)

        loader = None

        if repr_format == 'fbx':
            loader = self._get_fbx_loader(
                loaders, product_base_type)
        elif repr_format == 'abc':
            loader = self._get_abc_loader(
                loaders, product_base_type)

        if not loader:
            if repr_format == "ma":
                msg = (
                    f"No valid {product_base_type} loader found "
                    f"for {repre_id} ({repr_format}), "
                    f"consider using {product_base_type} loader "
                    "(fbx/abc) instead."
                )
                self.log.warning(msg)
            else:
                self.log.error(
                    f"No valid loader found for {repre_id} "
                    f"({repr_format}) "
                    f"{product_base_type}")
            return

        import_options = {
            "layout": True
        }
        assets = load_container(
            loader,
            repre_id,
            namespace=instance_name,
            options=import_options
        )
        return assets

    def _remove_Loaded_asset(self, container):
        """
        Delete the layout. First, check if the assets loaded with the layout
        are used by other layouts. If not, delete the assets.
        """
        if self.remove_loaded_assets:
            remove_asset_confirmation_dialog = unreal.EditorDialog.show_message(
                "The removal of the loaded assets",
                "The layout will be removed. Do you want to delete all associated assets as well?",
                unreal.AppMsgType.YES_NO)
            if (remove_asset_confirmation_dialog == unreal.AppReturnType.YES):
                remove_loaded_asset(container)

class Loader(LoaderPlugin, ABC):
    """This serves as skeleton for future Ayon specific functionality"""

class UnrealActorCreator(UnrealBaseCreator):
    """Base class for Unreal creator plugins based on actors."""

    def create(self, product_name, instance_data, pre_create_data):
        """Create instance of the asset.

        Args:
            product_name (str): Name of the product.
            instance_data (dict): Data for the instance.
            pre_create_data (dict): Data for the instance.

        Returns:
            CreatedInstance: Created instance.
        """
        try:
            if UNREAL_VERSION.major == 5:
                world = unreal.UnrealEditorSubsystem().get_editor_world()
            else:
                world = unreal.EditorLevelLibrary.get_editor_world()

            # Check if the level is saved
            if world.get_path_name().startswith("/Temp/"):
                raise CreatorError(
                    "Level must be saved before creating instances.")

            # Check if instance data has members, filled by the plugin.
            # If not, use selection.
            if not instance_data.get("members"):
                actor_subsystem = unreal.EditorActorSubsystem()
                sel_actors = actor_subsystem.get_selected_level_actors()
                selection = [a.get_path_name() for a in sel_actors]

                instance_data["members"] = selection
            instance_data["level"] = world.get_path_name()

            super(UnrealActorCreator, self).create(
                product_name,
                instance_data,
                pre_create_data)

        except Exception as exc:
            raise CreatorError(f"Creator error: {exc}") from exc

    def get_pre_create_attr_defs(self):
        return [
            UILabelDef("Select actors to create instance from them."),
        ]

class UnrealAssetCreator(UnrealBaseCreator):
    """Base class for Unreal creator plugins based on assets."""

    def create(self, product_name, instance_data, pre_create_data):
        """Create instance of the asset.

        Args:
            product_name (str): Name of the product.
            instance_data (dict): Data for the instance.
            pre_create_data (dict): Data for the instance.

        Returns:
            CreatedInstance: Created instance.
        """
        try:
            # Check if instance data has members, filled by the plugin.
            # If not, use selection.
            if not pre_create_data.get("members"):
                pre_create_data["members"] = []

                if pre_create_data.get("use_selection"):
                    utilib = unreal.EditorUtilityLibrary
                    sel_objects = utilib.get_selected_assets()
                    pre_create_data["members"] = [
                        a.get_path_name() for a in sel_objects]

            super(UnrealAssetCreator, self).create(
                product_name,
                instance_data,
                pre_create_data)

        except Exception as exc:
            raise CreatorError(f"Creator error: {exc}") from exc

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

class UnrealBaseAutoCreator(AutoCreator, UnrealCreateLogic):
    """Base class for Unreal auto creator plugins."""

    settings_category = "unreal"

    def collect_instances(self):
        return self._default_collect_instances()

    def update_instances(self, update_list):
        return self._default_update_instances(update_list)

    def remove_instances(self, instances):
        return self._default_remove_instances(instances)

class UnrealBaseCreator(UnrealCreateLogic, Creator):
    """Base class for Unreal creator plugins."""

    settings_category = "unreal"

    def create(self, product_name, instance_data, pre_create_data):
        self.create_unreal(product_name, instance_data, pre_create_data)

    def collect_instances(self):
        return self._default_collect_instances()

    def update_instances(self, update_list):
        return self._default_update_instances(update_list)

    def remove_instances(self, instances):
        return self._default_remove_instances(instances)

class UnrealCreateLogic:
    """Universal class for logic that Unreal creators could inherit from."""
    root = "/Game/Ayon/AyonPublishInstances"
    suffix = "_INS"


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

        Create `unreal_cached_products` 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
        `unreal_cached_legacy_products` there and fill it with
        all legacy products under product_base_type as a key.

        Args:
            Dict[str, Any]: Shared data.

        Return:
            Dict[str, Any]: Shared data dictionary.

        """
        if shared_data.get("unreal_cached_products") is None:
            unreal_cached_products = collections.defaultdict(list)
            unreal_cached_legacy_products = collections.defaultdict(list)
            for instance in ls_inst():
                creator_id = instance.get("creator_identifier")
                if creator_id:
                    unreal_cached_products[creator_id].append(instance)
                else:
                    # Handle legacy instances that may use "product_type"
                    # instead of "product_base_type" to avoid KeyError.
                    product_base_type = (
                            instance.get("product_base_type")
                            or instance.get("product_type")
                    )
                    if product_base_type is None:
                        unreal.log_warning(
                            f"Legacy instance without product_base_type or "
                            f"product_type: {instance}"
                        )
                        continue
                    unreal_cached_legacy_products[product_base_type].append(
                        instance)

            shared_data["unreal_cached_products"] = unreal_cached_products
            shared_data["unreal_cached_legacy_products"] = (
                unreal_cached_legacy_products
            )
        return shared_data

    def _default_collect_instances(self):
        # cache instances if missing
        self.get_cached_instances(self.collection_shared_data)
        for instance in self.collection_shared_data[
                "unreal_cached_products"].get(self.identifier, []):
            # Unreal saves metadata as string, so we need to convert it back
            instance['creator_attributes'] = ast.literal_eval(
                instance.get('creator_attributes', '{}'))
            instance['publish_attributes'] = ast.literal_eval(
                instance.get('publish_attributes', '{}'))
            instance['members'] = ast.literal_eval(
                instance.get('members', '[]'))
            instance['families'] = ast.literal_eval(
                instance.get('families', '[]'))
            instance['active'] = ast.literal_eval(
                instance.get('active', ''))
            created_instance = CreatedInstance.from_existing(instance, self)
            self._add_instance_to_context(created_instance)

    def _default_update_instances(self, update_list):
        for created_inst, changes in update_list:
            instance_node = created_inst.get("instance_path", "")

            if not instance_node:
                unreal.log_warning(
                    f"Instance node not found for {created_inst}")
                continue

            new_values = {
                key: changes[key].new_value
                for key in changes.changed_keys
            }
            imprint(
                instance_node,
                new_values
            )

    def _default_remove_instances(self, instances):
        for instance in instances:
            instance_node = instance.data.get("instance_path", "")
            if instance_node:
                unreal.EditorAssetLibrary.delete_asset(instance_node)

            self._remove_instance_from_context(instance)


    def create_unreal(self, product_name, instance_data, pre_create_data):
        try:
            instance_name = f"{product_name}{self.suffix}"
            pub_instance = create_publish_instance(instance_name, self.root)

            instance_data["product_name"] = product_name
            instance_data["instance_path"] = f"{self.root}/{instance_name}"

            product_type: str = instance_data.get("product_type")
            if not product_type:
                product_type = self.product_base_type

            instance = CreatedInstance(
                product_type=product_type,
                product_base_type=self.product_base_type,
                product_name=product_name,
                data=instance_data,
                creator=self,
            )
            self._add_instance_to_context(instance)

            pub_instance.set_editor_property('add_external_assets', True)
            assets = pub_instance.get_editor_property('asset_data_external')

            ar = unreal.AssetRegistryHelpers.get_asset_registry()

            for member in pre_create_data.get("members", []):
                obj = ar.get_asset_by_object_path(member).get_asset()
                assets.add(obj)

            imprint(f"{self.root}/{instance_name}",
                    instance.data_to_store())

            return instance

        except Exception as exc:
            raise CreatorError(f"Creator error: {exc}") from exc