Trixy3/trixy_core/ml/container_integration.py

"""
Application Container Integration for Trixy ML Systems.

This module provides comprehensive integration between the ML components and the
main Trixy application container, enabling dependency injection, service management,
lifecycle coordination, and seamless integration with the main application.

Key Features:
- Dependency injection for ML components
- Service registry and discovery
- Lifecycle management and coordination
- Configuration management integration
- Event bus integration
- Health checking and monitoring
- Resource management and cleanup
- Plugin architecture support

Usage:
    from trixy_core.ml import MLContainerManager, MLServiceRegistry
    
    # Create container manager
    container_manager = MLContainerManager(app_container)
    
    # Register ML services
    await container_manager.register_ml_services()
    
    # Start ML integration
    await container_manager.start_ml_services()
"""

import asyncio
import logging
import weakref
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Dict, List, Optional, Type, Callable, Union, TypeVar, Generic
import inspect
from contextlib import asynccontextmanager


class ServiceState(Enum):
    """Service lifecycle states."""
    UNINITIALIZED = "uninitialized"
    INITIALIZING = "initializing"
    INITIALIZED = "initialized"
    STARTING = "starting"
    RUNNING = "running"
    STOPPING = "stopping"
    STOPPED = "stopped"
    ERROR = "error"


class ServiceScope(Enum):
    """Service scope definitions."""
    SINGLETON = "singleton"
    TRANSIENT = "transient"
    SCOPED = "scoped"


class DependencyType(Enum):
    """Dependency injection types."""
    REQUIRED = "required"
    OPTIONAL = "optional"
    LAZY = "lazy"


T = TypeVar('T')


@dataclass
class ServiceDefinition:
    """Service definition for dependency injection."""
    service_id: str
    service_type: Type
    implementation: Optional[Type] = None
    scope: ServiceScope = ServiceScope.SINGLETON
    factory: Optional[Callable] = None
    dependencies: List[str] = field(default_factory=list)
    config_key: Optional[str] = None
    auto_start: bool = True
    health_check: Optional[Callable] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ServiceInstance:
    """Service instance metadata."""
    definition: ServiceDefinition
    instance: Any
    state: ServiceState = ServiceState.UNINITIALIZED
    created_at: datetime = field(default_factory=datetime.now)
    started_at: Optional[datetime] = None
    last_health_check: Optional[datetime] = None
    health_status: bool = True
    error: Optional[Exception] = None


class ServiceError(Exception):
    """Base exception for service-related errors."""
    pass


class DependencyError(ServiceError):
    """Exception for dependency resolution errors."""
    pass


class LifecycleError(ServiceError):
    """Exception for service lifecycle errors."""
    pass


class IService(ABC):
    """Interface for managed services."""
    
    @abstractmethod
    async def initialize(self, context: Dict[str, Any]) -> None:
        """Initialize the service."""
        pass
    
    @abstractmethod
    async def start(self) -> None:
        """Start the service."""
        pass
    
    @abstractmethod
    async def stop(self) -> None:
        """Stop the service."""
        pass
    
    @abstractmethod
    async def health_check(self) -> bool:
        """Check service health."""
        pass


class IMLService(IService):
    """Interface for ML-specific services."""
    
    @abstractmethod
    async def process_audio(self, audio_data: Any) -> Any:
        """Process audio data."""
        pass
    
    @abstractmethod
    def get_performance_metrics(self) -> Dict[str, Any]:
        """Get performance metrics."""
        pass


