"""Functions to handle the fit round in mt_dao server."""

import ast
import copy
from collections import deque
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from functools import partial
from queue import Queue
from typing import TYPE_CHECKING, Any, cast

from flwr.common import (
    Message,
    MessageType,
    Scalar,
)
from flwr.common.typing import ConfigRecordValues

from mt_dao.callbacks.fit_eval_callback import FitEvalCallback
from mt_dao.clients.configs import (
    get_mt_dao_evaluate_config_fn,
    get_mt_dao_fit_config_fn,
)
from mt_dao.conf.base_schema import BaseConfig, CommStack
from mt_dao.constants import CLIENT_STATE_ACCUMULATOR, EVAL_INS, FIT_INS
from mt_dao.masks_utils import ModelStateNames, generate_full_mask, generate_mask
from mt_dao.server.fit_eval.incoming_message import (
    create_gather_messages,
    handle_evaluate_message,
    queue_gather_results_incoming,
)
from mt_dao.server.fit_eval.outgoing_message import (
    handle_outgoing_message,
)
from mt_dao.server.fit_eval.utils import (
    mt_daoMessageType,
    interpret_unrecognized_message,
)
from mt_dao.server.load_balancer import ServerLoadBalancer
from mt_dao.server.server_state import ServerState
from mt_dao.server.utils import construct_message_for_client
from mt_dao.strategy_functional.aggregation import (
    submit_aggregation_tasks_to_executor,
    wait_for_all_futures,
    weighted_average,
)
from mt_dao.strategy_functional.dispatcher import ServerOptimizer
from mt_dao.utils import ClientState

if TYPE_CHECKING:
    import numpy as np


def handle_client_state_accumulator(
    server_state: ServerState,
    metrics: list[tuple[int, dict[str, Scalar]]],
) -> None:
    """Handle the client state accumulator from the metrics.

    This function collects the client states from the metrics and updates the
    server state with the accumulated client states. It also updates the
    server steps cumulative based on the maximum number of local steps done
    across the clients sampled in this round. Finally, it resets the steps_done
    for all clients to zero as it is just meant to be an ephemeral record of
    the local steps done during one federated round.

    Parameters
    ----------
    server_state : ServerState
        The current state of the server, which includes client states and
        server steps cumulative.
    metrics : list[tuple[int, dict[str, Scalar]]]
        A list of tuples containing the number of examples and the metrics
        for each client. The metrics should include the client state
        accumulator under the key `CLIENT_STATE_ACCUMULATOR`.

    """
    # Collect client states
    for _num_examples, inner_metrics in metrics:
        # NOTE: If the communication of the gather parameters is batched, only one
        # reply will contain the client state accumulator
        if CLIENT_STATE_ACCUMULATOR not in inner_metrics:
            continue
        # TODO(anon): Check why we need the double ast.literal_eval
        acc: dict[str | int, dict[str, Any]] = ast.literal_eval(
            ast.literal_eval(
                cast("str", inner_metrics.pop(CLIENT_STATE_ACCUMULATOR)),
            ),
        )
        # NOTE: When using partial participation we need to accumulate the keys of the
        # old state and the new state
        server_state.client_states |= {k: ClientState(**v) for k, v in acc.items()}

    # NOTE: Update the server steps cumulative by adding to the previous
    # value the maximum number of local steps done across the clients sample
    # in this round
    max_steps = max(
        *[
            _client_state.steps_done
            for _client_state in server_state.client_states.values()
        ],
        0,
    )
    server_state.server_steps_cumulative += max_steps
    # NOTE: Reset the steps_done for all clients to zero as it is just meant
    # to be an ephemeral record of the local steps done during one federated
    # round
    for c_state in server_state.client_states.values():
        c_state.steps_done = 0


def get_aggregation_mask(
    server_state: ServerState,
) -> tuple[tuple[bool, ...], list[str], list[str]]:
    """Get the aggregation mask for the current round.

    This function generates the aggregation mask based on the current
    server state. If the sampled clients have changed since the last round,
    it generates a full mask. Otherwise, it uses the aggregation mask
    generator to create a mask based on the current sampled clients and
    the current round.

    Parameters
    ----------
    server_state : ServerState
        The current state of the server, which includes the sampled clients,
        layer names and types, and the aggregation mask scheduler.

    Returns
    -------
    tuple[tuple[bool, ...], list[str], list[str]]
        A tuple describing the aggregation mask and containing:
        - A tuple of booleans indicating which parameters to aggregate.
        - A list of layer names.
        - A list of layer types.

    """
    # Construct aggregation mask generator
    aggregation_mask_generator = partial(
        generate_mask,
        server_state.layer_names_and_types,
        server_state.aggregation_mask_scheduler,
    )
    # Obtain the current aggregation 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
    return (
        generate_full_mask(server_state.layer_names_and_types)
        if set(server_state.sampled_clients)
        != set(
            server_state.previously_sampled_clients,
        )
        else aggregation_mask_generator(
            server_state.sampled_clients[-1],
            server_state.current_round,
        )
    )


