selfbot-legacy/bot/cogs/cog_manager.py

import os
import sys
import importlib
import importlib.util
import inspect
from typing import Dict, Any, Optional

class BaseCog:
    """Base cog class that all cogs should inherit from"""
    def __init__(self, bot):
        self.bot = bot
        self.name = self.__class__.__name__
        
    def cleanup(self):
        """Called when cog is unloaded"""
        pass

class CogManager:
    def __init__(self, bot):
        self.bot = bot
        self.cogs: Dict[str, BaseCog] = {}
        self.cogs_dir = os.path.join("bot", "cogs")
        os.makedirs(self.cogs_dir, exist_ok=True)
        
    def load_cog(self, cog_name: str) -> bool:
        """
        Load a cog by name
        
        Args:
            cog_name: Name of the cog file (without .py)
            
        Returns:
            bool: True if loaded successfully, False otherwise
        """
        try:
            # If already loaded, unload first
            if cog_name in self.cogs:
                self.unload_cog(cog_name)
                
            # Get full path to cog file
            cog_path = os.path.join(self.cogs_dir, f"{cog_name}.py")
            if not os.path.exists(cog_path):
                print(f"Cog file not found: {cog_path}")
                return False
                
            # Create a unique module name to avoid conflicts
            module_name = f"bot.cogs.{cog_name}"
            
            # Load the module
            spec = importlib.util.spec_from_file_location(module_name, cog_path)
            if not spec:
                print(f"Failed to create spec for {cog_path}")
                return False
                
            module = importlib.util.module_from_spec(spec)
            sys.modules[module_name] = module
            spec.loader.exec_module(module)
            
            # Find and instantiate the cog class
            # A cog class is a class that inherits from BaseCog
            cog_class = None
            for name, obj in inspect.getmembers(module):
                if (inspect.isclass(obj) and 
                    issubclass(obj, BaseCog) and 
                    obj is not BaseCog):
                    cog_class = obj
                    break
                    
            if not cog_class:
                print(f"No cog class found in {cog_path}")
                return False
                
            # Instantiate the cog
            cog = cog_class(self.bot)
            self.cogs[cog_name] = cog
            
            # Register commands
            for name, method in inspect.getmembers(cog, inspect.ismethod):
                if name.startswith("cmd_"):
                    cmd_name = name[4:]  # Remove "cmd_" prefix
                    self.bot.loaded_commands[cmd_name] = method
                    
            print(f"Loaded cog: {cog_name}")
            return True
            
        except Exception as e:
            print(f"Error loading cog {cog_name}: {e}")
            import traceback
            traceback.print_exc()
            return False
            
    def unload_cog(self, cog_name: str) -> bool:
        """
        Unload a cog by name
        
        Args:
            cog_name: Name of the cog
            
        Returns:
            bool: True if unloaded successfully, False otherwise
        """
        try:
            if cog_name not in self.cogs:
                print(f"Cog not loaded: {cog_name}")
                return False
                
            cog = self.cogs[cog_name]
            
            # Call cleanup method
            cog.cleanup()
            
            # Remove commands from loaded_commands
            for name, method in inspect.getmembers(cog, inspect.ismethod):
                if name.startswith("cmd_"):
                    cmd_name = name[4:]
                    if cmd_name in self.bot.loaded_commands:
                        del self.bot.loaded_commands[cmd_name]
            
            # Remove cog
            del self.cogs[cog_name]
            
            # Remove module
            module_name = f"bot.cogs.{cog_name}"
            if module_name in sys.modules:
                del sys.modules[module_name]
                
            print(f"Unloaded cog: {cog_name}")
            return True
            
        except Exception as e:
            print(f"Error unloading cog {cog_name}: {e}")
            return False
            
    def reload_cog(self, cog_name: str) -> bool:
        """
        Reload a cog by name
        
        Args:
            cog_name: Name of the cog
            
        Returns:
            bool: True if reloaded successfully, False otherwise
        """
        return self.load_cog(cog_name)
        
    def load_all_cogs(self) -> int:
        """
        Load all cogs from the cogs directory
        
        Returns:
            int: Number of cogs loaded successfully
        """
        count = 0
        for filename in os.listdir(self.cogs_dir):
            if filename.endswith(".py") and not filename.startswith("_"):
                cog_name = filename[:-3]
                if self.load_cog(cog_name):
                    count += 1
        return count
        
    def unload_all_cogs(self) -> int:
        """
        Unload all cogs
        
        Returns:
            int: Number of cogs unloaded successfully
        """
        cog_names = list(self.cogs.keys())
        count = 0
        for cog_name in cog_names:
            if self.unload_cog(cog_name):
                count += 1
        return count