import uuid
import struct
from typing import Dict, Optional, List
from mlagents_envs.side_channel import SideChannel, IncomingMessage
from mlagents_envs.exception import UnityEnvironmentException
from mlagents_envs.logging_util import get_logger


class SideChannelManager:
    def __init__(self, side_channels=Optional[List[SideChannel]]):
        self._side_channels_dict = self._get_side_channels_dict(side_channels)

    def process_side_channel_message(self, data: bytes) -> None:
        """
        Separates the data received from Python into individual messages for each
        registered side channel and calls on_message_received on them.
        :param data: The packed message sent by Unity
        """
        offset = 0
        while offset < len(data):
            try:
                channel_id = uuid.UUID(bytes_le=bytes(data[offset : offset + 16]))
                offset += 16
                message_len, = struct.unpack_from("<i", data, offset)
                offset = offset + 4
                message_data = data[offset : offset + message_len]
                offset = offset + message_len
            except (struct.error, ValueError, IndexError):
                raise UnityEnvironmentException(
                    "There was a problem reading a message in a SideChannel. "
                    "Please make sure the version of MLAgents in Unity is "
                    "compatible with the Python version."
                )
            if len(message_data) != message_len:
                raise UnityEnvironmentException(
                    "The message received by the side channel {} was "
                    "unexpectedly short. Make sure your Unity Environment "
                    "sending side channel data properly.".format(channel_id)
                )
            if channel_id in self._side_channels_dict:
                incoming_message = IncomingMessage(message_data)
                self._side_channels_dict[channel_id].on_message_received(
                    incoming_message
                )
            else:
                get_logger(__name__).warning(
                    f"Unknown side channel data received. Channel type: {channel_id}."
                )

    def generate_side_channel_messages(self) -> bytearray:
        """
        Gathers the messages that the registered side channels will send to Unity
        and combines them into a single message ready to be sent.
        """
        result = bytearray()
        for channel_id, channel in self._side_channels_dict.items():
            for message in channel.message_queue:
                result += channel_id.bytes_le
                result += struct.pack("<i", len(message))
                result += message
            channel.message_queue = []
        return result

    @staticmethod
    def _get_side_channels_dict(
        side_channels: Optional[List[SideChannel]]
    ) -> Dict[uuid.UUID, SideChannel]:
        """
        Converts a list of side channels into a dictionary of channel_id to SideChannel
        :param side_channels: The list of side channels.
        """
        side_channels_dict: Dict[uuid.UUID, SideChannel] = {}
        if side_channels is not None:
            for _sc in side_channels:
                if _sc.channel_id in side_channels_dict:
                    raise UnityEnvironmentException(
                        f"There cannot be two side channels with "
                        f"the same channel id {_sc.channel_id}."
                    )
                side_channels_dict[_sc.channel_id] = _sc
        return side_channels_dict