"""
Multi-Signal Handler Middleware

Handles multiple types of signals with different callbacks and optional chunk emission.
"""

from collections.abc import AsyncIterator, AsyncGenerator
from typing import Dict, Callable, Awaitable, Optional, Union
from vixio.core.middleware import SignalMiddleware, NextHandler
from vixio.core.chunk import Chunk, ChunkType


# Type aliases for signal handlers
# Can be either:
# 1. Async callback that returns None
# 2. Async generator that yields chunks
SignalHandler = Union[
    Callable[[Chunk], Awaitable[None]],  # Simple callback
    Callable[[Chunk], AsyncIterator[Chunk]]  # Generator that yields chunks
]


class MultiSignalHandlerMiddleware(SignalMiddleware):
    """
    Handles multiple signal types with different callbacks.
    
    Data chunks are passed through unchanged.
    
    More flexible than SignalHandlerMiddleware which only handles CONTROL_STATE_RESET.
    Supports both simple callbacks and generators that yield chunks.
    
    Useful for stations that need to:
    - React to multiple signal types (e.g., TTS)
    - Emit events during signal processing
    - Update state and generate output chunks
    """
    
    def __init__(
        self,
        signal_handlers: Optional[Dict[ChunkType, SignalHandler]] = None,
        passthrough_signals: bool = True,
        name: str = "MultiSignalHandler"
    ):
        """
        Initialize multi-signal handler middleware.
        
        Args:
            signal_handlers: Dict mapping signal types to handlers
                - Handler can be async callback: async def handler(chunk) -> None
                - Or async generator: async def handler(chunk) -> AsyncIterator[Chunk]
            passthrough_signals: Whether to passthrough signals after handling
            name: Middleware name
        """
        super().__init__(name)
        self.signal_handlers = signal_handlers or {}
        self.passthrough_signals = passthrough_signals
    
    async def process_signal(self, chunk: Chunk, next_handler: NextHandler) -> AsyncGenerator[Chunk, None]:
        """
        Process signal chunk with type-specific handlers.
        
        Args:
            chunk: Signal chunk (not data)
            next_handler: Next handler in chain
            
        Yields:
            Chunks generated by signal handlers, plus signal passthrough
        """
        # Check if we have a handler for this signal type
        if chunk.type in self.signal_handlers:
            handler = self.signal_handlers[chunk.type]
            if handler:
                try:
                    # Try to iterate as async generator first
                    result = handler(chunk)
                    
                    # Check if it's an async generator
                    if hasattr(result, '__aiter__'):
                        # It's a generator, yield all chunks
                        async for output_chunk in result:
                            yield output_chunk
                        self.logger.debug(f"Handled signal (generator): {chunk.type.name}")
                    else:
                        # It's a simple coroutine, await it
                        await result
                        self.logger.debug(f"Handled signal (callback): {chunk.type.name}")
            
                except Exception as e:
                    self.logger.error(f"Error in signal handler for {chunk.type.name}: {e}")
        
            # Passthrough signal if enabled
            if self.passthrough_signals:
                yield chunk

