LNXSDK/leenkx/blender/lnx/nodes_logic.py

from typing import Any, Callable
import webbrowser

import bl_operators
import bpy
import blf
from bpy.props import BoolProperty, CollectionProperty, StringProperty

import lnx.logicnode.lnx_nodes as lnx_nodes
import lnx.logicnode.replacement
import lnx.logicnode.tree_variables
import lnx.logicnode.lnx_node_group
import lnx.logicnode
import lnx.props_traits
import lnx.ui_icons as ui_icons
import lnx.utils

if lnx.is_reload(__name__):
    lnx_nodes = lnx.reload_module(lnx_nodes)
    lnx.logicnode.replacement = lnx.reload_module(lnx.logicnode.replacement)
    lnx.logicnode.tree_variables = lnx.reload_module(lnx.logicnode.tree_variables)
    lnx.logicnode = lnx.reload_module(lnx.logicnode)
    lnx.props_traits = lnx.reload_module(lnx.props_traits)
    ui_icons = lnx.reload_module(ui_icons)
    lnx.utils = lnx.reload_module(lnx.utils)
else:
    lnx.enable_reload(__name__)

INTERNAL_GROUPS_MENU_ID = 'LNX_INTERNAL_GROUPS'
internal_groups_menu_class: bpy.types.Menu

registered_nodes = []
registered_categories = []


class LnxLogicTree(bpy.types.NodeTree):
    """Logic nodes"""
    bl_idname = 'LnxLogicTreeType'
    bl_label = 'Logic Node Editor'
    bl_icon = 'NODETREE'

    def update(self):
        pass


class LNX_MT_NodeAddOverride(bpy.types.Menu):
    """
    Overrides the `Add node` menu. If called from the logic node
    editor, the custom menu is drawn, otherwise the default one is drawn.

    TODO: Find a better solution to custom menus, this will conflict
     with other add-ons overriding this menu.
    """
    bl_idname = "NODE_MT_add"
    bl_label = "Add"
    bl_translation_context = bpy.app.translations.contexts.operator_default

    overridden_menu: bpy.types.Menu = None
    overridden_draw: Callable = None

    def draw(self, context):
        if context.space_data.tree_type == 'LnxLogicTreeType':
            layout = self.layout

            # Invoke the search
            layout.operator_context = "INVOKE_DEFAULT"
            layout.operator('lnx.node_search', icon="VIEWZOOM")

            for category_section in lnx_nodes.category_items.values():
                layout.separator()

                for category in category_section:
                    safe_category_name = lnx.utils.safesrc(category.name.lower())
                    layout.menu(f'LNX_MT_{safe_category_name}_menu', text=category.name, icon=category.icon)

            if lnx.logicnode.lnx_node_group.LnxGroupTree.has_linkable_group_trees():
                layout.separator()
                layout.menu(f'LNX_MT_{INTERNAL_GROUPS_MENU_ID}_menu', text=internal_groups_menu_class.bl_label, icon='OUTLINER_OB_GROUP_INSTANCE')

        elif context.space_data.tree_type == 'ShaderNodeTree' and bpy.app.version > (4, 0, 0):
            # TO DO - Recursively gather nodes and draw them to menu

            LNX_MT_NodeAddOverride.overridden_draw(self, context)
            
            layout = self.layout
            layout.separator()
            layout.separator()
            col = layout.column()
            col.label(text="Custom")
            
            shader_data_op = col.operator("node.add_node", text="Shader Data")
            shader_data_op.type = "LnxShaderDataNode"
            shader_data_op.use_transform = True
            
            particle_op = col.operator("node.add_node", text="Custom Particle")
            particle_op.type = "LnxCustomParticleNode"
            particle_op.use_transform = True
            
        else:
            LNX_MT_NodeAddOverride.overridden_draw(self, context)


