api

class NukeCreator(NewCreator):
    node_class_name = None

    def _pass_pre_attributes_to_instance(
        self,
        instance_data,
        pre_create_data,
        keys=None
    ):
        if keys is None:
            keys = pre_create_data.keys()
        creator_attrs = instance_data["creator_attributes"] = {}

        creator_attrs.update({
            key: value
            for key, value in pre_create_data.items()
            if key in keys
        })

    def check_existing_product(self, product_name):
        """Make sure product name is unique.

        It search within all nodes recursively
        and checks if product name is found in
        any node having instance data knob.

        Arguments:
            product_name (str): Product name
        """

        for node in nuke.allNodes(recurseGroups=True):
            # make sure testing node is having instance knob
            if INSTANCE_DATA_KNOB not in node.knobs().keys():
                continue
            node_data = get_node_data(node, INSTANCE_DATA_KNOB)

            if not node_data:
                # a node has no instance data
                continue

            # test if product name is matching
            if node_data.get("productType") == product_name:
                raise NukeCreatorError(
                    (
                        "A publish instance for '{}' already exists "
                        "in nodes! Please change the variant "
                        "name to ensure unique output."
                    ).format(product_name)
                )

    def create_instance_node(
        self,
        node_name,
        knobs=None,
        parent=None,
        node_type=None,
        node_selection=None,
    ):
        """Create node representing instance.

        Arguments:
            node_name (str): Name of the new node.
            knobs (OrderedDict): node knobs name and values
            parent (str): Name of the parent node.
            node_type (str, optional): Nuke node Class.
            node_selection (Optional[list[nuke.Node]]): The node selection.

        Returns:
            nuke.Node: Newly created instance node.

        """
        node_type = node_type or "NoOp"

        node_knobs = knobs or {}

        # set parent node
        parent_node = nuke.root()
        if parent:
            parent_node = nuke.toNode(parent)

        try:
            with parent_node:
                created_node = nuke.createNode(node_type)
                created_node["name"].setValue(node_name)

                for key, values in node_knobs.items():
                    if key in created_node.knobs():
                        created_node["key"].setValue(values)
        except Exception as _err:
            raise NukeCreatorError("Creating have failed: {}".format(_err))

        return created_node

    def _get_current_selected_nodes(
        self,
        pre_create_data,
        class_name: str = None,
    ):
        """ Get current node selection.

        Arguments:
            pre_create_data (dict): The creator initial data.
            class_name (Optional[str]): Filter on a class name.

        Returns:
            list[nuke.Node]: node selection.
        """
        class_name = class_name or self.node_class_name
        use_selection = pre_create_data.get("use_selection")

        if use_selection:
            selected_nodes = nuke.selectedNodes()
        else:
            selected_nodes = nuke.allNodes()

        if class_name:
            # Allow class name implicit last versions of class names like
            # `Camera` to match any of its explicit versions, e.g.
            # `Camera3` or `Camera4`.
            if not class_name[-1].isdigit():
                # Match name with any digit
                pattern = rf"^{class_name}\d*$"
            else:
                pattern = class_name
            regex = re.compile(pattern)
            selected_nodes = [
                node
                for node in selected_nodes
                if regex.match(node.Class())
            ]

        if class_name and use_selection and not selected_nodes:
            raise NukeCreatorError(f"Select a {class_name} node.")

        return selected_nodes

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

        # make sure selected nodes are detected early on.
        # we do not want any further Nuke operation to change the selection.
        node_selection = self._get_current_selected_nodes(pre_create_data)

        # make sure product name is unique
        self.check_existing_product(product_name)

        try:
            instance_node = self.create_instance_node(
                product_name,
                node_type=instance_data.pop("node_type", None),
                node_selection=node_selection,
            )
            instance = CreatedInstance(
                self.product_type,
                product_name,
                instance_data,
                self
            )

            self.apply_staging_dir(instance)
            instance.transient_data["node"] = instance_node

            self._add_instance_to_context(instance)

            set_node_data(
                instance_node, INSTANCE_DATA_KNOB, instance.data_to_store())

            return instance

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

    def collect_instances(self):
        cached_instances = _collect_and_cache_nodes(self)
        attr_def_keys = {
            attr_def.key
            for attr_def in self.get_instance_attr_defs()
        }
        attr_def_keys.discard(None)

        for (node, data) in cached_instances[self.identifier]:
            created_instance = CreatedInstance.from_existing(
                data, self
            )

            self.apply_staging_dir(created_instance)
            created_instance.transient_data["node"] = node
            self._add_instance_to_context(created_instance)

            for key in (
                set(created_instance["creator_attributes"].keys())
                - attr_def_keys
            ):
                created_instance["creator_attributes"].pop(key)

    def update_instances(self, update_list):
        for created_inst, changes in update_list:
            instance_node = created_inst.transient_data["node"]

            # in case node is not existing anymore (user erased it manually)
            try:
                instance_node.fullName()
            except ValueError:
                self._remove_instance_from_context(created_inst)
                continue

            # update instance node name if product name changed
            if "productName" in changes.changed_keys:
                instance_node["name"].setValue(
                    changes["productName"].new_value
                )

            set_node_data(
                instance_node,
                INSTANCE_DATA_KNOB,
                created_inst.data_to_store()
            )

    def remove_instances(self, instances):
        for instance in instances:
            remove_instance(instance)
            self._remove_instance_from_context(instance)

    def get_pre_create_attr_defs(self):
        return [
            BoolDef(
                "use_selection",
                default=not self.create_context.headless,
                label="Use selection"
            )
        ]

    def get_creator_settings(self, project_settings, settings_key=None):
        if not settings_key:
            settings_key = self.__class__.__name__
        return project_settings["nuke"]["create"][settings_key]

