import sys
import inspect
import logging
from .metadata_registry import MetadataRegistry

logger = logging.getLogger(__name__)

class MetadataHook:
    """Install hooks for metadata collection"""
    
    @staticmethod
    def install():
        """Install hooks to collect metadata during execution"""
        try:
            # Import ComfyUI's execution module
            execution = None
            try:
                # Try direct import first
                import execution # pyright: ignore[reportMissingImports]
            except ImportError:
                # Try to locate from system modules
                for module_name in sys.modules:
                    if module_name.endswith('.execution'):
                        execution = sys.modules[module_name]
                        break
                    
            # If we can't find the execution module, we can't install hooks
            if execution is None:
                logger.warning("Could not locate ComfyUI execution module, metadata collection disabled")
                return
            
            # Detect whether we're using the new async version of ComfyUI
            is_async = False
            map_node_func_name = '_map_node_over_list'
            
            if hasattr(execution, '_async_map_node_over_list'):
                is_async = inspect.iscoroutinefunction(execution._async_map_node_over_list)
                map_node_func_name = '_async_map_node_over_list'
            elif hasattr(execution, '_map_node_over_list'):
                is_async = inspect.iscoroutinefunction(execution._map_node_over_list)
            
            if is_async:
                logger.info("Detected async ComfyUI execution, installing async metadata hooks")
                MetadataHook._install_async_hooks(execution, map_node_func_name)
            else:
                logger.info("Detected sync ComfyUI execution, installing sync metadata hooks")
                MetadataHook._install_sync_hooks(execution)
            
            logger.info("Metadata collection hooks installed for runtime values")
            
        except Exception as e:
            logger.error(f"Error installing metadata hooks: {str(e)}")
    
    @staticmethod
    def _install_sync_hooks(execution):
        """Install hooks for synchronous execution model"""
        # Store the original _map_node_over_list function
        original_map_node_over_list = execution._map_node_over_list
        
        # Define the wrapped _map_node_over_list function
        def map_node_over_list_with_metadata(obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None):
            # Only collect metadata when calling the main function of nodes
            if func == obj.FUNCTION and hasattr(obj, '__class__'):
                try:
                    # Get the current prompt_id from the registry
                    registry = MetadataRegistry()
                    prompt_id = registry.current_prompt_id
                    
                    if prompt_id is not None:
                        # Get node class type
                        class_type = obj.__class__.__name__
                        
                        # Unique ID might be available through the obj if it has a unique_id field
                        node_id = getattr(obj, 'unique_id', None)
                        if node_id is None and pre_execute_cb:
                            # Try to extract node_id through reflection on GraphBuilder.set_default_prefix
                            frame = inspect.currentframe()
                            while frame:
                                if 'unique_id' in frame.f_locals:
                                    node_id = frame.f_locals['unique_id']
                                    break
                                frame = frame.f_back
                        
                        # Record inputs before execution
                        if node_id is not None:
                            return_types = getattr(obj, 'RETURN_TYPES', None)
                            registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
                except Exception as e:
                    logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
            
            # Execute the original function
            results = original_map_node_over_list(obj, input_data_all, func, allow_interrupt, execution_block_cb, pre_execute_cb)
            
            # After execution, collect outputs for relevant nodes
            if func == obj.FUNCTION and hasattr(obj, '__class__'):
                try:
                    # Get the current prompt_id from the registry
                    registry = MetadataRegistry()
                    prompt_id = registry.current_prompt_id
                    
                    if prompt_id is not None:
                        # Get node class type
                        class_type = obj.__class__.__name__
                        
                        # Unique ID might be available through the obj if it has a unique_id field
                        node_id = getattr(obj, 'unique_id', None)
                        if node_id is None and pre_execute_cb:
                            # Try to extract node_id through reflection
                            frame = inspect.currentframe()
                            while frame:
                                if 'unique_id' in frame.f_locals:
                                    node_id = frame.f_locals['unique_id']
                                    break
                                frame = frame.f_back
                        
                        # Record outputs after execution
                        if node_id is not None:
                            return_types = getattr(obj, 'RETURN_TYPES', None)
                            registry.update_node_execution(node_id, class_type, results, return_types=return_types)
                except Exception as e:
                    logger.error(f"Error collecting metadata (post-execution): {str(e)}")
            
            return results
            
        # Also hook the execute function to track the current prompt_id
        original_execute = execution.execute
        
        def execute_with_prompt_tracking(*args, **kwargs):
            if len(args) >= 7:  # Check if we have enough arguments
                server, prompt, caches, node_id, extra_data, executed, prompt_id = args[:7]
                registry = MetadataRegistry()
                
                # Start collection if this is a new prompt
                if not registry.current_prompt_id or registry.current_prompt_id != prompt_id:
                    registry.start_collection(prompt_id)
                    
                # Store the dynprompt reference for node lookups
                if hasattr(prompt, 'original_prompt'):
                    registry.set_current_prompt(prompt)

                # Store extra_data for accessing full workflow node properties
                registry.set_extra_data(extra_data)

            # Execute the original function
            return original_execute(*args, **kwargs)

        # Replace the functions
        execution._map_node_over_list = map_node_over_list_with_metadata
        execution.execute = execute_with_prompt_tracking
    
    @staticmethod
    def _install_async_hooks(execution, map_node_func_name='_async_map_node_over_list'):
        """Install hooks for asynchronous execution model"""
        # Store the original _async_map_node_over_list function
        original_map_node_over_list = getattr(execution, map_node_func_name)

        # Wrapped async function - signature must exactly match _async_map_node_over_list
        async def async_map_node_over_list_with_metadata(
            prompt_id, unique_id, obj, input_data_all, func,
            allow_interrupt=False, execution_block_cb=None,
            pre_execute_cb=None, v3_data=None
        ):
            # Only collect metadata when calling the main function of nodes
            if func == obj.FUNCTION and hasattr(obj, '__class__'):
                try:
                    registry = MetadataRegistry()
                    if prompt_id is not None:
                        class_type = obj.__class__.__name__
                        node_id = unique_id
                        if node_id is not None:
                            return_types = getattr(obj, 'RETURN_TYPES', None)
                            registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
                except Exception as e:
                    logger.error(f"Error collecting metadata (pre-execution): {str(e)}")

            # Call original function with exact parameters
            results = await original_map_node_over_list(
                prompt_id, unique_id, obj, input_data_all, func,
                allow_interrupt, execution_block_cb, pre_execute_cb, v3_data=v3_data
            )

            if func == obj.FUNCTION and hasattr(obj, '__class__'):
                try:
                    registry = MetadataRegistry()
                    if prompt_id is not None:
                        class_type = obj.__class__.__name__
                        node_id = unique_id
                        if node_id is not None:
                            return_types = getattr(obj, 'RETURN_TYPES', None)
                            registry.update_node_execution(node_id, class_type, results, return_types=return_types)
                except Exception as e:
                    logger.error(f"Error collecting metadata (post-execution): {str(e)}")

            return results

        # Also hook the execute function to track the current prompt_id
        original_execute = execution.execute

        async def async_execute_with_prompt_tracking(*args, **kwargs):
            if len(args) >= 7:  # Check if we have enough arguments
                server, prompt, caches, node_id, extra_data, executed, prompt_id = args[:7]
                registry = MetadataRegistry()

                # Start collection if this is a new prompt
                if not registry.current_prompt_id or registry.current_prompt_id != prompt_id:
                    registry.start_collection(prompt_id)

                # Store the dynprompt reference for node lookups
                if hasattr(prompt, 'original_prompt'):
                    registry.set_current_prompt(prompt)

                # Store extra_data for accessing full workflow node properties
                registry.set_extra_data(extra_data)

            # Execute the original function
            return await original_execute(*args, **kwargs)

        # Replace the functions with async versions
        setattr(execution, map_node_func_name, async_map_node_over_list_with_metadata)
        execution.execute = async_execute_with_prompt_tracking