class LNX_OT_AddNodeOverride(bpy.types.Operator):
    bl_idname = "lnx.add_node_override"
    bl_label = "Add Node"
    bl_property = "type"
    bl_options = {'INTERNAL'}

    type: StringProperty(name="NodeItem type")
    use_transform: BoolProperty(name="Use Transform")
    settings: CollectionProperty(
        name="Settings",
        description="Settings to be applied on the newly created node",
        type=bl_operators.node.NodeSetting,
        options={'SKIP_SAVE'},
    )

    def invoke(self, context, event):
        # Passing collection properties as operator parameters only
        # works via raw sequences of dicts:
        # https://blender.stackexchange.com/a/298977/58208
        # https://github.com/blender/blender/blob/cf1e1ed46b7ec80edb0f43cb514d3601a1696ec1/source/blender/python/intern/bpy_rna.c#L2033-L2043
        setting_dicts = []
        for setting in self.settings.values():
            setting_dicts.append({
                "name": setting.name,
                "value": setting.value,
                "array_index": setting.array_index
            })

        bpy.ops.node.add_node('INVOKE_DEFAULT', type=self.type, use_transform=self.use_transform, settings=setting_dicts)
        return {'FINISHED'}

    @classmethod
    def description(cls, context, properties):
        """Show the node's bl_description attribute as a tooltip or, if
        it doesn't exist, its docstring."""
        nodetype = lnx.utils.type_name_to_type(properties.type)

        if hasattr(nodetype, 'bl_description'):
            return nodetype.bl_description.split('.')[0]

        if nodetype.__doc__ is None:
            return ""

        return nodetype.__doc__.split('.')[0].strip()

    @classmethod
    def poll(cls, context):
        return context.space_data.tree_type == 'LnxLogicTreeType' and context.space_data.edit_tree


def get_category_draw_func(category: lnx_nodes.LnxNodeCategory):
    def draw_category_menu(self, context):
        layout = self.layout

        for index, node_section in enumerate(category.node_sections.values()):
            if index != 0:
                layout.separator()

            for node_item in node_section:
                op = layout.operator("lnx.add_node_override", text=node_item.label)
                op.type = node_item.nodetype
                op.use_transform = True

    return draw_category_menu


def register_nodes():
    global registered_nodes, internal_groups_menu_class

    # Re-register all nodes for now..
    if len(registered_nodes) > 0 or len(registered_categories) > 0:
        unregister_nodes()

    lnx.logicnode.init_nodes(subpackages_only=True)

    for node_type in lnx_nodes.nodes:
        # Don't register internal nodes, they are already registered
        if not issubclass(node_type, bpy.types.NodeInternal):
            registered_nodes.append(node_type)
            bpy.utils.register_class(node_type)

    # Also add Blender's layout nodes
    lnx_nodes.add_node(bpy.types.NodeReroute, 'Layout')
    lnx_nodes.add_node(bpy.types.NodeFrame, 'Layout')

    # Generate and register category menus
    for category_section in lnx_nodes.category_items.values():
        for category in category_section:
            category.sort_nodes()
            safe_category_name = lnx.utils.safesrc(category.name.lower())
            assert(safe_category_name != INTERNAL_GROUPS_MENU_ID)  # see below
            menu_class = type(f'LNX_MT_{safe_category_name}Menu', (bpy.types.Menu, ), {
                'bl_space_type': 'NODE_EDITOR',
                'bl_idname': f'LNX_MT_{safe_category_name}_menu',
                'bl_label': category.name,
                'bl_description': category.description,
                'draw': get_category_draw_func(category)
            })
            registered_categories.append(menu_class)

            bpy.utils.register_class(menu_class)

    # Generate and register group menu
    def draw_nodegroups_menu(self, context):
        layout = self.layout

        tree: lnx.logicnode.lnx_node_group.LnxGroupTree
        for tree in lnx.logicnode.lnx_node_group.LnxGroupTree.get_linkable_group_trees():
            op = layout.operator('lnx.add_node_override', text=tree.name)
            op.type = 'LNCallGroupNode'
            op.use_transform = True
            item = op.settings.add()
            item.name = "group_tree"
            item.value = f'bpy.data.node_groups["{tree.name}"]'

    # Don't name categories like the content of the INTERNAL_GROUPS_MENU_ID variable!
    menu_class = type(f'LNX_MT_{INTERNAL_GROUPS_MENU_ID}Menu', (bpy.types.Menu,), {
        'bl_space_type': 'NODE_EDITOR',
        'bl_idname': f'LNX_MT_{INTERNAL_GROUPS_MENU_ID}_menu',
        'bl_label': 'Node Groups',
        'bl_description': 'List of node groups that can be added to the current tree',
        'draw': draw_nodegroups_menu
    })
    internal_groups_menu_class = menu_class
    bpy.utils.register_class(menu_class)