class ServiceRegistry:
    """Service registry for dependency injection and service discovery."""
    
    def __init__(self):
        self.services: Dict[str, ServiceDefinition] = {}
        self.instances: Dict[str, ServiceInstance] = {}
        self.factories: Dict[str, Callable] = {}
        self._locks: Dict[str, asyncio.Lock] = {}
        self.logger = logging.getLogger(__name__)
    
    def register_service(self, definition: ServiceDefinition) -> None:
        """Register a service definition."""
        if definition.service_id in self.services:
            raise ServiceError(f"Service {definition.service_id} already registered")
        
        self.services[definition.service_id] = definition
        self._locks[definition.service_id] = asyncio.Lock()
        
        self.logger.info(f"Registered service: {definition.service_id}")
    
    def register_factory(self, service_id: str, factory: Callable) -> None:
        """Register a factory function for a service."""
        self.factories[service_id] = factory
        self.logger.debug(f"Registered factory for service: {service_id}")
    
    async def get_service(self, service_id: str, context: Optional[Dict[str, Any]] = None) -> Any:
        """Get a service instance."""
        if service_id not in self.services:
            raise ServiceError(f"Service {service_id} not registered")
        
        definition = self.services[service_id]
        
        # Handle singleton scope
        if definition.scope == ServiceScope.SINGLETON:
            if service_id in self.instances:
                instance_data = self.instances[service_id]
                if instance_data.state in [ServiceState.RUNNING, ServiceState.INITIALIZED]:
                    return instance_data.instance
        
        # Create new instance
        async with self._locks[service_id]:
            return await self._create_instance(service_id, context)
    
    async def _create_instance(self, service_id: str, context: Optional[Dict[str, Any]] = None) -> Any:
        """Create a new service instance."""
        definition = self.services[service_id]
        context = context or {}
        
        try:
            # Resolve dependencies
            dependencies = {}
            for dep_id in definition.dependencies:
                dependencies[dep_id] = await self.get_service(dep_id, context)
            
            # Create instance
            if definition.factory:
                instance = await self._call_factory(definition.factory, dependencies, context)
            elif service_id in self.factories:
                instance = await self._call_factory(self.factories[service_id], dependencies, context)
            elif definition.implementation:
                instance = self._create_from_type(definition.implementation, dependencies, context)
            else:
                instance = self._create_from_type(definition.service_type, dependencies, context)
            
            # Store instance
            instance_data = ServiceInstance(
                definition=definition,
                instance=instance,
                state=ServiceState.INITIALIZED
            )
            
            if definition.scope == ServiceScope.SINGLETON:
                self.instances[service_id] = instance_data
            
            self.logger.info(f"Created instance for service: {service_id}")
            return instance
            
        except Exception as e:
            self.logger.error(f"Error creating instance for service {service_id}: {e}")
            raise ServiceError(f"Failed to create service {service_id}: {e}")
    
    async def _call_factory(self, factory: Callable, dependencies: Dict[str, Any], 
                          context: Dict[str, Any]) -> Any:
        """Call a factory function with dependency injection."""
        signature = inspect.signature(factory)
        kwargs = {}
        
        for param_name, param in signature.parameters.items():
            if param_name in dependencies:
                kwargs[param_name] = dependencies[param_name]
            elif param_name in context:
                kwargs[param_name] = context[param_name]
        
        if asyncio.iscoroutinefunction(factory):
            return await factory(**kwargs)
        else:
            return factory(**kwargs)
    
    def _create_from_type(self, service_type: Type, dependencies: Dict[str, Any], 
                         context: Dict[str, Any]) -> Any:
        """Create instance from type with dependency injection."""
        try:
            signature = inspect.signature(service_type.__init__)
            kwargs = {}
            
            for param_name, param in signature.parameters.items():
                if param_name == 'self':
                    continue
                if param_name in dependencies:
                    kwargs[param_name] = dependencies[param_name]
                elif param_name in context:
                    kwargs[param_name] = context[param_name]
            
            return service_type(**kwargs)
            
        except Exception as e:
            raise ServiceError(f"Failed to create instance of {service_type}: {e}")
    
    async def start_service(self, service_id: str) -> None:
        """Start a service."""
        if service_id not in self.instances:
            await self.get_service(service_id)
        
        instance_data = self.instances[service_id]
        if instance_data.state == ServiceState.RUNNING:
            return
        
        try:
            instance_data.state = ServiceState.STARTING
            
            if isinstance(instance_data.instance, IService):
                await instance_data.instance.start()
            
            instance_data.state = ServiceState.RUNNING
            instance_data.started_at = datetime.now()
            
            self.logger.info(f"Started service: {service_id}")
            
        except Exception as e:
            instance_data.state = ServiceState.ERROR
            instance_data.error = e
            self.logger.error(f"Error starting service {service_id}: {e}")
            raise LifecycleError(f"Failed to start service {service_id}: {e}")
    
    async def stop_service(self, service_id: str) -> None:
        """Stop a service."""
        if service_id not in self.instances:
            return
        
        instance_data = self.instances[service_id]
        if instance_data.state != ServiceState.RUNNING:
            return
        
        try:
            instance_data.state = ServiceState.STOPPING
            
            if isinstance(instance_data.instance, IService):
                await instance_data.instance.stop()
            
            instance_data.state = ServiceState.STOPPED
            
            self.logger.info(f"Stopped service: {service_id}")
            
        except Exception as e:
            instance_data.state = ServiceState.ERROR
            instance_data.error = e
            self.logger.error(f"Error stopping service {service_id}: {e}")
    
    async def health_check_service(self, service_id: str) -> bool:
        """Perform health check on a service."""
        if service_id not in self.instances:
            return False
        
        instance_data = self.instances[service_id]
        
        try:
            if isinstance(instance_data.instance, IService):
                health = await instance_data.instance.health_check()
            elif instance_data.definition.health_check:
                health = await instance_data.definition.health_check(instance_data.instance)
            else:
                health = instance_data.state == ServiceState.RUNNING
            
            instance_data.health_status = health
            instance_data.last_health_check = datetime.now()
            
            return health
            
        except Exception as e:
            self.logger.error(f"Health check failed for service {service_id}: {e}")
            instance_data.health_status = False
            return False
    
    def get_service_info(self, service_id: str) -> Optional[Dict[str, Any]]:
        """Get service information."""
        if service_id not in self.services:
            return None
        
        definition = self.services[service_id]
        instance_data = self.instances.get(service_id)
        
        info = {
            'service_id': service_id,
            'service_type': definition.service_type.__name__,
            'scope': definition.scope.value,
            'auto_start': definition.auto_start,
            'dependencies': definition.dependencies,
            'metadata': definition.metadata
        }
        
        if instance_data:
            info.update({
                'state': instance_data.state.value,
                'health_status': instance_data.health_status,
                'created_at': instance_data.created_at.isoformat(),
                'started_at': instance_data.started_at.isoformat() if instance_data.started_at else None,
                'last_health_check': instance_data.last_health_check.isoformat() if instance_data.last_health_check else None,
                'error': str(instance_data.error) if instance_data.error else None
            })
        
        return info
    
    def list_services(self) -> List[str]:
        """List all registered services."""
        return list(self.services.keys())


