plugin

class SilhouetteCreator(Creator):
    """Base class for Silhouette creators.

    This base creator can be applied multiple times to a single node, where
    each instance is stored as a separate imprint on the node, inside an
    `AYON_instances` state that is a `dict[str, dict]` of instance data per
    uuid. The `instance_id` will then be defined by the node's id and this uuid
    as `{node_id}|{uuid}`.

    This way, a single RotoNode can have multiple instances of different
    product types (or even the same product type) to allow exporting e.g.
    track points and matte shapes from the same RotoNode.

    """
    default_variants = ["Main"]
    settings_category = "silhouette"

    create_node_type = "OutputNode"

    # When `valid_node_types` is set, all these node types are allowed to get
    # imprinted by this creator
    valid_node_types = set()

    @lib.undo_chunk("Create instance")
    def create(self, product_name, instance_data, pre_create_data):

        session = fx.activeSession()
        if not session:
            return

        instance_node = None
        use_selection = pre_create_data.get("use_selection")
        selected_nodes = []
        if use_selection:
            # Allow to imprint on a currently selected node of the same type
            # as this creator would generate. If the node is already imprinted
            # by a Creator then raise an error - otherwise use it as the
            # instance node.
            valid_node_types = self.valid_node_types or {self.create_node_type}
            selected_nodes = [
                node for node in fx.selection() if isinstance(node, fx.Node)]
            for node in selected_nodes:
                if node.type in valid_node_types:
                    data = lib.read(node)
                    if data and data.get("creator_identifier"):
                        raise CreatorError(
                            "Selected node is already imprinted by a Creator."
                        )
                    instance_node = node
                    self.log.info(
                        f"Using selected node as instance node: {node.label}")
                    break

        if instance_node is None:
            # Create new node and place it in the scene
            instance_node = fx.Node(self.create_node_type)
            session.addNode(instance_node)
            lib.set_new_node_position(instance_node)

            # When generating a new instance node and use selection is enabled,
            # connect to the first selected node with a matching output type
            if use_selection and selected_nodes:
                self._connect_input_to_first_matching_candidate(
                    instance_node, selected_nodes)

            instance_node.label = session.uniqueLabel(product_name)
        fx.activate(instance_node)

        # Use the uniqueness of the node in Silhouette as part of the instance
        # id, but because we support multiple instances per node, we also add
        # an uuid within the node to make duplicates of nodes still unique in
        # the full create context.
        instance_id = f"{instance_node.id}|{uuid.uuid4()}"
        instance_data["instance_id"] = instance_id
        instance_data["label"] = self._define_label(
            instance_node, product_name)
        instance = CreatedInstance(
            product_type=self.product_type,
            product_name=product_name,
            data=instance_data,
            creator=self,
            transient_data={
                "instance_node": instance_node
            }
        )

        # Store the instance data
        data = instance.data_to_store()
        self._imprint(instance_node, data)

        self._add_instance_to_context(instance)

        return instance

    def collect_instances(self):
        shared_data = cache_instance_data(self.collection_shared_data)
        cached_instances = shared_data["silhouette_cached_instances"]
        for obj, instance_uuid, data in cached_instances.get(
                self.identifier, []):
            data["instance_id"] = f"{obj.id}|{instance_uuid}"
            data["label"] = self._define_label(obj, data["productName"])

            # Add instance
            created_instance = CreatedInstance.from_existing(
                data,
                self,
                transient_data={"instance_node": obj}
            )
            self._add_instance_to_context(created_instance)

    @lib.undo_chunk("Update instances")
    def update_instances(self, update_list):
        for created_inst, _changes in update_list:
            new_data = created_inst.data_to_store()
            node = created_inst.transient_data["instance_node"]
            self._imprint(node, new_data)

    @lib.undo_chunk("Remove instances")
    def remove_instances(self, instances):
        for instance in instances:

            # Remove node from the scene
            node = instance.transient_data["instance_node"]
            if node:
                instance_uuid = self._get_uuid_from_instance_id(instance.id)
                instances_by_uuid = lib.read(node,
                                             key=INSTANCES_DATA_KEY) or {}
                instances_by_uuid.pop(instance_uuid, None)
                if not instances_by_uuid:
                    # Remove the node, because it was the last imprinted value
                    session = node.session
                    session.removeNode(node)
                else:
                    # Update the node's imprinted value by removing the entry
                    lib.imprint(
                        node,
                        instances_by_uuid,
                        key=INSTANCES_DATA_KEY)

            # Remove the collected CreatedInstance to remove from UI directly
            self._remove_instance_from_context(instance)

    def _imprint(self, node, data):
        data.pop("label", None)  # do not store the label
        # Do not store instance id since it's the Silhouette node id
        instance_id = data.pop("instance_id")

        instance_uuid = self._get_uuid_from_instance_id(instance_id)
        instances_by_uuid = lib.read(node, key=INSTANCES_DATA_KEY) or {}
        instances_by_uuid[instance_uuid] = data
        lib.imprint(node, instances_by_uuid, key=INSTANCES_DATA_KEY)

    def _get_uuid_from_instance_id(self, instance_id: str) -> str:
        """Return uuid for instance's key on the node data from instance id."""
        return instance_id.rsplit("|", 1)[-1]

    def _define_label(self, obj: fx.Node, product_name: str) -> str:
        return f"{product_name} ({obj.label})"

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

    def _connect_input_to_first_matching_candidate(self, node, candidates):
        """Connect the primary input of `node` to the first candidate with
        an output that has a matching data type."""
        primary_input = node.primaryInput
        if not primary_input:
            return

        allowed_types = set(primary_input.dataTypes)
        for candidate in candidates:
            for output in candidate.outputs:
                if allowed_types.intersection(output.dataTypes):
                    output.connect(primary_input)
                    return