def unregister_nodes():
    global registered_nodes, registered_categories, internal_groups_menu_class

    for n in registered_nodes:
        if issubclass(n, lnx_nodes.LnxLogicTreeNode):
            n.on_unregister()
        bpy.utils.unregister_class(n)
    registered_nodes = []

    for c in registered_categories:
        bpy.utils.unregister_class(c)
    registered_categories = []

    if internal_groups_menu_class is not None:
        bpy.utils.unregister_class(internal_groups_menu_class)
        internal_groups_menu_class = None


class LNX_PT_LogicNodePanel(bpy.types.Panel):
    bl_label = 'Leenkx Logic Node'
    bl_idname = 'LNX_PT_LogicNodePanel'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'UI'
    bl_category = 'Leenkx'

    @classmethod
    def poll(cls, context):
        return context.space_data.tree_type == 'LnxLogicTreeType' and context.space_data.edit_tree

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False
        if context.active_node is not None and context.active_node.bl_idname.startswith('LN'):
            layout.prop(context.active_node, 'lnx_watch')

            layout.separator()
            layout.operator('lnx.open_node_documentation', icon='HELP')
            column = layout.column(align=True)
            column.operator('lnx.open_node_python_source', icon='FILE_SCRIPT')
            column.operator('lnx.open_node_haxe_source', icon_value=ui_icons.get_id("haxe"))


class LnxOpenNodeHaxeSource(bpy.types.Operator):
    """Expose Haxe source"""
    bl_idname = 'lnx.open_node_haxe_source'
    bl_label = 'Open Node Haxe Source'

    def execute(self, context):
        if context.selected_nodes is not None:
            if len(context.selected_nodes) == 1:
                if context.selected_nodes[0].bl_idname.startswith('LN'):
                    name = context.selected_nodes[0].bl_idname[2:]
                    version = lnx.utils.get_last_commit()
                    if version == '':
                        version = 'main'
                    webbrowser.open(f'https://dev.leenkx.com/LeenkxTeam/LNXSDK/src/branch/{version}/leenkx/Sources/leenkx/logicnode/{name}.hx')
        return{'FINISHED'}


class LnxOpenNodePythonSource(bpy.types.Operator):
    """Expose Python source"""
    bl_idname = 'lnx.open_node_python_source'
    bl_label = 'Open Node Python Source'

    def execute(self, context):
        if context.selected_nodes is not None:
            if len(context.selected_nodes) == 1:
                node = context.selected_nodes[0]
                if node.bl_idname.startswith('LN') and node.lnx_version is not None:
                    version = lnx.utils.get_last_commit()
                    if version == '':
                        version = 'main'
                    rel_path = node.__module__.replace('.', '/')
                    webbrowser.open(f'https://dev.leenkx.com/LeenkxTeam/LNXSDK/src/branch/{version}/leenkx/blender/{rel_path}.py')
        return{'FINISHED'}