class MLServiceRegistry(ServiceRegistry):
    """Specialized service registry for ML components."""
    
    def __init__(self):
        super().__init__()
        self._register_ml_services()
    
    def _register_ml_services(self) -> None:
        """Register default ML services."""
        
        # ML Manager
        self.register_service(ServiceDefinition(
            service_id="ml_manager",
            service_type=Any,  # Will be replaced with actual MLManager
            scope=ServiceScope.SINGLETON,
            dependencies=["ml_config", "event_handler", "performance_monitor"],
            auto_start=True,
            metadata={"component": "core", "category": "ml"}
        ))
        
        # Audio Pipeline
        self.register_service(ServiceDefinition(
            service_id="audio_pipeline",
            service_type=Any,  # Will be replaced with actual AudioPipeline
            scope=ServiceScope.SINGLETON,
            dependencies=["audio_config", "ml_manager"],
            auto_start=True,
            metadata={"component": "audio", "category": "ml"}
        ))
        
        # Performance Monitor
        self.register_service(ServiceDefinition(
            service_id="performance_monitor",
            service_type=Any,  # Will be replaced with actual PerformanceMonitor
            scope=ServiceScope.SINGLETON,
            dependencies=["performance_config"],
            auto_start=True,
            metadata={"component": "monitoring", "category": "system"}
        ))
        
        # Event Integration
        self.register_service(ServiceDefinition(
            service_id="event_integration",
            service_type=Any,  # Will be replaced with actual MLEventManager
            scope=ServiceScope.SINGLETON,
            dependencies=["event_config", "ml_manager"],
            auto_start=True,
            metadata={"component": "events", "category": "integration"}
        ))
        
        # Error Handler
        self.register_service(ServiceDefinition(
            service_id="error_handler",
            service_type=Any,  # Will be replaced with actual MLErrorHandler
            scope=ServiceScope.SINGLETON,
            dependencies=["error_config"],
            auto_start=True,
            metadata={"component": "error_handling", "category": "system"}
        ))


