"""Wrapper registry and factory"""

from typing import Dict, List, Type, Optional
from diona.project.sandbox.wrappers.base import ModuleWrapper, BuiltinWrapper


class WrapperRegistry:
    """Registry for module and builtin wrappers"""
    
    def __init__(self):
        self._module_wrappers: Dict[str, Type[ModuleWrapper]] = {}
        self._builtin_wrappers: List[Type[BuiltinWrapper]] = []
    
    def register_module(self, wrapper_class: Type[ModuleWrapper]) -> None:
        """Register a module wrapper"""
        instance = wrapper_class()
        self._module_wrappers[instance.name] = wrapper_class
    
    def register_builtin(self, wrapper_class: Type[BuiltinWrapper]) -> None:
        """Register a builtin wrapper"""
        self._builtin_wrappers.append(wrapper_class)
    
    def get_module_wrapper(self, name: str) -> Optional[Type[ModuleWrapper]]:
        """Get a module wrapper by name"""
        return self._module_wrappers.get(name)
    
    def get_all_module_wrappers(self) -> List[ModuleWrapper]:
        """Get all registered module wrappers sorted by dependency"""
        wrappers = []
        added = set()
        
        def add_wrapper(name: str):
            if name in added:
                return
            wrapper_class = self._module_wrappers.get(name)
            if not wrapper_class:
                return
            
            instance = wrapper_class()
            for dep in instance.get_dependencies():
                add_wrapper(dep)
            
            wrappers.append(instance)
            added.add(name)
        
        for name in list(self._module_wrappers.keys()):
            add_wrapper(name)
        
        return wrappers
    
    def get_all_builtin_wrappers(self) -> List[BuiltinWrapper]:
        """Get all registered builtin wrappers"""
        return [wrapper_class() for wrapper_class in self._builtin_wrappers]
    
    def clear(self) -> None:
        """Clear all registered wrappers"""
        self._module_wrappers.clear()
        self._builtin_wrappers.clear()


_global_registry: Optional[WrapperRegistry] = None


def get_registry() -> WrapperRegistry:
    """Get the global wrapper registry"""
    global _global_registry
    if _global_registry is None:
        _global_registry = WrapperRegistry()
    return _global_registry


def reset_registry() -> None:
    """Reset the global registry"""
    global _global_registry
    _global_registry = None
