"""
Plugin system for extensibility and custom nodes.
"""

from typing import Any, Dict, Type, Optional
from ybagent.node import BaseNode, NodeType


class NodeRegistry:
    """
    Registry for node types.
    """
    
    _registry: Dict[str, Type[BaseNode]] = {}
    
    @classmethod
    def register(cls, node_type: str, node_class: Type[BaseNode]):
        """
        Register a node type.
        
        Args:
            node_type: Node type identifier
            node_class: Node class
        """
        if node_type in cls._registry:
            raise ValueError(f"Node type '{node_type}' already registered")
        
        cls._registry[node_type] = node_class
    
    @classmethod
    def unregister(cls, node_type: str):
        """
        Unregister a node type.
        
        Args:
            node_type: Node type identifier
        """
        if node_type in cls._registry:
            del cls._registry[node_type]
    
    @classmethod
    def get(cls, node_type: str) -> Optional[Type[BaseNode]]:
        """
        Get a node class by type.
        
        Args:
            node_type: Node type identifier
        
        Returns:
            Node class or None
        """
        return cls._registry.get(node_type)
    
    @classmethod
    def list_types(cls) -> list:
        """
        List all registered node types.
        
        Returns:
            List of node type identifiers
        """
        return list(cls._registry.keys())
    
    @classmethod
    def clear(cls):
        """Clear all registrations."""
        cls._registry.clear()


class PluginManager:
    """
    Manager for loading and managing plugins.
    """
    
    def __init__(self):
        self.plugins: Dict[str, Any] = {}
        self._auto_register_builtin_nodes()
    
    def _auto_register_builtin_nodes(self):
        """Auto-register built-in node types."""
        from ybagent.nodes.io_nodes import InputNode, OutputNode, VariableNode
        from ybagent.nodes.llm_nodes import OpenAINode, HuggingFaceNode, LLMPromptNode, OllamaNode
        from ybagent.nodes.control_nodes import ConditionalNode, SwitchNode, LoopNode, MergeNode
        from ybagent.nodes.action_nodes import APICallNode, TransformNode, ScriptNode
        from ybagent.nodes.data_nodes import WebScraperNode, FileReadNode, FileWriteNode, EmailNode, DatabaseNode
        from ybagent.nodes.ai_nodes import VectorSearchNode, SummarizationNode, TranslationNode, SentimentAnalysisNode, ClassificationNode, NERNode
        from ybagent.nodes.debug_nodes import LoggerNode, DebugNode, NotificationNode, FailNode
        from ybagent.nodes.agent_nodes import CodeAgentNode
        
        # Register built-in nodes
        NodeRegistry.register("input", InputNode)
        NodeRegistry.register("output", OutputNode)
        NodeRegistry.register("variable", VariableNode)
        NodeRegistry.register("openai", OpenAINode)
        NodeRegistry.register("huggingface", HuggingFaceNode)
        NodeRegistry.register("llm_prompt", LLMPromptNode)
        NodeRegistry.register("ollama", OllamaNode)
        NodeRegistry.register("conditional", ConditionalNode)
        NodeRegistry.register("switch", SwitchNode)
        NodeRegistry.register("loop", LoopNode)
        NodeRegistry.register("merge", MergeNode)
        NodeRegistry.register("api_call", APICallNode)
        NodeRegistry.register("transform", TransformNode)
        NodeRegistry.register("script", ScriptNode)
        
        # Register data nodes
        NodeRegistry.register("web_scraper", WebScraperNode)
        NodeRegistry.register("file_read", FileReadNode)
        NodeRegistry.register("file_write", FileWriteNode)
        NodeRegistry.register("email", EmailNode)
        NodeRegistry.register("database", DatabaseNode)
        
        # Register AI nodes
        NodeRegistry.register("vector_search", VectorSearchNode)
        NodeRegistry.register("summarization", SummarizationNode)
        NodeRegistry.register("translation", TranslationNode)
        NodeRegistry.register("sentiment", SentimentAnalysisNode)
        NodeRegistry.register("classification", ClassificationNode)
        NodeRegistry.register("ner", NERNode)
        
        # Register agent nodes
        NodeRegistry.register("code_agent", CodeAgentNode)
        
        # Register debug nodes
        NodeRegistry.register("logger", LoggerNode)
        NodeRegistry.register("debug", DebugNode)
        NodeRegistry.register("notification", NotificationNode)
        NodeRegistry.register("fail", FailNode)
    
    def load_plugin(self, plugin_name: str, plugin_module: Any):
        """
        Load a plugin.
        
        Args:
            plugin_name: Plugin identifier
            plugin_module: Plugin module or object
        """
        if plugin_name in self.plugins:
            raise ValueError(f"Plugin '{plugin_name}' already loaded")
        
        self.plugins[plugin_name] = plugin_module
        
        # If plugin has a register function, call it
        if hasattr(plugin_module, "register"):
            plugin_module.register(NodeRegistry)
    
    def unload_plugin(self, plugin_name: str):
        """
        Unload a plugin.
        
        Args:
            plugin_name: Plugin identifier
        """
        if plugin_name in self.plugins:
            plugin = self.plugins[plugin_name]
            
            # If plugin has an unregister function, call it
            if hasattr(plugin, "unregister"):
                plugin.unregister(NodeRegistry)
            
            del self.plugins[plugin_name]
    
    def get_plugin(self, plugin_name: str) -> Optional[Any]:
        """
        Get a loaded plugin.
        
        Args:
            plugin_name: Plugin identifier
        
        Returns:
            Plugin module or None
        """
        return self.plugins.get(plugin_name)
    
    def list_plugins(self) -> list:
        """
        List all loaded plugins.
        
        Returns:
            List of plugin names
        """
        return list(self.plugins.keys())


# Global plugin manager instance
_plugin_manager = PluginManager()


def get_plugin_manager() -> PluginManager:
    """
    Get the global plugin manager instance.
    
    Returns:
        PluginManager instance
    """
    return _plugin_manager