class NukeHost(
    HostBase, IWorkfileHost, ILoadHost, IPublishHost
):
    name = "nuke"

    def open_workfile(self, filepath):
        return open_file(filepath)

    def save_workfile(self, filepath=None):
        return save_file(filepath)

    def work_root(self, session):
        return work_root(session)

    def get_current_workfile(self):
        return current_file()

    def workfile_has_unsaved_changes(self):
        return has_unsaved_changes()

    def get_workfile_extensions(self):
        return file_extensions()

    def get_containers(self):
        return ls()

    def install(self):
        """Installing all requirements for Nuke host"""

        pyblish.api.register_host("nuke")

        self.log.info("Registering Nuke plug-ins..")
        pyblish.api.register_plugin_path(PUBLISH_PATH)
        register_loader_plugin_path(LOAD_PATH)
        register_creator_plugin_path(CREATE_PATH)
        register_inventory_action_path(INVENTORY_PATH)
        register_workfile_build_plugin_path(WORKFILE_BUILD_PATH)

        # Register AYON event for workfiles loading.
        register_event_callback("workio.open_file", check_inventory_versions)
        register_event_callback("taskChanged", change_context_label)

        _install_menu()

        # add script menu
        add_scripts_menu()
        add_scripts_gizmo()

        add_nuke_callbacks()

        launch_workfiles_app()

    def get_context_data(self):
        root_node = nuke.root()
        return get_node_data(root_node, ROOT_DATA_KNOB)

    def update_context_data(self, data, changes):
        root_node = nuke.root()
        set_node_data(root_node, ROOT_DATA_KNOB, data)