class LnxOpenNodeWikiEntry(bpy.types.Operator):
    """Open the logic node's documentation in the Leenkx wiki"""
    bl_idname = 'lnx.open_node_documentation'
    bl_label = 'Open Node Documentation'

    def execute(self, context):
        if context.selected_nodes is not None:
            if len(context.selected_nodes) == 1:
                node = context.selected_nodes[0]
                if node.bl_idname.startswith('LN') and node.lnx_version is not None:
                    anchor = node.bl_label.lower().replace(" ", "-")

                    category = lnx_nodes.eval_node_category(node)
                    category_section = lnx_nodes.get_category(category).category_section

                    webbrowser.open(f'https://github.com/leenkx3d/leenkx/wiki/reference_{category_section}#{anchor}')

        return {'FINISHED'}


class LNX_PT_NodeDevelopment(bpy.types.Panel):
    """Sidebar panel to ease development of logic nodes."""
    bl_label = 'Node Development'
    bl_idname = 'LNX_PT_NodeDevelopment'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'UI'
    bl_category = 'Leenkx'

    @classmethod
    def poll(cls, context):
        return context.space_data.tree_type == 'LnxLogicTreeType' and context.space_data.edit_tree

    def draw(self, context):
        layout = self.layout
        layout.use_property_split = True
        layout.use_property_decorate = False

        node = context.active_node
        if node is not None and node.bl_idname.startswith('LN'):
            box = layout.box()
            box.label(text='Selected Node')
            col = box.column(align=True)

            self._draw_row(col, 'bl_idname', node.bl_idname)
            self._draw_row(col, 'Category', lnx_nodes.eval_node_category(node))
            self._draw_row(col, 'Section', node.lnx_section)
            self._draw_row(col, 'Specific Version', node.lnx_version)
            self._draw_row(col, 'Class Version', node.__class__.lnx_version)
            self._draw_row(col, 'Is Deprecated', node.lnx_is_obsolete)

            is_var_node = isinstance(node, lnx_nodes.LnxLogicVariableNodeMixin)
            self._draw_row(col, 'Is Variable Node', is_var_node)
            self._draw_row(col, 'Logic ID', node.lnx_logic_id)
            if is_var_node:
                self._draw_row(col, 'Is Master Node', node.is_master_node)

            layout.separator()
            layout.operator('lnx.node_replace_all')
            layout.operator('lnx.recalculate_rotations')


    @staticmethod
    def _draw_row(col: bpy.types.UILayout, text: str, val: Any):
        split = col.split(factor=0.4)
        split.label(text=text)
        split.label(text=str(val))


class LNX_OT_ReplaceNodesOperator(bpy.types.Operator):
    bl_idname = "lnx.node_replace_all"
    bl_label = "Replace Deprecated Nodes"
    bl_description = "Replace all deprecated nodes in the active node tree"
    bl_options = {'REGISTER'}

    def execute(self, context):
        lnx.logicnode.replacement.replace_all()
        return {'FINISHED'}

    @classmethod
    def poll(cls, context):
        return context.space_data is not None and context.space_data.type == 'NODE_EDITOR'

class LNX_OT_RecalculateRotations(bpy.types.Operator):
    """Recalculates internal rotation values for all rotation sockets in the tree"""
    bl_idname = "lnx.recalculate_rotations"
    bl_label = "Recalculate Rotations"
    bl_description = "Forces recalculation of internal quaternion values for all LnxRotationSockets in the active tree using their current settings. Useful for fixing old files."
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        return (context.space_data is not None and
                context.space_data.type == 'NODE_EDITOR' and
                context.space_data.tree_type == 'LnxLogicTreeType' and
                context.space_data.edit_tree is not None)

    def execute(self, context):
        tree = context.space_data.edit_tree
        if not tree:
            self.report({'WARNING'}, "No active Logic Node tree found")
            return {'CANCELLED'}
        recalculated_count = 0
        for node in tree.nodes:
            for socket in list(node.inputs) + list(node.outputs):
                if hasattr(socket, 'do_update_raw') and callable(socket.do_update_raw):
                    try:
                        socket.do_update_raw(context)
                        recalculated_count += 1
                    except Exception as e:
                        print(f"Error recalculating socket '{socket.name}' on node '{node.name}': {e}")

        self.report({'INFO'}, f"Recalculated {recalculated_count} rotation sockets in tree '{tree.name}'")
        tree.update_tag()
        return {'FINISHED'}