class SilhouetteImportLoader(SilhouetteLoader):
    """Import using the chosen Silhouette IO module."""
    # TODO: Should we allow importing multiple containers to one node?
    #  It may be needed for importing both track points and matte shapes to
    #  a single RotoNode

    io_module: str
    node_type = "RotoNode"

    options = [
        BoolDef(
            "use_selection",
            label="Use selection",
            default=True,
            tooltip=(
                "Use the selected node if it supports the import type. "
                "Otherwise create a new node."
            )
        )
    ]

    @lib.undo_chunk("Load")
    @lib.maintained_selection()
    def load(self, context, name=None, namespace=None, options=None):
        """Merge the Alembic into the scene."""
        if not fx.activeProject():
            raise RuntimeError("No active project found.")
        if not fx.activeSession():
            raise RuntimeError("No active session found.")
        if options is None:
            options = {}

        # Use selected node or create a new one to import to
        selection = [
            _node for _node in fx.selection() if self.can_import_to_node(_node)
        ]
        if options.get("use_selection", True) and selection:
            node = selection[0]
        else:
            # Create a new node
            node = fx.Node(self.node_type)
            fx.activeSession().addNode(node)
            lib.set_new_node_position(node)


        # Import the file
        fx.activate(node)
        filepath = self.filepath_from_context(context)

        try:
            fx.io_modules[self.io_module].importFile(filepath)
        except AssertionError as exc:
            # Provide better message than "importer not ready" when it fails
            # to import
            if str(exc) == "importer not ready":
                raise LoadError(
                    f"Failed to import as '{self.io_module}':\n{filepath}"
                    f"\n\n"
                    f"Most likely the Tracker format is unsupported."
                ) from exc
            raise
        self._process_loaded(context, node)

        # property.hidden = True  # hide the attribute
        lib.imprint(node, data={
            "name": str(name),
            "namespace": str(namespace),
            "loader": str(self.__class__.__name__),
            "representation": context["representation"]["id"],
        })

    def _process_loaded(self, context, node):
        # For overrides on inherited classes
        pass

    @lib.undo_chunk("Update Source")
    @lib.maintained_selection()
    def update(self, container, context):
        item: fx.Node = container["_item"]

        # Remove existing children
        item.objects.removeObjects(item.children)

        # Import the file
        fx.activate(item)
        filepath = self.filepath_from_context(context)
        fx.io_modules[self.io_module].importFile(filepath)

        # Update representation id
        data = lib.read(item)
        data["representation"] = context["representation"]["id"]
        lib.imprint(item, data)

    @lib.undo_chunk("Remove container")
    def remove(self, container):
        """Remove node from session"""
        node: fx.Node = container["_item"]
        session = container["_session"]
        session.removeNode(node)

    def switch(self, container, context):
        """Support switch to another representation."""
        self.update(container, context)

    def can_import_to_node(self, node) -> bool:

        if isinstance(node, fx.Node):
            return True

        # Do not allow load to node that already has import
        # TODO: Support multiple containers per node
        if parse_container(node):
            return False

        return False