class MLContainerManager:
    """Container manager for ML integration with main application."""
    
    def __init__(self, app_container: Optional[Any] = None):
        self.app_container = app_container
        self.ml_registry = MLServiceRegistry()
        self.logger = logging.getLogger(__name__)
        self._running = False
        self._startup_tasks: List[Callable] = []
        self._shutdown_tasks: List[Callable] = []
        
        # Integration state
        self._event_bus_integration = False
        self._config_integration = False
        self._health_check_integration = False
    
    async def initialize(self) -> None:
        """Initialize the ML container integration."""
        try:
            self.logger.info("Initializing ML container integration...")
            
            # Setup integration with main container
            await self._setup_container_integration()
            
            # Register ML service factories
            await self._register_ml_factories()
            
            # Setup cross-cutting concerns
            await self._setup_cross_cutting_concerns()
            
            self.logger.info("ML container integration initialized")
            
        except Exception as e:
            self.logger.error(f"Error initializing ML container: {e}")
            raise ServiceError(f"ML container initialization failed: {e}")
    
    async def _setup_container_integration(self) -> None:
        """Setup integration with main application container."""
        if not self.app_container:
            self.logger.warning("No application container provided - running in standalone mode")
            return
        
        try:
            # Integrate with configuration system
            if hasattr(self.app_container, 'get_config'):
                self._config_integration = True
                self.logger.debug("Configuration integration enabled")
            
            # Integrate with event bus
            if hasattr(self.app_container, 'event_bus'):
                self._event_bus_integration = True
                self.logger.debug("Event bus integration enabled")
            
            # Integrate with health checking
            if hasattr(self.app_container, 'health_checker'):
                self._health_check_integration = True
                self.logger.debug("Health check integration enabled")
            
        except Exception as e:
            self.logger.error(f"Error setting up container integration: {e}")
    
    async def _register_ml_factories(self) -> None:
        """Register factory functions for ML services."""
        
        # ML Manager factory
        async def create_ml_manager(ml_config: Any, event_handler: Any, 
                                   performance_monitor: Any) -> Any:
            from .ml_manager import MLManager
            return MLManager(ml_config, event_handler, None, self.logger)
        
        self.ml_registry.register_factory("ml_manager", create_ml_manager)
        
        # Audio Pipeline factory
        async def create_audio_pipeline(audio_config: Any, ml_manager: Any) -> Any:
            from .audio_pipeline import AudioPipeline
            return AudioPipeline(audio_config, ml_manager, self.logger)
        
        self.ml_registry.register_factory("audio_pipeline", create_audio_pipeline)
        
        # Performance Monitor factory
        async def create_performance_monitor(performance_config: Any) -> Any:
            from .performance_monitor import PerformanceMonitor
            return PerformanceMonitor(performance_config)
        
        self.ml_registry.register_factory("performance_monitor", create_performance_monitor)
        
        # Event Integration factory
        async def create_event_integration(event_config: Any, ml_manager: Any) -> Any:
            from .event_integration import MLEventManager
            return MLEventManager(event_config, ml_manager, self.logger)
        
        self.ml_registry.register_factory("event_integration", create_event_integration)
        
        # Error Handler factory
        async def create_error_handler(error_config: Any) -> Any:
            from .error_handling import MLErrorHandler
            return MLErrorHandler(error_config, self.logger)
        
        self.ml_registry.register_factory("error_handler", create_error_handler)
    
    async def _setup_cross_cutting_concerns(self) -> None:
        """Setup cross-cutting concerns like logging, monitoring, etc."""
        
        # Setup logging integration
        if self.app_container and hasattr(self.app_container, 'logger'):
            self.logger = self.app_container.logger.get_child("ml")
        
        # Setup monitoring integration
        if self._health_check_integration:
            self._startup_tasks.append(self._register_health_checks)
        
        # Setup event integration
        if self._event_bus_integration:
            self._startup_tasks.append(self._setup_event_integration)
    
    async def start_ml_services(self) -> None:
        """Start all ML services."""
        if self._running:
            return
        
        try:
            self.logger.info("Starting ML services...")
            
            # Execute startup tasks
            for task in self._startup_tasks:
                await task()
            
            # Start services in dependency order
            service_order = [
                "performance_monitor",
                "error_handler", 
                "ml_manager",
                "audio_pipeline",
                "event_integration"
            ]
            
            for service_id in service_order:
                await self.ml_registry.start_service(service_id)
            
            self._running = True
            self.logger.info("ML services started successfully")
            
        except Exception as e:
            self.logger.error(f"Error starting ML services: {e}")
            await self.stop_ml_services()
            raise LifecycleError(f"Failed to start ML services: {e}")
    
    async def stop_ml_services(self) -> None:
        """Stop all ML services."""
        if not self._running:
            return
        
        try:
            self.logger.info("Stopping ML services...")
            
            # Stop services in reverse dependency order
            service_order = [
                "event_integration",
                "audio_pipeline",
                "ml_manager",
                "error_handler",
                "performance_monitor"
            ]
            
            for service_id in service_order:
                await self.ml_registry.stop_service(service_id)
            
            # Execute shutdown tasks
            for task in self._shutdown_tasks:
                await task()
            
            self._running = False
            self.logger.info("ML services stopped successfully")
            
        except Exception as e:
            self.logger.error(f"Error stopping ML services: {e}")
    
    async def _register_health_checks(self) -> None:
        """Register health checks with main application."""
        if not self.app_container or not hasattr(self.app_container, 'health_checker'):
            return
        
        try:
            health_checker = self.app_container.health_checker
            
            for service_id in self.ml_registry.list_services():
                async def check_service_health(sid=service_id):
                    return await self.ml_registry.health_check_service(sid)
                
                health_checker.register_check(f"ml_{service_id}", check_service_health)
            
            self.logger.info("ML health checks registered")
            
        except Exception as e:
            self.logger.error(f"Error registering health checks: {e}")
    
    async def _setup_event_integration(self) -> None:
        """Setup event bus integration."""
        if not self.app_container or not hasattr(self.app_container, 'event_bus'):
            return
        
        try:
            event_bus = self.app_container.event_bus
            
            # Subscribe to relevant application events
            event_bus.subscribe("app.shutdown", self._handle_app_shutdown)
            event_bus.subscribe("app.config_changed", self._handle_config_change)
            
            self.logger.info("Event bus integration setup")
            
        except Exception as e:
            self.logger.error(f"Error setting up event integration: {e}")
    
    async def _handle_app_shutdown(self, event: Dict[str, Any]) -> None:
        """Handle application shutdown event."""
        self.logger.info("Received application shutdown event")
        await self.stop_ml_services()
    
    async def _handle_config_change(self, event: Dict[str, Any]) -> None:
        """Handle configuration change event."""
        self.logger.info("Received configuration change event")
        # Restart services that depend on configuration
        # Implementation would reload configurations and restart affected services
    
    def get_service(self, service_id: str) -> Any:
        """Get a service instance."""
        return asyncio.create_task(self.ml_registry.get_service(service_id))
    
    def get_ml_manager(self) -> Any:
        """Get the ML manager instance."""
        return self.get_service("ml_manager")
    
    def get_audio_pipeline(self) -> Any:
        """Get the audio pipeline instance."""
        return self.get_service("audio_pipeline")
    
    def get_performance_monitor(self) -> Any:
        """Get the performance monitor instance."""
        return self.get_service("performance_monitor")
    
    def get_service_status(self) -> Dict[str, Any]:
        """Get status of all ML services."""
        status = {
            'running': self._running,
            'services': {}
        }
        
        for service_id in self.ml_registry.list_services():
            service_info = self.ml_registry.get_service_info(service_id)
            if service_info:
                status['services'][service_id] = service_info
        
        return status
    
    @asynccontextmanager
    async def ml_context(self):
        """Context manager for ML services lifecycle."""
        try:
            await self.initialize()
            await self.start_ml_services()
            yield self
        finally:
            await self.stop_ml_services()


