node_kit/operators.py

"""
This module contains all addons operators

:author: Autour de Minuit
:maintainers: Florentin LUCE
:date: 2024
"""

import json
from pprint import pprint
from pathlib import Path

import bpy
from bpy.props import BoolProperty, EnumProperty
from bpy.types import Operator

# from node_kit.core.node_tree import NodeTree
from .core.dumper import dump_nodes, load_nodes
from .core.node_utils import remap_node_group_duplicates
from .core.pack_nodes import combine_objects, extract_objects


class NODEKIT_OT_copy(Operator):
    bl_idname = "node_kit.copy_node_tree"
    bl_label = "Copy nodes"
    bl_description = "Copy nodes to system clipboard"
    bl_options = {"REGISTER", "UNDO"}

    select_only: BoolProperty(default=True)  # TODO: Expose/rework this property properly - No F3 panel in Node Editor - Only F9

    def execute(self, context):
        ntree = context.space_data.edit_tree
        if self.select_only:
            ntree_data = {
                "nodes": dump_nodes(
                    [n for n in ntree.nodes if n.select]
                ),  # [dump(n) for n in ntree.nodes if n.select],
                "links": dump_nodes(
                    [l for l in ntree.links if l.from_node.select and l.to_node.select]
                ),
            }
        else:
            ntree_data = dump_nodes(ntree)

        pprint(ntree_data)

        context.window_manager.clipboard = json.dumps(ntree_data)

        self.report({"INFO"}, f"Copied 5 selected nodes to system clipboard")
        return {"FINISHED"}


class NODEKIT_OT_paste(Operator):
    bl_idname = "node_kit.paste_node_tree"
    bl_label = "Paste nodes"
    bl_description = "Paste nodes from system clipboard"
    bl_options = {"REGISTER", "UNDO"}

    def execute(self, context):
        ntree_data = json.loads(context.window_manager.clipboard)
        load_nodes(ntree_data, context.space_data.edit_tree)

        self.report({"INFO"}, f"5 node(s) pasted from system clipboard")
        return {"FINISHED"}


class NODEKIT_OT_remap_node_group_duplicates(Operator):
    bl_idname = "node_kit.remap_node_group_duplicates"
    bl_label = "Clean nodes"
    bl_options = {"REGISTER", "UNDO"}

    selection: EnumProperty(
        items=[(s, s.title(), "") for s in ("ALL", "SELECTED", "CURRENT")],
        default="CURRENT",
        name="All Nodes",
    )
    force: BoolProperty(
        default=False,
        description="Remap nodes even if there are different",
        name="Force",
    )

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        nodes = None
        if self.selection == "SELECTED":
            nodes = [
                n.node_tree
                for n in context.space_data.edit_tree.nodes
                if n.type == "GROUP" and n.select
            ]
        elif self.selection == "ACTIVE":
            active_node = context.space_data.edit_tree
            nodes = [active_node]

        merged, failed = remap_node_group_duplicates(nodes=nodes, force=self.force)

        if failed and not merged:
            self.report({"ERROR"}, "No duplicates remapped, Node Group are differents")
            return {"CANCELLED"}

        self.report(
            {"INFO"},
            f"{len(merged)} Node Groups Remapped, {len(failed)} Node Groups failed",
        )

        return {"FINISHED"}

    def draw(self, context):
        layout = self.layout
        layout.prop(self, "selection", expand=True)
        layout.prop(self, "force")
        if self.force and self.selection == "CURRENT":
            ntree = context.space_data.edit_tree
            layout.label(text=f"Remap node {ntree.name} to others")
        elif self.force and self.selection == "SELECTED":
            layout.label(text="Selected nodes will override others")
        elif self.selection == "SELECTED":
            layout.label(text="Remap last .*** nodes")
            layout.label(text="Ex: Node.001 will override Node")
        elif self.selection in ("CURRENT", "ALL"):
            layout.label(text="Remap last .*** nodes")
            layout.label(text="Ex: Node.001 will override Node")


