"""Functions handling the outgoing messages in the mt_dao server."""

from collections import deque
from logging import DEBUG

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

from mt_dao.callbacks.fit_eval_callback import FitEvalCallback
from mt_dao.constants import EVAL_INS, FIT_INS
from mt_dao.messages.generic import get_client_id_from_train_eval_ins_recorddict
from mt_dao.server.fit_eval.utils import mt_daoMessageType
from mt_dao.server.load_balancer import ServerLoadBalancer


def handle_outgoing_message(
    outgoing_messages: deque,
    server_load_balancer: ServerLoadBalancer,
    incoming_message_ids: deque,
    fit_eval_callback: FitEvalCallback,
) -> None:
    """Handle the next outgoing message in the queue.

    This function retrieves the next message from the outgoing messages queue,
    dispatches it to the appropriate client using the server load balancer,
    and adds the returned message back to the outgoing messages queue if it is
    a valid message. If the returned message is not a valid message, it is
    assumed to be an incoming message ID and is added to the incoming message
    IDs queue for further processing.

    Parameters
    ----------
    outgoing_messages : deque
        A deque containing the outgoing messages to be processed.
    server_load_balancer : ServerLoadBalancer
        The server load balancer used to route the messages to the appropriate clients.
    incoming_message_ids : deque
        A deque to store incoming message IDs that need to be processed later.
    fit_eval_callback : FitEvalCallback
        The callback used to record metrics related to message sending and receiving.

    """
    # Get the next message to send
    message = outgoing_messages.pop()
    # Dispatch the message to its correct pipeline
    returned_message = route_dispatch_outgoing(message, server_load_balancer)
    if isinstance(returned_message, Message):
        # If we have returned a message, it means that the message was not sent since
        # the node is not available and must be put back at the end of the queue, so it
        # won't slow down the processing of other messages
        outgoing_messages.appendleft(returned_message)
    else:
        # If we have returned an incoming message ID, it means that the message was sent
        # and we can start asking for a reply to the grid. We also record the message
        # sending time for the original message.
        fit_eval_callback.record_send_message(message)
        incoming_message_ids.appendleft(returned_message)


def route_dispatch_outgoing(
    message: Message,
    server_load_balancer: ServerLoadBalancer,
) -> Message | str:
    """Route the outgoing message based on its type.

    Parameters
    ----------
    message : Message
        The outgoing message to be routed.
    server_load_balancer : ServerLoadBalancer
        The server load balancer used to route the message.

    Returns
    -------
    Message | str
        The message to be sent to the client or an error message if the message type is
        unknown.

    Raises
    ------
    ValueError
        If the message type is unknown.

    """
    # Get the message type
    message_type = message.metadata.message_type
    # Route the message based on its type
    if message_type in {mt_daoMessageType.TRAIN, mt_daoMessageType.EVALUATE}:
        return dispatch_train_and_eval_message_outgoing(
            message,
            server_load_balancer,
            message_type=message_type,
        )
    if message_type == mt_daoMessageType.GATHER:
        return dispatch_gather_message_outgoing(message, server_load_balancer)
    msg = f"Unknown message type: {message_type}"
    raise ValueError(msg)


def dispatch_train_and_eval_message_outgoing(
    message: Message,
    server_load_balancer: ServerLoadBalancer,
    message_type: str,
) -> Message | str:
    """Dispatch a train or eval message to the appropriate client.

    Parameters
    ----------
    message : Message
        The train or eval message to be dispatched.
    server_load_balancer : ServerLoadBalancer
        The server load balancer used to route the message.
    message_type : str
        The type of the message being dispatched.

    Returns
    -------
    Message | str
        The message to be sent to the client or an error message if no client is found.

    """
    # Get the client ID from the message
    client_id = get_client_id_from_train_eval_ins_recorddict(
        recorddict=message.content,
        msg_str=FIT_INS if message_type == mt_daoMessageType.TRAIN else EVAL_INS,
    )
    # Get the node ID for the client
    node_id = server_load_balancer.get_node_id_for_client(client_id)
    # Loop over the available nodes
    if node_id in server_load_balancer.get_available_nodes():
        # Flag node as busy
        server_load_balancer.nodes_states[node_id] = False
        # Return the message to be sent
        log(
            DEBUG,
            "ServerApp: Sending message of type %s to node %s for client %s",
            message_type,
            node_id,
            client_id,
        )
        return next(
            iter(
                server_load_balancer.grid.push_messages(
                    [message],
                ),
            ),
        )
    # If no node is found, return the message as is
    return message


def dispatch_gather_message_outgoing(
    message: Message,
    server_load_balancer: ServerLoadBalancer,
) -> Message | str:
    """Dispatch a gather message to the appropriate client.

    Parameters
    ----------
    message : Message
        The gather message to be dispatched.
    server_load_balancer : ServerLoadBalancer
        The server load balancer used to route the message.

    Returns
    -------
    Message | str
        The message to be sent to the client or an error message if no client is found.

    """
    # Send out the message immediately as the message is already in the correct format
    log(
        DEBUG,
        "ServerApp: Sending message of type %s to node %s.",
        mt_daoMessageType.GATHER,
        message.metadata.dst_node_id,
    )
    return next(
        iter(
            server_load_balancer.grid.push_messages(
                [message],
            ),
        ),
    )