# Convenience functions for common integration patterns
async def create_ml_container(app_container: Optional[Any] = None, 
                            auto_start: bool = True) -> MLContainerManager:
    """
    Create and initialize ML container manager.
    
    Args:
        app_container: Main application container
        auto_start: Whether to auto-start services
        
    Returns:
        Initialized ML container manager
    """
    manager = MLContainerManager(app_container)
    await manager.initialize()
    
    if auto_start:
        await manager.start_ml_services()
    
    return manager


def register_ml_service(registry: ServiceRegistry, service_id: str, 
                       service_type: Type, **kwargs) -> None:
    """
    Register a custom ML service.
    
    Args:
        registry: Service registry
        service_id: Unique service identifier
        service_type: Service type/class
        **kwargs: Additional service definition parameters
    """
    definition = ServiceDefinition(
        service_id=service_id,
        service_type=service_type,
        **kwargs
    )
    registry.register_service(definition)


class MLServiceDecorator:
    """Decorator for ML service classes."""
    
    def __init__(self, service_id: str, scope: ServiceScope = ServiceScope.SINGLETON,
                 dependencies: Optional[List[str]] = None, auto_start: bool = True):
        self.service_id = service_id
        self.scope = scope
        self.dependencies = dependencies or []
        self.auto_start = auto_start
    
    def __call__(self, cls: Type) -> Type:
        """Apply decorator to class."""
        # Store service metadata on the class
        cls._ml_service_id = self.service_id
        cls._ml_service_scope = self.scope
        cls._ml_service_dependencies = self.dependencies
        cls._ml_service_auto_start = self.auto_start
        
        return cls


# Example usage decorator
@MLServiceDecorator("custom_ml_service", dependencies=["ml_manager"])
class CustomMLService(IMLService):
    """Example custom ML service."""
    
    def __init__(self, ml_manager: Any):
        self.ml_manager = ml_manager
        self.logger = logging.getLogger(__name__)
    
    async def initialize(self, context: Dict[str, Any]) -> None:
        """Initialize the service."""
        self.logger.info("Initializing custom ML service")
    
    async def start(self) -> None:
        """Start the service."""
        self.logger.info("Starting custom ML service")
    
    async def stop(self) -> None:
        """Stop the service."""
        self.logger.info("Stopping custom ML service")
    
    async def health_check(self) -> bool:
        """Check service health."""
        return True
    
    async def process_audio(self, audio_data: Any) -> Any:
        """Process audio data."""
        return await self.ml_manager.process_audio_chunk(audio_data)
    
    def get_performance_metrics(self) -> Dict[str, Any]:
        """Get performance metrics."""
        return {"processed_samples": 0, "latency_ms": 0.0}