class NODEKIT_OT_update_nodes(Operator):
    bl_idname = "node_kit.update_nodes"
    bl_label = "Update node"
    bl_options = {"REGISTER", "UNDO"}

    selection: EnumProperty(
        items=[(s, s.title(), "") for s in ("ALL", "SELECTED", "ACTIVE")],
        default="ACTIVE",
        name="All Nodes",
    )

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        asset_libraries = context.preferences.filepaths.asset_libraries

        ntree = context.space_data.edit_tree
        ntree_name = ntree.name
        new_ntree = None

        if self.selection == "SELECTED":
            nodes = [
                n.node_tree
                for n in context.space_data.edit_tree.nodes
                if n.type == "GROUP" and n.select
            ]
        elif self.selection == "ACTIVE":
            active_node = context.space_data.edit_tree
            nodes = [active_node]
        else:
            nodes = list(bpy.data.node_groups)

        node_names = set(n.name for n in nodes)
        # new_node_groups = []

        # print("node_names", node_names)

        for asset_library in asset_libraries:
            library_path = Path(asset_library.path)
            blend_files = [fp for fp in library_path.glob("**/*.blend") if fp.is_file()]

            node_groups = list(
                bpy.data.node_groups
            )  # Storing original node_geoup to compare with imported

            link = asset_library.import_method == "LINK"
            for blend_file in blend_files:
                print(blend_file)
                with bpy.data.libraries.load(
                    str(blend_file), assets_only=True, link=link
                ) as (data_from, data_to):
                    print(data_from.node_groups)
                    import_node_groups = [
                        n for n in data_from.node_groups if n in node_names
                    ]
                    print("import_node_groups", import_node_groups)
                    data_to.node_groups = import_node_groups

                    # print(data_from.node_groups)
                    # print("data_to.node_groups", data_to.node_groups)
                    node_names -= set(import_node_groups)  # Store already updated nodes

                # new_ntree = data_to.node_groups[0]
        new_node_groups = [n for n in bpy.data.node_groups if n not in node_groups]

        # break

        # if new_ntree:
        #    break
        new_node_groups = list(set(new_node_groups))
        # print(new_node_groups)

        # if new_node_groups:
        for new_node_group in new_node_groups:
            new_node_group_name = new_node_group.library_weak_reference.id_name[2:]
            local_node_group = next(
                (
                    n
                    for n in bpy.data.node_groups
                    if n.name == new_node_group_name and n != new_node_group
                ),
                None,
            )

            if not local_node_group:
                print(f"No local node_group {new_node_group_name}")
                continue

            print(f"Merge node {local_node_group.name} into {new_node_group.name}")

            local_node_group.user_remap(new_node_group)
            new_node_group.interface_update(context)
            bpy.data.node_groups.remove(local_node_group)

            new_node_group.name = new_node_group_name
            new_node_group.asset_clear()

        # self.report({"INFO"}, f"Node updated from {blend_file}")
        return {"FINISHED"}

        # else:
        #     self.report({"ERROR"}, f'No Node Group named "{ntree_name}" in the library')
        #     return {'CANCELLED'}

    def draw(self, context):
        layout = self.layout
        layout.prop(self, "selection", expand=True)


class NODEKIT_OT_pack_nodes(Operator):
    bl_idname = "node_kit.pack_nodes"
    bl_label = "Update node"
    bl_options = {"REGISTER", "UNDO"}

    def execute(self, context):
        combine_objects(context.selected_objects)
        return {"FINISHED"}


class NODEKIT_OT_unpack_nodes(Operator):
    bl_idname = "node_kit.unpack_nodes"
    bl_label = "Update node"
    bl_options = {"REGISTER", "UNDO"}

    def execute(self, context):
        extract_objects(context.active_object)
        return {"FINISHED"}


classes = (
    NODEKIT_OT_copy,
    NODEKIT_OT_paste,
    NODEKIT_OT_remap_node_group_duplicates,
    NODEKIT_OT_update_nodes,
    NODEKIT_OT_pack_nodes,
    NODEKIT_OT_unpack_nodes,
)


def register():
    for c in classes:
        bpy.utils.register_class(c)


def unregister():
    for c in reversed(classes):
        bpy.utils.unregister_class(c)