def get_fit_config_fn(
    cfg: BaseConfig,
    server_state: ServerState,
) -> Callable[[int, str | int], dict[str, ConfigRecordValues]]:
    """Return a function to generate the fit configuration for clients.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration for the mt_dao server, which includes the model,
        tokenizer, and other settings.
    server_state : ServerState
        The current state of the server, which includes model states,
        layer names and types, and sampled clients.

    Returns
    -------
    Callable[[int, str | int], dict[str, ConfigRecordValues]]
        A function that generates the fit configuration for a client.

    """
    return get_mt_dao_fit_config_fn(
        cfg=cfg,
        aggregation_mask=partial(
            generate_mask,
            server_state.layer_names_and_types,
            server_state.aggregation_mask_scheduler,
        ),
        layer_names_and_types=server_state.layer_names_and_types,
    )


def create_eval_messages(
    cfg: BaseConfig,
    server_state: ServerState,
    server_load_balancer: ServerLoadBalancer,
) -> list[Message]:
    """Create evaluation messages for the clients.

    This function constructs the per-client evaluation instructions
    based on the configuration and the current server state. It generates
    messages that instruct each client to perform an evaluation, using the
    node ID assigned to each client by the load balancer.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration for the mt_dao server, which includes the model,
        tokenizer, and other settings.
    server_state : ServerState
        The current state of the server, which includes model states,
        layer names and types, and sampled clients.
    server_load_balancer : ServerLoadBalancer
        The load balancer that manages the distribution of clients across nodes.

    Returns
    -------
    list[Message]
        A list of messages containing the per-client evaluation instructions.

    """
    # Construct the per-client fit and evaluate instructions with a fake node ID (-1)
    # NOTE: This is done so we can assign the correct one later when we have mapped it
    # (almost) statically to the node IDs. This allows for efficiency gains when in need
    # to recover local checkpoints.
    per_client_evaluate_ins_messages: list[Message] = []
    # Add the per-client evaluate instructions if the evaluation period is set
    # and the current round is a multiple of the evaluation period
    if (
        cfg.fl.eval_period is not None
        and server_state.current_round % cfg.fl.eval_period == 0
    ):
        per_client_evaluate_ins_messages.extend(
            [
                construct_message_for_client(
                    ids=(
                        server_load_balancer.get_node_id_for_client(client_id),
                        client_id,
                    ),
                    gen_ins_function=get_mt_dao_evaluate_config_fn(cfg),
                    message_constants=(MessageType.EVALUATE, EVAL_INS),
                    server_state=server_state,
                )
                for client_id in copy.deepcopy(server_state.sampled_clients)
            ],
        )
    return per_client_evaluate_ins_messages


