"""Utility functions for mt_dao server messages."""

from collections import deque
from logging import DEBUG, WARNING

from flwr.common import MessageType, log
from flwr.common.message import Message

from mt_dao.constants import (
    BROADCAST_MSG_TYPE,
    FREE_RESOURCES_MSG_TYPE,
    GATHER_MSG_TYPE,
)


class mt_daoMessageType(MessageType):
    """mt_dao-specific message types."""

    GATHER = f"{GATHER_MSG_TYPE}"
    BROADCAST = f"{BROADCAST_MSG_TYPE}"
    FREE_RESOURCES = f"{FREE_RESOURCES_MSG_TYPE}"


class FailureMessageError(Exception):
    """Exception raised when a failure message is encountered."""

    def __init__(self, message: Message) -> None:
        """Initialize the FailureMessageError with the message.

        Parameters
        ----------
        message : Message
            The message that indicates a failure.

        """
        super().__init__(f"Failure message received: {message}")
        self.message = message


def check_message_success(
    message: Message,
) -> bool:
    """Check if the message was successful.

    Parameters
    ----------
    message : Message
        The message to check.

    Returns
    -------
    bool
        True if the message was successful, False otherwise.

    """
    # TODO(anon): Improve
    log(
        DEBUG,
        "Got %s from %s",
        "result" if message.has_content() else "error",
        message.metadata.src_node_id,
    )
    # Check if the message was successful
    return message.has_content()


def interpret_unrecognized_message(
    message_id: str,
    message: Message,
    incoming_message_ids: deque[str],
    *,
    log_info: bool = False,
) -> None:
    """Interpret an unrecognized message.

    The weird 'system' message triggers a warning.

    Parameters
    ----------
    message_id : str
        The ID of the unrecognized message.
    message : Message
        The unrecognized message to interpret.
    incoming_message_ids : deque[str]
        The deque of incoming message IDs.
    log_info : bool
        Whether to log information about the unrecognized message.

    Raises
    ------
    ValueError
        If the message type is not recognized.

    """
    if "system" in message.metadata.message_type:
        # If there are no messages, we need to put the message_id back at the
        # end of the queue
        incoming_message_ids.appendleft(message_id)
    else:
        msg = f"Unknown message type: {message.metadata.message_type}"
        raise ValueError(msg)
    if log_info:
        log(
            WARNING,
            "Received an unrecognized message of type %s: %s",
            message.metadata.message_type,
            message,
        )
