"""Async streaming receiver for NMDB channels."""

import asyncio
from typing import Callable, Dict

from saltts.nmdb_integration.channel_manager import ChannelManager
from saltts.nmdb_integration.data_embedding_pair import DataEmbeddingPair


class StreamingReceiver:
    """Asynchronous streaming receiver for NMDB channels."""

    def __init__(self, channel_manager: ChannelManager):
        """Initialize streaming receiver.

        Args:
            channel_manager: Channel manager instance
        """
        self.channel_manager = channel_manager
        self.running = False
        self.callbacks: Dict[str, Callable[[DataEmbeddingPair], None]] = {}

    def register_callback(
        self,
        channel_name: str,
        callback: Callable[[DataEmbeddingPair], None],
    ) -> None:
        """Register callback for channel.

        Args:
            channel_name: Channel name
            callback: Callback function
        """
        self.callbacks[channel_name] = callback

    async def start_streaming(self, poll_interval_ms: int = 50) -> None:
        """Start streaming from all connected channels.

        Args:
            poll_interval_ms: Polling interval in milliseconds
        """
        self.running = True
        timeout_seconds = poll_interval_ms / 1000.0

        while self.running:
            # Poll all channels
            pairs = self.channel_manager.receive_all_channels(timeout=timeout_seconds)

            # Process received data
            for channel_name, pair in pairs.items():
                if pair and channel_name in self.callbacks:
                    # Run callback
                    self.callbacks[channel_name](pair)

            # Small delay to prevent busy waiting
            await asyncio.sleep(poll_interval_ms / 1000.0)

    def stop_streaming(self) -> None:
        """Stop streaming."""
        self.running = False