def eval_fit_round(
    cfg: BaseConfig,
    server_state: ServerState,
    server_load_balancer: ServerLoadBalancer,
    comm_stack: CommStack,
    server_optimizer: ServerOptimizer,
) -> None:
    """Run a evalfit round on the mt_dao server.

    This function orchestrates the fit round by sending fit instructions to
    the clients, gathering results, and applying the server optimizer.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration for the mt_dao server, which includes the model,
        tokenizer, and other settings.
    server_state : ServerState
        The current state of the server, which includes model states,
        layer names and types, and sampled clients.
    server_load_balancer : ServerLoadBalancer
        The load balancer that manages the distribution of clients across nodes.
    comm_stack : CommStack
        The communication stack used for sending and receiving messages.
    server_optimizer : ServerOptimizer
        A function that applies the server optimizer to the model states.

    """
    # Assign clients to nodes to allow for keeping consistent mapping throughout the
    # entire federated learning process
    server_load_balancer.assign_clients_to_nodes(
        copy.deepcopy(server_state.sampled_clients),
    )
    # Construct the per-client fit and evaluate instructions
    # NOTE: Adding this first so that they will be processed first
    per_client_ins_messages = create_eval_messages(
        cfg=cfg,
        server_state=server_state,
        server_load_balancer=server_load_balancer,
    )
    # NOTE: Adding this secondly so that they will be processed after eval
    per_client_ins_messages.extend(
        [
            construct_message_for_client(
                ids=(server_load_balancer.get_node_id_for_client(client_id), client_id),
                gen_ins_function=get_fit_config_fn(cfg, server_state),
                message_constants=(MessageType.TRAIN, FIT_INS),
                server_state=server_state,
            )
            for client_id in copy.deepcopy(server_state.sampled_clients)
        ],
    )
    # Initialize the queue for outgoing and incoming messages
    outgoing_messages: deque[Message] = deque()
    incoming_message_ids: deque[str] = deque()
    # Add the per-client fit evaluate instructions to the outgoing messages queue
    for message in per_client_ins_messages:
        outgoing_messages.appendleft(message)
    # Instantiate the metrics accumulator
    fit_metrics_accumulator: list[tuple[int, dict[str, Scalar]]] = []
    eval_metrics_accumulator: list[tuple[int, dict[str, Scalar]]] = []
    # Initialize the aggregation executor
    aggregation_executor = ThreadPoolExecutor()
    accumulator: dict[tuple[str, ModelStateNames], tuple[np.ndarray, int]] = {}
    futures: dict[tuple[str, ModelStateNames], Future[tuple[np.ndarray, int]]] = {}
    aggregation_queue: Queue[tuple[np.ndarray, int, str, ModelStateNames]] = Queue()
    # Create a fit eval callback
    fit_eval_callback = FitEvalCallback()
    # Loop until there are no more outgoing and incoming messages to process
    log_unrecognized_msg_info = True
    while outgoing_messages or incoming_message_ids:
        if outgoing_messages:
            handle_outgoing_message(
                outgoing_messages=outgoing_messages,
                server_load_balancer=server_load_balancer,
                incoming_message_ids=incoming_message_ids,
                fit_eval_callback=fit_eval_callback,
            )

        if incoming_message_ids:
            # Get the next incoming message
            message_id = incoming_message_ids.pop()
            # Check if the Grid has the message reply or not
            messages = server_load_balancer.grid.pull_messages([message_id])
            if not messages:
                # 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:
                # Get the first (and unique) message from the messages
                message = next(iter(messages))
                fit_eval_callback.record_recv_message(message)
                # Submit aggregation tasks to the executor if the message is a gather
                if message.metadata.message_type == mt_daoMessageType.EVALUATE:
                    # Handle the evaluate message
                    eval_metrics_accumulator.append(
                        handle_evaluate_message(message=message),
                    )
                    # Flag node as available (True)
                    server_load_balancer.nodes_states[message.metadata.src_node_id] = (
                        True
                    )
                elif message.metadata.message_type == mt_daoMessageType.GATHER:
                    # Put the results in the aggregation queue
                    queue_gather_results_incoming(
                        message=message,
                        aggregation_queue=aggregation_queue,
                        server_state=server_state,
                        comm_stack=comm_stack,
                    )
                    # Submit the aggregation tasks in the queue to the executor
                    fit_eval_callback.start_record_submit_aggregation()
                    accumulator = submit_aggregation_tasks_to_executor(
                        executor=aggregation_executor,
                        accumulator=accumulator,
                        aggregation_queue=aggregation_queue,
                        futures=futures,
                    )
                    fit_eval_callback.end_record_submit_aggregation()
                elif message.metadata.message_type == mt_daoMessageType.TRAIN:
                    fit_metrics_accumulator.append(
                        create_gather_messages(
                            message=message,
                            server_state=server_state,
                            comm_stack=comm_stack,
                            aggregation_mask=get_aggregation_mask(server_state),
                            outgoing_messages=outgoing_messages,
                        ),
                    )
                    # Flag node as available (True)
                    server_load_balancer.nodes_states[message.metadata.src_node_id] = (
                        True
                    )
                else:
                    interpret_unrecognized_message(
                        message_id=message_id,
                        message=message,
                        incoming_message_ids=incoming_message_ids,
                        log_info=log_unrecognized_msg_info,
                    )
                    log_unrecognized_msg_info = False
    # Collect the remaining results from futures
    fit_eval_callback.start_record_waiting_for_aggregation()
    wait_for_all_futures(
        accumulator=accumulator,
        futures=futures,
    )
    fit_eval_callback.end_record_waiting_for_aggregation()
    # Execute the outer optimizer
    fit_eval_callback.start_record_server_optimization()
    server_optimizer(
        server_state,
        {key: item[0] for key, item in accumulator.items()},
    )
    fit_eval_callback.end_record_server_optimization()
    # Handle the client state accumulator
    handle_client_state_accumulator(
        server_state=server_state,
        metrics=fit_metrics_accumulator,
    )
    # Aggregate the fit metrics
    server_state.history.add_metrics_distributed_fit(
        server_round=server_state.current_round,
        metrics=weighted_average(fit_metrics_accumulator),
    )
    # Aggregate the evaluation metrics
    server_state.history.add_metrics_distributed(
        server_round=server_state.current_round,
        metrics=weighted_average(eval_metrics_accumulator),
    )
    # Add server metrics to the history
    server_state.history.add_metrics_centralized(
        server_round=server_state.current_round,
        metrics=fit_eval_callback.compose_metrics(),
    )
