LNXSDK/leenkx/blender/lnx/handlers.py

import importlib
import os
import queue
import sys
import threading
import time
import types
from typing import Dict, Tuple, Callable, Set

import bpy
from bpy.app.handlers import persistent

import lnx
import lnx.api
import lnx.nodes_logic
import lnx.make_state as state
import lnx.utils
import lnx.utils_vs
from lnx import live_patch, log, make, props
from lnx.logicnode import lnx_nodes

if lnx.is_reload(__name__):
    lnx.api = lnx.reload_module(lnx.api)
    live_patch = lnx.reload_module(live_patch)
    log = lnx.reload_module(log)
    lnx_nodes = lnx.reload_module(lnx_nodes)
    lnx.nodes_logic = lnx.reload_module(lnx.nodes_logic)
    make = lnx.reload_module(make)
    state = lnx.reload_module(state)
    props = lnx.reload_module(props)
    lnx.utils = lnx.reload_module(lnx.utils)
    lnx.utils_vs = lnx.reload_module(lnx.utils_vs)
else:
    lnx.enable_reload(__name__)

# Module-level storage for active threads (eliminates re-queuing overhead)
_active_threads: Dict[threading.Thread, Callable] = {}
_last_poll_time = 0.0
_consecutive_empty_polls = 0

@persistent
def on_depsgraph_update_post(self):
    if state.proc_build is not None:
        return

    # Recache
    depsgraph = bpy.context.evaluated_depsgraph_get()

    for update in depsgraph.updates:
        uid = update.id
        if hasattr(uid, 'lnx_cached'):
            # uid.lnx_cached = False # TODO: does not trigger update
            if isinstance(uid, bpy.types.Mesh) and uid.name in bpy.data.meshes:
                bpy.data.meshes[uid.name].lnx_cached = False
            elif isinstance(uid, bpy.types.Curve) and uid.name in bpy.data.curves:
                bpy.data.curves[uid.name].lnx_cached = False
            elif isinstance(uid, bpy.types.MetaBall) and uid.name in bpy.data.metaballs:
                bpy.data.metaballs[uid.name].lnx_cached = False
            elif isinstance(uid, bpy.types.Armature) and uid.name in bpy.data.armatures:
                bpy.data.armatures[uid.name].lnx_cached = False
            elif isinstance(uid, bpy.types.NodeTree) and uid.name in bpy.data.node_groups:
                bpy.data.node_groups[uid.name].lnx_cached = False
            elif isinstance(uid, bpy.types.Material) and uid.name in bpy.data.materials:
                bpy.data.materials[uid.name].lnx_cached = False

    # Send last operator to Krom
    wrd = bpy.data.worlds['Lnx']
    if state.proc_play is not None and state.target == 'krom' and wrd.lnx_live_patch:
        ops = bpy.context.window_manager.operators
        if len(ops) > 0 and ops[-1] is not None:
            live_patch.on_operator(ops[-1].bl_idname)

    # Hacky solution to update leenkx props after operator executions.
    # bpy.context.active_operator doesn't always exist, in some cases
    # like marking assets for example, this code is also executed before
    # the operator actually finishes and sets the variable
    last_operator = getattr(bpy.context, 'active_operator', None)
    if last_operator is not None:
        on_operator_post(last_operator.bl_idname)


def on_operator_post(operator_id: str) -> None:
    """Called after operator execution. Does not work for operators
    executed in another context. Warning: this function is also called
    when the operator execution raised an exception!"""
    # 3D View > Object > Rigid Body > Copy from Active
    if operator_id == "RIGIDBODY_OT_object_settings_copy":
        # Copy leenkx rigid body settings
        source_obj = bpy.context.active_object
        for target_obj in bpy.context.selected_objects:
            target_obj.lnx_rb_linear_factor = source_obj.lnx_rb_linear_factor
            target_obj.lnx_rb_angular_factor = source_obj.lnx_rb_angular_factor
            target_obj.lnx_rb_angular_friction = source_obj.lnx_rb_angular_friction
            target_obj.lnx_rb_trigger = source_obj.lnx_rb_trigger
            target_obj.lnx_rb_deactivation_time = source_obj.lnx_rb_deactivation_time
            target_obj.lnx_rb_ccd = source_obj.lnx_rb_ccd
            target_obj.lnx_rb_interpolate = source_obj.lnx_rb_interpolate
            target_obj.lnx_rb_collision_filter_mask = source_obj.lnx_rb_collision_filter_mask

    elif operator_id == "NODE_OT_new_node_tree":
        if bpy.context.space_data.tree_type == lnx.nodes_logic.LnxLogicTree.bl_idname:
            # In Blender 3.5+, new node trees are no longer called "NodeTree"
            # but follow the bl_label attribute by default. New logic trees
            # are thus called "Leenkx Logic Editor" which conflicts with Haxe's
            # class naming convention. To avoid this, we listen for the
            # creation of a node tree and then rename it.
            # Unfortunately, manually naming the tree has the unfortunate
            # side effect of not basing the new name on the name of the
            # previously opened node tree, as it is the case for Blender trees...
            bpy.context.space_data.edit_tree.name = "LogicTree"


