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, load
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_options = {'REGISTER', 'UNDO'}

    select_only: BoolProperty(default=True)

    def execute(self, context):

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

        pprint(ntree_data)

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

        return {'FINISHED'}


class NODEKIT_OT_paste(Operator):
    bl_idname = 'node_kit.paste_node_tree'
    bl_label = 'Paste nodes'

    def execute(self, context):

        ntree_data = json.loads(context.window_manager.clipboard)
        load(ntree_data, context.space_data.edit_tree)

        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)