class NukeWriteCreator(NukeCreator):
    """Add Publishable Write node"""

    identifier = "create_write"
    label = "Create Write"
    product_type = "write"
    icon = "sign-out"

    temp_rendering_path_template = (  # default to be applied if settings is missing
        "{work}/renders/nuke/{product[name]}/{product[name]}.{frame}.{ext}")

    render_target = "local"  # default to be applied if settings is missing

    def get_linked_knobs(self):
        linked_knobs = []
        if "channels" in self.instance_attributes:
            linked_knobs.append("channels")
        if "ordered" in self.instance_attributes:
            linked_knobs.append("render_order")
        if "use_range_limit" in self.instance_attributes:
            linked_knobs.extend(["___", "first", "last", "use_limit"])

        return linked_knobs

    def integrate_links(self, node_selection, node, outputs=True):
        # skip if no selection
        if not node_selection:  # selection should contain either 1 or no node.
            return

        # collect dependencies
        input_nodes = node_selection
        dependent_nodes = node_selection[0].dependent() if outputs else []

        # relinking to collected connections
        for i, input in enumerate(input_nodes):
            node.setInput(i, input)

        # make it nicer in graph
        node.autoplace()

        # relink also dependent nodes
        for dep_nodes in dependent_nodes:
            dep_nodes.setInput(0, node)

    def _get_current_selected_nodes(
        self,
        pre_create_data,
    ):
        """ Get current node selection.

        Arguments:
            pre_create_data (dict): The creator initial data.
            class_name (Optional[str]): Filter on a class name.

        Returns:
            list[nuke.Node]: node selection.

        Raises:
            NukeCreatorError. When the selection contains more than 1 Write node.
        """
        if not pre_create_data.get("use_selection"):
            return []

        selected_nodes = super()._get_current_selected_nodes(
            pre_create_data,
            class_name=None,
        )

        if not selected_nodes:
            raise NukeCreatorError("No active selection")

        elif len(selected_nodes) > 1:
            raise NukeCreatorError("Select only one node")

        return selected_nodes

    def update_instances(self, update_list):
        super().update_instances(update_list)
        for created_inst, changes in update_list:
            # ensure was not deleted by super()
            if self.create_context.get_instance_by_id(created_inst.id):
                self._update_write_node_filepath(created_inst, changes)

    def _update_write_node_filepath(self, created_inst, changes):
        """Update instance node on context changes.

        Whenever any of productName, folderPath, task or productType
        changes then update:
        - output filepath of the write node
        - instance node's name to the product name
        """
        keys = ("productName", "folderPath", "task", "productType")
        if not any(key in changes.changed_keys for key in keys):
            # No relevant changes, no need to update
            return

        data = created_inst.data_to_store()
        # Update values with new formatted path
        instance_node = created_inst.transient_data["node"]
        formatting_data = copy.deepcopy(data)
        write_node = nuke.allNodes(group=instance_node, filter="Write")[0]
        _, ext = os.path.splitext(write_node["file"].value())
        formatting_data.update({"ext": ext[1:]})

        # Retieve render template and staging directory.
        fpath_template = self.temp_rendering_path_template
        formatting_data["work"] = get_work_default_directory(formatting_data)
        fpath = StringTemplate(fpath_template).format_strict(formatting_data)

        staging_dir = self.apply_staging_dir(created_inst)
        if staging_dir:
            basename = os.path.basename(fpath)
            staging_path = pathlib.Path(staging_dir)/ basename
            fpath = staging_path.as_posix()

        write_node["file"].setValue(fpath)

    def get_pre_create_attr_defs(self):
        attrs_defs = super().get_pre_create_attr_defs()
        attrs_defs.append(self._get_render_target_enum())

        return attrs_defs

    def get_instance_attr_defs(self):
        attr_defs = [self._get_render_target_enum()]

        # add reviewable attribute
        if "reviewable" in self.instance_attributes:
            attr_defs.append(
                BoolDef(
                    "review",
                    default=True,
                    label="Review"
                )
            )

        return attr_defs

    def _get_render_target_enum(self):
        rendering_targets = {
            "local": "Local machine rendering",
            "frames": "Use existing frames"
        }

        if "farm_rendering" in self.instance_attributes:
            rendering_targets.update({
                "frames_farm": "Use existing frames - farm",
                "farm": "Farm rendering",
            })

        return EnumDef(
            "render_target",
            items=rendering_targets,
            default=self.render_target,
            label="Render target",
            tooltip="Define the render target."
        )

    def create(self, product_name, instance_data, pre_create_data):
        if not pre_create_data:
            # add no selection for headless
            pre_create_data = {
                "use_selection": False
            }

        # pass values from precreate to instance
        self._pass_pre_attributes_to_instance(
            instance_data,
            pre_create_data,
            [
                "active_frame",
                "render_target"
            ]
        )
        # make sure selected nodes are added
        node_selection = self._get_current_selected_nodes(pre_create_data)

        # make sure product name is unique
        self.check_existing_product(product_name)

        try:
            instance = CreatedInstance(
                self.product_type,
                product_name,
                instance_data,
                self
            )

            staging_dir = self.apply_staging_dir(instance)
            instance_node = self.create_instance_node(
                product_name,
                instance_data,
                staging_dir=staging_dir,
                node_selection=node_selection,
            )

            instance.transient_data["node"] = instance_node

            self._add_instance_to_context(instance)

            set_node_data(
                instance_node,
                INSTANCE_DATA_KNOB,
                instance.data_to_store()
            )

            exposed_write_knobs(
                self.project_settings, self.__class__.__name__, instance_node
            )

            return instance

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

    def apply_settings(self, project_settings):
        """Method called on initialization of plugin to apply settings."""

        # plugin settings
        plugin_settings = self.get_creator_settings(project_settings)
        temp_rendering_path_template = (
            plugin_settings.get("temp_rendering_path_template")
            or self.temp_rendering_path_template
        )
        # TODO remove template key replacements
        temp_rendering_path_template = (
            temp_rendering_path_template
            .replace("{product[name]}", "{subset}")
            .replace("{product[type]}", "{family}")
            .replace("{task[name]}", "{task}")
            .replace("{folder[name]}", "{asset}")
        )
        # individual attributes
        self.instance_attributes = plugin_settings.get(
            "instance_attributes") or self.instance_attributes
        self.prenodes = plugin_settings["prenodes"]
        self.default_variants = plugin_settings.get(
            "default_variants") or self.default_variants
        self.render_target = plugin_settings.get(
            "render_target") or self.render_target
        self.temp_rendering_path_template = temp_rendering_path_template

