plugin_system.py

"""
Plugin System for Anonka Bot
============================
This module provides a simple plugin architecture for extending bot functionality.

How to create a plugin:
1. Create a folder in the 'plugins/' directory (e.g., 'plugins/ban_plugin/')
2. Create a 'plugin.py' file inside that folder
3. Define a class that inherits from Plugin
4. Implement the setup() method to register handlers
5. Optionally implement cleanup() method for cleanup logic

Example plugin structure:
plugins/
├── ban_plugin/
│   ├── plugin.py
│   └── config.json (optional)
"""

import importlib.util
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, List, Optional

from aiogram import Bot, Dispatcher


class Plugin(ABC):
    """Base class for all plugins.
    
    Every plugin must inherit from this class and implement the setup() method.
    """
    
    def __init__(self, bot: Bot, dp: Dispatcher, plugin_dir: Path):
        """
        Initialize the plugin.
        
        Args:
            bot: The bot instance
            dp: The dispatcher instance
            plugin_dir: Path to the plugin directory
        """
        self.bot = bot
        self.dp = dp
        self.plugin_dir = plugin_dir
        self.name = self.__class__.__name__
        self.logger = logging.getLogger(f"plugin.{self.name}")
    
    @abstractmethod
    async def setup(self) -> bool:
        """
        Set up the plugin (register handlers, initialize resources, etc.).
        
        Returns:
            True if setup was successful, False otherwise
        """
        pass
    
    async def cleanup(self):
        """
        Clean up resources when the plugin is being disabled/unloaded.
        Override this method if your plugin needs cleanup logic.
        """
        pass
    
    def get_config_path(self) -> Path:
        """Get the path to the plugin's config.json file."""
        return self.plugin_dir / "config.json"


class PluginManager:
    """Manager for loading and managing plugins."""
    
    def __init__(self, bot: Bot, dp: Dispatcher, plugins_dir: str = "plugins"):
        """
        Initialize the plugin manager.
        
        Args:
            bot: The bot instance
            dp: The dispatcher instance
            plugins_dir: Path to the plugins directory
        """
        self.bot = bot
        self.dp = dp
        self.plugins_dir = Path(plugins_dir)
        self.plugins: Dict[str, Plugin] = {}
        self.logger = logging.getLogger("PluginManager")
        
        # Create plugins directory if it doesn't exist
        self.plugins_dir.mkdir(exist_ok=True)
        
        # Create __init__.py in plugins directory to make it a package
        init_file = self.plugins_dir / "__init__.py"
        if not init_file.exists():
            init_file.touch()
    
    def discover_plugins(self) -> List[Path]:
        """
        Discover all plugins in the plugins directory.
        
        Returns:
            List of paths to plugin directories
        """
        plugin_dirs = []
        
        if not self.plugins_dir.exists():
            self.logger.warning(f"Plugins directory '{self.plugins_dir}' does not exist")
            return plugin_dirs
        
        for item in self.plugins_dir.iterdir():
            if item.is_dir() and not item.name.startswith("_"):
                plugin_file = item / "plugin.py"
                if plugin_file.exists():
                    plugin_dirs.append(item)
                else:
                    self.logger.warning(f"Plugin directory '{item.name}' missing plugin.py file")
        
        return plugin_dirs
    
    def load_plugin(self, plugin_dir: Path) -> Optional[Plugin]:
        """
        Load a single plugin from a directory.
        
        Args:
            plugin_dir: Path to the plugin directory
            
        Returns:
            Plugin instance if successful, None otherwise
        """
        plugin_name = plugin_dir.name
        plugin_file = plugin_dir / "plugin.py"
        
        try:
            # Load the module
            spec = importlib.util.spec_from_file_location(
                f"plugins.{plugin_name}.plugin",
                plugin_file
            )
            if spec is None or spec.loader is None:
                self.logger.error(f"Failed to load plugin spec: {plugin_name}")
                return None
            
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            
            # Find the Plugin class in the module
            plugin_class = None
            for attr_name in dir(module):
                attr = getattr(module, attr_name)
                if (isinstance(attr, type) and 
                    issubclass(attr, Plugin) and 
                    attr is not Plugin):
                    plugin_class = attr
                    break
            
            if plugin_class is None:
                self.logger.error(f"No Plugin class found in {plugin_name}/plugin.py")
                return None
            
            # Instantiate the plugin
            plugin_instance = plugin_class(self.bot, self.dp, plugin_dir)
            return plugin_instance
            
        except Exception as e:
            self.logger.exception(f"Error loading plugin '{plugin_name}': {e}")
            return None
    
    async def load_all_plugins(self):
        """Discover and load all plugins."""
        self.logger.info("Discovering plugins...")
        plugin_dirs = self.discover_plugins()
        
        if not plugin_dirs:
            self.logger.info("No plugins found")
            return
        
        self.logger.info(f"Found {len(plugin_dirs)} plugin(s)")
        
        for plugin_dir in plugin_dirs:
            plugin_name = plugin_dir.name
            self.logger.info(f"Loading plugin: {plugin_name}")
            
            plugin = self.load_plugin(plugin_dir)
            if plugin is None:
                self.logger.error(f"Failed to load plugin: {plugin_name}")
                continue
            
            # Setup the plugin
            try:
                success = await plugin.setup()
                if success:
                    self.plugins[plugin_name] = plugin
                    self.logger.info(f"✓ Plugin '{plugin_name}' loaded successfully")
                else:
                    self.logger.error(f"✗ Plugin '{plugin_name}' setup failed")
            except Exception as e:
                self.logger.exception(f"✗ Error setting up plugin '{plugin_name}': {e}")
    
    async def unload_all_plugins(self):
        """Unload all plugins and clean up resources."""
        self.logger.info("Unloading plugins...")
        
        for plugin_name, plugin in self.plugins.items():
            try:
                await plugin.cleanup()
                self.logger.info(f"✓ Plugin '{plugin_name}' unloaded")
            except Exception as e:
                self.logger.exception(f"Error cleaning up plugin '{plugin_name}': {e}")
        
        self.plugins.clear()
    
    def get_plugin(self, name: str) -> Optional[Plugin]:
        """Get a loaded plugin by name."""
        return self.plugins.get(name)
    
    def list_plugins(self) -> List[str]:
        """Get a list of all loaded plugin names."""
        return list(self.plugins.keys())