def send_operator(op):
    if hasattr(bpy.context, 'object') and bpy.context.object is not None:
        obj = bpy.context.object.name
        if op.name == 'Move':
            vec = bpy.context.object.location
            js = 'var o = iron.Scene.active.getChild("' + obj + '"); o.transform.loc.set(' + str(vec[0]) + ', ' + str(vec[1]) + ', ' + str(vec[2]) + '); o.transform.dirty = true;'
            make.write_patch(js)
        elif op.name == 'Resize':
            vec = bpy.context.object.scale
            js = 'var o = iron.Scene.active.getChild("' + obj + '"); o.transform.scale.set(' + str(vec[0]) + ', ' + str(vec[1]) + ', ' + str(vec[2]) + '); o.transform.dirty = true;'
            make.write_patch(js)
        elif op.name == 'Rotate':
            vec = bpy.context.object.rotation_euler.to_quaternion()
            js = 'var o = iron.Scene.active.getChild("' + obj + '"); o.transform.rot.set(' + str(vec[1]) + ', ' + str(vec[2]) + ', ' + str(vec[3]) + ' ,' + str(vec[0]) + '); o.transform.dirty = true;'
            make.write_patch(js)
        else: # Rebuild
            make.patch()


def always() -> float:
    # Force ui redraw
    if state.redraw_ui:
        for area in bpy.context.screen.areas:
            if area.type in ('NODE_EDITOR', 'PROPERTIES', 'VIEW_3D'):
                area.tag_redraw()
        state.redraw_ui = False

    return 0.5



def poll_threads() -> float:
    """
    Improved thread polling with:
    - No re-queuing overhead
    - Batch processing of completed threads
    - Adaptive timing based on activity
    - Better memory management
    - Simplified logic flow
    """
    global _last_poll_time, _consecutive_empty_polls
    current_time = time.time()
    
    # Process all new threads from queue at once (batch processing)
    new_threads_added = 0
    try:
        while True:
            thread, callback = make.thread_callback_queue.get(block=False)
            _active_threads[thread] = callback
            new_threads_added += 1
    except queue.Empty:
        pass
    
    # Early return if no active threads
    if not _active_threads:
        _consecutive_empty_polls += 1
        # Adaptive timing: longer intervals when consistently empty
        if _consecutive_empty_polls > 10:
            return 0.5  # Back off when no activity
        return 0.25
    
    # Reset empty poll counter when we have active threads
    _consecutive_empty_polls = 0
    
    # Find completed threads (single pass, no re-queuing)
    completed_threads = []
    for thread in list(_active_threads.keys()):
        if not thread.is_alive():
            completed_threads.append(thread)
    
    # Batch process all completed threads
    if completed_threads:
        _process_completed_threads(completed_threads)
    
    # Adaptive timing based on activity level
    active_count = len(_active_threads)
    if active_count == 0:
        return 0.25
    elif active_count <= 3:
        return 0.05  # Medium frequency for low activity
    else:
        return 0.01  # High frequency for high activity

def _process_completed_threads(completed_threads: list) -> None:
    """Process a batch of completed threads with robust error handling."""
    for thread in completed_threads:
        callback = _active_threads.pop(thread)  # Remove from tracking
        
        try:
            thread.join()  # Should be instant since thread is dead
            callback()
        except Exception as e:
            # Robust error recovery
            _handle_callback_error(e)
            continue  # Continue processing other threads
        
        # Explicit cleanup for better memory management
        del thread, callback

def _handle_callback_error(exception: Exception) -> None:
    """Centralized error handling with better recovery."""
    try:
        # Try to unregister existing timer
        bpy.app.timers.unregister(poll_threads)
    except ValueError:
        pass  # Timer wasn't registered, that's fine
    
    # Re-register timer with slightly longer interval for stability
    bpy.app.timers.register(poll_threads, first_interval=0.1, persistent=True)
    
    # Re-raise the original exception after ensuring timer continuity
    raise exception

def cleanup_polling_system() -> None:
    """Optional cleanup function for proper shutdown."""
    global _active_threads, _consecutive_empty_polls
    
    # Wait for remaining threads to complete (with timeout)
    for thread in list(_active_threads.keys()):
        if thread.is_alive():
            thread.join(timeout=1.0)  # 1 second timeout
    
    # Clear tracking structures
    _active_threads.clear()
    _consecutive_empty_polls = 0
    
    # Unregister timer
    try:
        bpy.app.timers.unregister(poll_threads)
    except ValueError:
        pass

def get_polling_stats() -> dict:
    """Get statistics about the polling system for monitoring."""
    return {
        'active_threads': len(_active_threads),
        'consecutive_empty_polls': _consecutive_empty_polls,
        'thread_ids': [t.ident for t in _active_threads.keys()]
    }


loaded_py_libraries: dict[str, types.ModuleType] = {}
context_screen = None


