"""Functions handling the incoming messages to the mt_dao server."""

import ast
from collections import deque
from logging import DEBUG
from queue import Queue
from typing import cast

import numpy as np
from flwr.common import Scalar, log
from flwr.common.message import Message

from mt_dao.conf.base_schema import CommStack
from mt_dao.constants import (
    COMM_STACK,
    EVAL_RES,
    FIT_RES,
    GATHER_RES,
    MASK,
)
from mt_dao.masks_utils import ModelStateNames
from mt_dao.messages.gather import gather_res_recorddict_to_num_samples
from mt_dao.messages.generic import get_metrics_from_recorddict
from mt_dao.server.comms import load_recorddict_parameters_from_remote
from mt_dao.server.fit_eval.utils import (
    FailureMessageError,
    check_message_success,
)
from mt_dao.server.gather_utils import finalize_gather_message
from mt_dao.server.server_state import ServerState


def handle_evaluate_message(
    message: Message,
) -> tuple[int, dict[str, Scalar]]:
    """Handle the incoming evaluation message.

    Parameters
    ----------
    message : Message
        The incoming message containing evaluation results.

    Returns
    -------
    tuple[int, dict[str, Scalar]]
        The metrics extracted from the evaluation message.

    Raises
    ------
    FailureMessageError
        If the incoming message is not successful.

    """
    # Check message success
    if check_message_success(message):
        log(
            DEBUG,
            "ServerApp: Collecting message of type %s from node %s.",
            EVAL_RES,
            message.metadata.src_node_id,
        )
        # Collect the evaluation metrics
        return get_metrics_from_recorddict(
            recorddict=message.content,
            msg_str=EVAL_RES,
        )
    raise FailureMessageError(message=message)


def create_gather_messages(
    message: Message,
    server_state: ServerState,
    comm_stack: CommStack,
    aggregation_mask: tuple[tuple[bool, ...], list[str], list[str]],
    outgoing_messages: deque,
) -> tuple[int, dict[str, Scalar]]:
    """Create gather messages based on the incoming message.

    This function creates gather messages based on the incoming message and
    adds them to the outgoing messages deque.

    Parameters
    ----------
    message : Message
        The incoming message containing the gather results.
    server_state : ServerState
        The current state of the server.
    comm_stack : CommStack
        The communication stack used for messaging.
    aggregation_mask : tuple[tuple[bool, ...], list[str], list[str]]
        The aggregation mask to apply to the incoming message.
    outgoing_messages : deque
        The deque to which the outgoing messages will be added.

    Returns
    -------
    tuple[int, dict[str, Scalar]]
        The metrics extracted from the incoming message and the number of
        outgoing messages created.

    Raises
    ------
    FailureMessageError
        If the incoming message is not successful.

    """
    # Check message success
    if check_message_success(message):
        # Get node ID
        node_id = message.metadata.src_node_id
        log(
            DEBUG,
            "ServerApp: Collecting message of type %s from node %s.",
            FIT_RES,
            node_id,
        )
        # Send gather message
        new_messages = finalize_gather_message(
            node_id=node_id,
            server_state=server_state,
            comm_stack=comm_stack,
            aggregation_mask=aggregation_mask,
        )
        # NOTE: We want the gather message to be the first message in the queue
        # so we append it to the start of the deque
        for new_message in new_messages:
            outgoing_messages.append(new_message)
        return get_metrics_from_recorddict(
            recorddict=message.content,
            msg_str=FIT_RES,
        )
    raise FailureMessageError(message=message)


def queue_gather_results_incoming(
    message: Message,
    aggregation_queue: Queue[tuple[np.ndarray, int, str, ModelStateNames]],
    server_state: ServerState,
    comm_stack: CommStack,
) -> None:
    """Queue the gather results from the incoming message.

    This function extracts the gather results from the incoming message and
    adds them to the aggregation queue for further processing.

    Parameters
    ----------
    message : Message
        The incoming message containing the gather results.
    aggregation_queue : Queue[tuple[np.ndarray, int, str, ModelStateNames]]
        The queue to which the gather results will be added.
    server_state : ServerState
        The current state of the server.
    comm_stack : CommStack
        The communication stack used for messaging.

    Raises
    ------
    FailureMessageError
        If the incoming message is not successful.

    """
    # Check message success
    if check_message_success(message):
        log(
            DEBUG,
            "ServerApp: Collecting message of type %s from node %s.",
            GATHER_RES,
            message.metadata.src_node_id,
        )
        message, current_params = load_recorddict_parameters_from_remote(
            remote_uploader_downloader=server_state.remote_up_down,
            incoming_message=message,
            comm_stack=comm_stack,
            msg_str=GATHER_RES,
            params_garbage_queue=server_state.params_garbage_queue,
        )
        num_samples = gather_res_recorddict_to_num_samples(
            gather_res_recorddict=message.content,
        )
        # Add the mask to the metrics to allow the strategy to
        # reconstruct the parameters correctly and execute partial updates
        _, current_layers, current_types = cast(
            "tuple[tuple[bool, ...], list[str], list[str]]",
            ast.literal_eval(
                cast(
                    "str",
                    message.content.config_records[f"{GATHER_RES}.{COMM_STACK}"][MASK],
                ),
            ),
        )
        # Fill the queue with the current parameters
        for param, name, typ in zip(
            current_params,
            current_layers,
            current_types,
            strict=True,
        ):
            aggregation_queue.put((param, num_samples, name, ModelStateNames(typ)))
    else:
        raise FailureMessageError(message=message)