class LNX_UL_InterfaceSockets(bpy.types.UIList):
    """UI List of input and output sockets"""
    def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
        socket = item
        color = socket.draw_color(context, context.active_node)

        if self.layout_type in {'DEFAULT', 'COMPACT'}:
            row = layout.row(align=True)

            row.template_node_socket(color=color)
            row.prop(socket, "display_label", text="", emboss=False, icon_value=icon)
        elif self.layout_type == 'GRID':
            layout.alignment = 'CENTER'
            layout.template_node_socket(color=color)

class DrawNodeBreadCrumbs():
    """A class to draw node tree breadcrumbs or context path"""
    draw_handler = None

    @classmethod
    def convert_array_to_string(cls, arr):
        return ' > '.join(arr)

    @classmethod
    def draw(cls, context):
        if  context.space_data.edit_tree and context.space_data.node_tree.bl_idname == "LnxLogicTreeType":
            height = context.area.height
            path_data = [path.node_tree.name for path in context.space_data.path]
            str = cls.convert_array_to_string(path_data)
            blf.position(0, 20, height-60, 0)
            if bpy.app.version < (4, 1, 0):
                blf.size(0, 15, 72)
            else:
                blf.size(15, 72)
            blf.draw(0, str)

    @classmethod
    def register_draw(cls):
        if cls.draw_handler is not None:
            cls.unregister_draw()
        cls.draw_handler = bpy.types.SpaceNodeEditor.draw_handler_add(cls.draw, tuple([bpy.context]), 'WINDOW', 'POST_PIXEL')

    @classmethod
    def unregister_draw(cls):
        if cls.draw_handler is not None:
            bpy.types.SpaceNodeEditor.draw_handler_remove(cls.draw_handler, 'WINDOW')
            cls.draw_handler = None


__REG_CLASSES = (
    LnxLogicTree,
    LnxOpenNodeHaxeSource,
    LnxOpenNodePythonSource,
    LnxOpenNodeWikiEntry,
    LNX_OT_ReplaceNodesOperator,
    LNX_OT_RecalculateRotations,
    LNX_MT_NodeAddOverride,
    LNX_OT_AddNodeOverride,
    LNX_UL_InterfaceSockets,
    LNX_PT_LogicNodePanel,
    LNX_PT_NodeDevelopment
)
__reg_classes, __unreg_classes = bpy.utils.register_classes_factory(__REG_CLASSES)


def register():
    lnx.logicnode.lnx_nodes.register()
    lnx.logicnode.lnx_sockets.register()
    lnx.logicnode.lnx_node_group.register()
    lnx.logicnode.tree_variables.register()

    LNX_MT_NodeAddOverride.overridden_menu = bpy.types.NODE_MT_add
    LNX_MT_NodeAddOverride.overridden_draw = bpy.types.NODE_MT_add.draw

    __reg_classes()

    lnx.logicnode.init_categories()
    DrawNodeBreadCrumbs.register_draw()
    register_nodes()


def unregister():
    unregister_nodes()
    DrawNodeBreadCrumbs.unregister_draw()
    # Ensure that globals are reset if the addon is enabled again in the same Blender session
    lnx_nodes.reset_globals()

    __unreg_classes()
    bpy.utils.register_class(LNX_MT_NodeAddOverride.overridden_menu)

    lnx.logicnode.tree_variables.unregister()
    lnx.logicnode.lnx_node_group.unregister()
    lnx.logicnode.lnx_sockets.unregister()
    lnx.logicnode.lnx_nodes.unregister()