"""Multi-channel management for NMDB integration."""

import time
from typing import Dict, Optional
from pathlib import Path

try:
    from nmdb import ChannelClient, CICData
except ImportError:
    # Fallback if nmdb-py not installed
    class ChannelClient:
        def __init__(self, socket_path):
            self.socket_path = socket_path
        def connect(self):
            return False
        def disconnect(self):
            pass
        def send_message(self, data):
            return False
        def receive_message(self, timeout=None):
            return None

    class CICData:
        def __init__(self, emb=None):
            self.emb = emb or []
        def to_bytes(self):
            return b""
        @classmethod
        def from_bytes(cls, data):
            return cls()

from saltts.nmdb_integration.data_embedding_pair import DataEmbeddingPair


class ChannelManager:
    """Manages multiple NMDB channels for SalTTS.

    Uses Unix Socket connections to communicate with NMDB channels.
    """

    def __init__(self, socket_dir: str = "/tmp/nmdb"):
        """Initialize channel manager.

        Args:
            socket_dir: Directory containing NMDB socket files
        """
        self.socket_dir = Path(socket_dir)
        self.channels = {
            "text1": {"socket": "text1.sock", "type": "text"},
            "text2": {"socket": "text2.sock", "type": "text"},
            "data3": {"socket": "data3.sock", "type": "binary"},
        }
        self.clients: Dict[str, ChannelClient] = {}

    def connect(self, channel_name: str) -> bool:
        """Connect to a specific channel.

        Args:
            channel_name: Channel name (text1, text2, data3)

        Returns:
            True if connected successfully
        """
        if channel_name not in self.channels:
            raise ValueError(f"Unknown channel: {channel_name}")

        socket_path = self.socket_dir / self.channels[channel_name]["socket"]
        client = ChannelClient(str(socket_path))

        if client.connect():
            self.clients[channel_name] = client
            return True
        return False

    def connect_all(self) -> Dict[str, bool]:
        """Connect to all channels.

        Returns:
            Dictionary mapping channel names to connection status
        """
        results = {}
        for channel_name in self.channels.keys():
            results[channel_name] = self.connect(channel_name)
        return results

    def disconnect(self, channel_name: str) -> None:
        """Disconnect from a specific channel.

        Args:
            channel_name: Channel name
        """
        if channel_name in self.clients:
            self.clients[channel_name].disconnect()
            del self.clients[channel_name]

    def disconnect_all(self) -> None:
        """Disconnect from all channels."""
        for channel_name in list(self.clients.keys()):
            self.disconnect(channel_name)

    def send_cic_data(self, channel_name: str, cic_data: CICData) -> bool:
        """Send CICData to channel.

        Args:
            channel_name: Channel name
            cic_data: CICData object to send

        Returns:
            True if sent successfully
        """
        if channel_name not in self.clients:
            raise RuntimeError(f"Not connected to channel: {channel_name}")

        data = cic_data.to_bytes()
        return self.clients[channel_name].send_message(data)

    def receive_cic_data(
        self,
        channel_name: str,
        timeout: float = 5.0
    ) -> Optional[DataEmbeddingPair]:
        """Receive CICData from channel.

        Args:
            channel_name: Channel name
            timeout: Timeout in seconds

        Returns:
            DataEmbeddingPair or None if timeout
        """
        if channel_name not in self.clients:
            raise RuntimeError(f"Not connected to channel: {channel_name}")

        try:
            data = self.clients[channel_name].receive_message(timeout=timeout)
            if data is None:
                return None

            cic_data = CICData.from_bytes(data)
            return DataEmbeddingPair(
                cic_data=cic_data,
                channel_name=channel_name,
                timestamp=time.time(),
            )
        except TimeoutError:
            return None

    def receive_all_channels(
        self,
        timeout: float = 5.0
    ) -> Dict[str, Optional[DataEmbeddingPair]]:
        """Receive data from all connected channels.

        Args:
            timeout: Timeout in seconds

        Returns:
            Dictionary mapping channel names to DataEmbeddingPairs
        """
        results = {}
        for channel_name in self.clients.keys():
            results[channel_name] = self.receive_cic_data(channel_name, timeout)
        return results

    def __enter__(self):
        """Context manager entry."""
        self.connect_all()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.disconnect_all()