class SelectInstanceNodeAction(pyblish.api.Action):
    """Select instance node for failed plugin."""
    label = "Select instance node"
    on = "failed"  # This action is only available on a failed plug-in
    icon = "mdi.cursor-default-click"

    def process(self, context, plugin):

        # Get the errored instances for the plug-in
        errored_instances = get_errored_instances_from_context(
            context, plugin)

        # Get the invalid nodes for the plug-ins
        self.log.info("Finding instance nodes..")
        nodes = set()
        for instance in errored_instances:
            instance_node = instance.data.get("transientData", {}).get("node")
            if not instance_node:
                raise RuntimeError(
                    "No transientData['node'] found on instance: {}".format(
                        instance
                    )
                )
            nodes.add(instance_node)

        if nodes:
            self.log.info("Selecting instance nodes: {}".format(nodes))
            reset_selection()
            select_nodes(nodes)
        else:
            self.log.info("No instance nodes found.")

class SelectInvalidAction(pyblish.api.Action):
    """Select invalid nodes in Nuke when plug-in failed.

    To retrieve the invalid nodes this assumes a static `get_invalid()`
    method is available on the plugin.

    """
    label = "Select invalid nodes"
    on = "failed"  # This action is only available on a failed plug-in
    icon = "search"  # Icon from Awesome Icon

    def process(self, context, plugin):

        errored_instances = get_errored_instances_from_context(context,
                                                               plugin=plugin)

        # Get the invalid nodes for the plug-ins
        self.log.info("Finding invalid nodes..")
        invalid = set()
        for instance in errored_instances:
            invalid_nodes = plugin.get_invalid(instance)

            if invalid_nodes:
                if isinstance(invalid_nodes, (list, tuple)):
                    invalid.update(invalid_nodes)
                else:
                    self.log.warning("Plug-in returned to be invalid, "
                                     "but has no selectable nodes.")

        if invalid:
            self.log.info("Selecting invalid nodes: {}".format(invalid))
            reset_selection()
            select_nodes(invalid)
        else:
            self.log.info("No invalid nodes found.")