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

import time
from logging import DEBUG

from flwr.common import (
    Message,
    log,
)

from mt_dao.conf.base_schema import CommStack
from mt_dao.constants import (
    BROADCAST_INS,
    BROADCAST_RES,
    BROADCAST_TTL,
    mt_dao_TTL,
)
from mt_dao.masks_utils import generate_full_mask, generate_mask, mask_to_batches
from mt_dao.messages.broadcast import build_broadcast_recorddict
from mt_dao.messages.generic import build_success_recorddict
from mt_dao.server.comms import (
    offload_recorddict_parameters_to_remote,
)
from mt_dao.server.fit_eval.utils import mt_daoMessageType
from mt_dao.server.load_balancer import ServerLoadBalancer
from mt_dao.server.s3_utils import delete_past_communication_states
from mt_dao.server.utils import ServerState
from mt_dao.shm.utils import compress_with_strict
from ray import ObjectRef


def compose_broadcast_messages(
    msg_str: str,
    server_load_balancer: ServerLoadBalancer,
    server_state: ServerState,
    comm_stack: CommStack,
    transmission_mask: tuple[tuple[bool, ...], list[str], list[str]],
) -> tuple[list[Message], list[ObjectRef | str] | None]:
    """Create messages for broadcasting parameters to all registered nodes.

    This function prepares messages for broadcasting model parameters to nodes in a
    federated learning system. It supports both batched and non-batched transmission
    modes based on the communication stack configuration. For batched transmission,
    it divides the parameters into multiple messages. It handles offloading parameters
    to remote storage when necessary and creates individualized messages for each node.

    Parameters
    ----------
    msg_str : str
        Identifier string used for record keys
    server_load_balancer : ServerLoadBalancer
        Load balancer managing node registration and messaging
    server_state : ServerState
        Server state containing model parameters and configuration
    comm_stack : CommStack
        Communication stack configuration controlling transmission behavior
    transmission_mask : tuple[tuple[bool, ...], list[str], list[str]]
        Tuple containing (boolean selectors, layer names, layer types) that
        specifies which parameters to transmit

    Returns
    -------
    tuple[list[Message], list[ObjectRef | str] | None]
        A tuple containing a list of messages to be sent and a list of Ray object
        references (or None if no Ray objects were created)

    """
    # 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:
        list_of_garbage_objects: list[ObjectRef | str] = []
        # Create chunks for the transmission mask based on the number of chunks
        batch_of_masks = mask_to_batches(
            full_mask=transmission_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_broadcast_recorddict(
                    transmission_mask=mask,
                    batch_id=i,
                ),
                message_type=mt_daoMessageType.BROADCAST,
                dst_node_id=0,
                group_id=str(server_state.current_round),
                ttl=mt_dao_TTL,
            )
            # Add the parameters to the message once per mask
            message.content, in_list_of_garbage_objects = (
                offload_recorddict_parameters_to_remote(
                    parameters=list(
                        compress_with_strict(
                            data=server_state.model_states,
                            selectors=mask[0],
                            strict=True,
                        ),
                    ),
                    remote_uploader_downloader=server_state.remote_up_down,
                    outgoing_recorddict=message.content,
                    comm_stack=comm_stack,
                    msg_str=msg_str,
                )
            )
            # Assign the template message to each node
            for node_id in server_load_balancer.get_registered_nodes():
                node_message = Message(
                    content=message.content,
                    message_type=mt_daoMessageType.BROADCAST,
                    dst_node_id=node_id,
                    group_id=str(server_state.current_round),
                    ttl=mt_dao_TTL,
                )
                list_of_messages.append(node_message)
            # Add the list of garbage objects to the list of garbage objects
            list_of_garbage_objects.extend(in_list_of_garbage_objects)
        return list_of_messages, list_of_garbage_objects

    # 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_broadcast_recorddict(
            transmission_mask=transmission_mask,
        ),
        message_type=mt_daoMessageType.BROADCAST,
        dst_node_id=0,
        group_id=str(server_state.current_round),
        ttl=mt_dao_TTL,
    )
    shared_message.content, in_list_of_garbage_objects = (
        offload_recorddict_parameters_to_remote(
            parameters=server_state.model_states,
            remote_uploader_downloader=server_state.remote_up_down,
            outgoing_recorddict=shared_message.content,
            comm_stack=comm_stack,
            msg_str=msg_str,
        )
    )
    for node_id in server_load_balancer.get_registered_nodes():
        node_message = Message(
            content=shared_message.content,
            message_type=mt_daoMessageType.BROADCAST,
            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, list(in_list_of_garbage_objects)


def broadcast_parameters_to_nodes(
    server_load_balancer: ServerLoadBalancer,
    server_state: ServerState,
    comm_stack: CommStack,
) -> None:
    """Broadcast model parameters to all registered nodes in the federated system.

    This function orchestrates the broadcasting of model parameters to all registered
    nodes in the federated learning system. It determines which parameters to send using
    transmission masks, compresses the model state data, and creates appropriate
    messages for each node. It handles both batched and non-batched transmission modes,
    sends the messages, waits for acknowledgment from all nodes, and manages cleanup of
    Ray object references.

    Parameters
    ----------
    server_load_balancer : ServerLoadBalancer
        Load balancer managing node registration and messaging
    server_state : ServerState
        Server state containing model parameters and configuration
    comm_stack : CommStack
        Communication stack configuration controlling transmission behavior

    Raises
    ------
    ValueError
        If the broadcast fails or if the message format is incorrect.
    TimeoutError
        If the expected number of replies is not received within the broadcast TTL.

    """
    # Construct transmission mask
    # NOTE: Taking one client as representative since we assume all clients have
    # the same mask. In the future we may have more complex masks and schedulers
    transmission_mask = (
        generate_full_mask(server_state.layer_names_and_types)
        if set(server_state.sampled_clients)
        != set(
            server_state.previously_sampled_clients,
        )
        else generate_mask(
            server_state.layer_names_and_types,
            server_state.transmission_scheduler,
            server_state.sampled_clients[-1],
            # NOTE: The transmission scheduler must look backward in time to the last
            # round's masks
            server_state.current_round - 1,
        )
    )
    # Message name
    msg_str = f"{BROADCAST_INS}"
    # Create a message for each node
    messages, list_of_garbage = compose_broadcast_messages(
        msg_str=msg_str,
        server_load_balancer=server_load_balancer,
        server_state=server_state,
        comm_stack=comm_stack,
        transmission_mask=transmission_mask,
    )
    # Push all messages to all nodes
    log(
        DEBUG,
        "Pushing %s %s messages to nodes %s",
        len(messages),
        mt_daoMessageType.BROADCAST,
        server_load_balancer.get_registered_nodes(),
    )
    message_ids: list[str] = []
    # NOTE: We must send over a message at a time to avoid the Driver to pack them
    # together thus creating problems with the ~500MB limit of gRPC
    for message in messages:
        message_ids.extend(server_load_balancer.grid.push_messages([message]))
    log(
        DEBUG,
        "Pushed %s %s messages to nodes %s",
        len(messages),
        mt_daoMessageType.BROADCAST,
        server_load_balancer.get_registered_nodes(),
    )
    # Set all nodes as unavailable until they have received the correct model state
    server_load_balancer.set_all_nodes_states(state=False)
    # Wait for results, ignore empty message_ids
    message_ids = [message_id for message_id in message_ids if message_id]
    all_replies: list[Message] = []

    start_time = time.time()
    while time.time() - start_time < BROADCAST_TTL:
        replies = server_load_balancer.grid.pull_messages(message_ids=message_ids)
        all_replies += replies
        if len(all_replies) == len(message_ids):
            break

    if len(all_replies) < len(message_ids):
        log(
            DEBUG,
            "Received %s replies out of %s expected messages",
            len(all_replies),
            len(message_ids),
        )
        msg = (
            f"Received {len(all_replies)} replies out of {len(message_ids)}."
            "Broadcast failed."
        )
        raise TimeoutError(
            msg,
        )

    # Filter correct results
    messages_with_content = [msg for msg in all_replies if msg.has_content()]
    # Elaborate results
    for message_with_content in messages_with_content:
        if message_with_content.content != build_success_recorddict(BROADCAST_RES):
            log(
                DEBUG,
                "Received failure message: %s",
                message_with_content.content,
            )
            msg = f"Received failure message: {message_with_content.content}"
            raise ValueError(
                msg,
            )
    if comm_stack.s3:
        assert server_state.remote_up_down is not None, (
            "Remote uploader/downloader is None. "
            "Please check if the server state is initialized correctly."
        )
        # Clean up all the past arrays stored in the S3 as they are stale
        delete_past_communication_states(
            bucket_name=server_state.remote_up_down.remote_bucket_name,
            prefix=server_state.remote_up_down.backend_kwargs["prefix"],
        )
    server_load_balancer.set_all_nodes_states(state=True)
    log(DEBUG, f"Received {len(messages_with_content)} correct results")
    # Append list_of_garbage to the garbage collector queue
    if list_of_garbage:
        assert server_state.params_garbage_queue is not None, (
            "Ray garbage queue is None. "
            "Please check if the server state is initialized correctly."
        )
        # Make it a set to avoid duplicates
        list_of_garbage = list(set(list_of_garbage))
        for obj_ref in list_of_garbage:
            server_state.params_garbage_queue.put(obj_ref)
