"""Wrapper code factory"""

from typing import List
from diona.project.sandbox.wrappers.base import ModuleWrapper, BuiltinWrapper
from diona.project.sandbox.wrappers.registry import get_registry


class WrapperCodeFactory:
    """Factory for generating wrapper code"""
    
    def __init__(self):
        self._module_wrappers: List[ModuleWrapper] = []
        self._builtin_wrappers: List[BuiltinWrapper] = []
    
    def add_module_wrapper(self, wrapper: ModuleWrapper) -> 'WrapperCodeFactory':
        self._module_wrappers.append(wrapper)
        return self
    
    def add_builtin_wrapper(self, wrapper: BuiltinWrapper) -> 'WrapperCodeFactory':
        self._builtin_wrappers.append(wrapper)
        return self
    
    def generate_code(self, env_vars: dict = None) -> str:
        lines = []
        
        lines.append("import sys")
        lines.append("import io")
        lines.append("import builtins")
        lines.append("")
        
        if env_vars:
            for key, value in env_vars.items():
                lines.append(f"os.environ['{key}'] = '{value}'")
            lines.append("")
        
        for wrapper in self._module_wrappers:
            import_code = wrapper.get_import_code()
            if import_code:
                lines.append(import_code)
        lines.append("")
        
        for wrapper in self._module_wrappers:
            wrapper_code = wrapper.get_wrapper_code()
            if wrapper_code:
                lines.append(wrapper_code)
        
        for wrapper in self._builtin_wrappers:
            wrapper_code = wrapper.get_wrapper_code()
            if wrapper_code:
                lines.append(wrapper_code)
        
        return '\n'.join(lines)
    
    @classmethod
    def from_registry(cls) -> 'WrapperCodeFactory':
        factory = cls()
        registry = get_registry()
        
        for wrapper in registry.get_all_module_wrappers():
            factory.add_module_wrapper(wrapper)
        
        for wrapper in registry.get_all_builtin_wrappers():
            factory.add_builtin_wrapper(wrapper)
        
        return factory


def generate_wrapper_code(env_vars: dict = None) -> str:
    """Generate wrapper code from registry"""
    return WrapperCodeFactory.from_registry().generate_code(env_vars)
