"""Utility functions for broadcasting to clients on main server loop in flwr next."""

from flwr.common import (
    Message,
)

from mt_dao.conf.base_schema import CommStack
from mt_dao.constants import (
    GATHER_INS,
    GATHER_MSG_TYPE,
    mt_dao_TTL,
)
from mt_dao.masks_utils import (
    mask_to_batches,
)
from mt_dao.messages.gather import (
    build_gather_base_recorddict,
)
from mt_dao.server.utils import ServerState


def finalize_gather_message(
    node_id: int,
    server_state: ServerState,
    comm_stack: CommStack,
    aggregation_mask: tuple[tuple[bool, ...], list[str], list[str]],
) -> list[Message]:
    """Construct gather messages for a client node, supporting batching if configured.

    This function creates one or more gather messages to request model parameters from a
    specific client node during the parameter collection phase of federated learning.
    If the communication stack is configured to use batching, the aggregation mask is
    split into batches and a message is created for each batch. Otherwise, a single
    message is created with the full aggregation mask. Each message is addressed to the
    specified node and includes the current round information.

    Parameters
    ----------
    node_id : int
        The identifier of the client node to which the gather message(s) will be sent.
    server_state : ServerState
        The current state of the server, including round number and layer structure.
    comm_stack : CommStack
        The communication stack configuration, specifying batching and other options.
    aggregation_mask : tuple[tuple[bool, ...], list[str], list[str]]
        The mask specifying which model parameters to request from the client, along
        with corresponding layer names and types.

    Returns
    -------
    list[Message]
        A list of gather messages to be sent to the specified client node. If batching
        is enabled, one message per batch is returned; otherwise, a single message is
        returned.

    """
    # Initialize the variables to return
    list_of_messages: list[Message] = []
    # If the communication stack is configured to use batches, we need to create
    # a batch of messages to send to each node
    if comm_stack.n_batches:
        # Create chunks fo the transmission mask based on the number of chunks
        batch_of_masks = mask_to_batches(
            full_mask=aggregation_mask,
            layer_names_and_types=server_state.layer_names_and_types,
            n_batches=comm_stack.n_batches,
        )
        # Iterate over the batches of masks
        for i, mask in enumerate(batch_of_masks):
            # Create the shared message of i-th batch for each node
            message = Message(
                content=build_gather_base_recorddict(
                    gather_msg=GATHER_INS,
                    aggregation_mask=mask,
                    batch_id=i,
                ),
                message_type=GATHER_MSG_TYPE,
                dst_node_id=0,
                group_id=str(server_state.current_round),
                ttl=mt_dao_TTL,
            )
            # Assign the template message to each node
            node_message = Message(
                content=message.content,
                message_type=GATHER_MSG_TYPE,
                dst_node_id=node_id,
                group_id=str(server_state.current_round),
                ttl=mt_dao_TTL,
            )
            list_of_messages.append(node_message)
        return list_of_messages

    # If the communication stack is not configured to use batches, we can
    # create a single message to send to each node
    shared_message = Message(
        content=build_gather_base_recorddict(
                    gather_msg=GATHER_INS,
            aggregation_mask=aggregation_mask,
        ),
        message_type=GATHER_MSG_TYPE,
        dst_node_id=0,
        group_id=str(server_state.current_round),
        ttl=mt_dao_TTL,
    )
    node_message = Message(
        content=shared_message.content,
        message_type=GATHER_MSG_TYPE,
        dst_node_id=node_id,
        group_id=str(server_state.current_round),
        ttl=mt_dao_TTL,
    )
    list_of_messages.append(node_message)
    return list_of_messages