@persistent
def on_save_pre(context):
    # Ensure that files are saved with the correct version number
    # (e.g. startup files with an "Arm" world may have old version numbers)
    wrd = bpy.data.worlds['Lnx']
    wrd.lnx_version = props.lnx_version
    wrd.lnx_commit = props.lnx_commit


@persistent
def on_load_pre(context):
    unload_py_libraries()
    log.clear(clear_warnings=True, clear_errors=True)


@persistent
def on_load_post(context):
    global context_screen
    context_screen = bpy.context.screen

    props.init_properties_on_load()
    reload_blend_data()
    lnx.utils.fetch_bundled_script_names()

    wrd = bpy.data.worlds['Lnx']
    wrd.lnx_recompile = True
    lnx.api.remove_drivers()

    load_py_libraries()

    # Show trait users as collections
    lnx.utils.update_trait_collections()
    props.update_leenkx_world()


def load_py_libraries():
    if bpy.data.filepath == '':
        # When a blend file is opened from the file explorer, Blender
        # first opens the default file and then the actual blend file,
        # so this function is called twice. Because the cwd is already
        # that of the folder containing the blend file, libraries would
        # be loaded/unloaded once for the default file which is not needed.
        return

    lib_path = os.path.join(lnx.utils.get_fp(), 'Libraries')
    if os.path.exists(lib_path):
        # Don't register nodes twice when calling register_nodes()
        lnx_nodes.reset_globals()

        # Make sure that Leenkx's categories are registered first (on top of the menu)
        lnx.logicnode.init_categories()

        libs = os.listdir(lib_path)
        for lib_name in libs:
            fp = os.path.join(lib_path, lib_name)
            if os.path.isdir(fp):
                if os.path.exists(os.path.join(fp, 'blender.py')):
                    sys.path.append(fp)

                    lib_module = importlib.import_module('blender')
                    importlib.reload(lib_module)
                    if hasattr(lib_module, 'register'):
                        lib_module.register()

                    log.debug(f'Leenkx: Loaded Python library {lib_name}')
                    loaded_py_libraries[lib_name] = lib_module

                    sys.path.remove(fp)

        # Register newly added nodes and node categories
        lnx.nodes_logic.register_nodes()


def unload_py_libraries():
    for lib_name, lib_module in loaded_py_libraries.items():
        if hasattr(lib_module, 'unregister'):
            lib_module.unregister()
        lnx.log.debug(f'Leenkx: Unloaded Python library {lib_name}')

    loaded_py_libraries.clear()


def reload_blend_data():
    leenkx_pbr = bpy.data.node_groups.get('Leenkx PBR')
    if leenkx_pbr is None:
        load_library('Leenkx PBR')
    custom_tilesheet = bpy.data.node_groups.get('CustomTilesheet')
    if custom_tilesheet is None:
        load_library('CustomTilesheet')


def load_library(asset_name):
    if bpy.data.filepath.endswith('lnx_data.blend'): # Prevent load in library itself
        return
    sdk_path = lnx.utils.get_sdk_path()
    data_path = sdk_path + '/leenkx/blender/data/lnx_data.blend'
    data_names = [asset_name]

    # Import
    data_refs = data_names.copy()
    with bpy.data.libraries.load(data_path, link=False) as (data_from, data_to):
        data_to.node_groups = data_refs

    for ref in data_refs:
        ref.use_fake_user = True


def post_register():
    """Called in start.py after all Leenkx modules have been registered.
    It is also called in case of add-on reloads. Put code here that
    needs to be run once at the beginning of each session.
    """
    if lnx.utils.get_os_is_windows():
        lnx.utils_vs.fetch_installed_vs(silent=True)


def register():
    bpy.app.handlers.save_pre.append(on_save_pre)
    bpy.app.handlers.load_pre.append(on_load_pre)
    bpy.app.handlers.load_post.append(on_load_post)
    bpy.app.handlers.depsgraph_update_post.append(on_depsgraph_update_post)
    # bpy.app.handlers.undo_post.append(on_undo_post)

    bpy.app.timers.register(always, persistent=True)
    bpy.app.timers.register(poll_threads, persistent=True)

    if lnx.utils.get_fp() != '':
        # TODO: On windows, on_load_post is not called when opening .blend file from explorer
        if lnx.utils.get_os() == 'win':
            on_load_post(None)
        else:
            # load_py_libraries() is called by on_load_post(). This call makes sure that libraries are also loaded
            # when a file is already opened during add-on registration
            load_py_libraries()

    reload_blend_data()


def unregister():
    unload_py_libraries()

    bpy.app.timers.unregister(poll_threads)
    bpy.app.timers.unregister(always)

    bpy.app.handlers.load_post.remove(on_load_post)
    bpy.app.handlers.load_pre.remove(on_load_pre)
    bpy.app.handlers.save_pre.remove(on_save_pre)
    bpy.app.handlers.depsgraph_update_post.remove(on_depsgraph_update_post)
    # bpy.app.handlers.undo_post.remove(on_undo_post)