"""Utility functions for running the main server loop in flwr next."""

import time
from collections.abc import Callable
from dataclasses import asdict
from logging import DEBUG

from flwr.common import (
    ConfigRecord,
    ConfigRecordValues,
    Message,
    RecordDict,
    log,
)
from flwr.server import Grid

from mt_dao.constants import (
    CLIENT_ID,
    CLIENT_STATE,
    mt_dao_TTL,
    SERVER_ROUND,
    SERVER_STEPS_CUMULATIVE,
)
from mt_dao.server.server_state import ServerState

DROPOUT_FUNCTIONS: list[str] = ["random"]


def wait_for_nodes_to_connect(grid: Grid, n_nodes: int, timeout: float = 3) -> None:
    """Wait for a specified number of client nodes to connect.

    This function continuously checks the number of client nodes that have connected by
    calling the `get_node_ids` method on the provided grid object. It logs the current
    number of connected client nodes at each check. The function will exit once the
    number of connected client nodes meets or exceeds the specified `n_nodes`. If the
    required number of nodes are not found, the function will wait for the specified
    `timeout` period before checking again.

    Parameters
    ----------
    grid : Grid
        The grid object used to interact with the client nodes.
    n_nodes : int
        The minimum number of client nodes that must be connected.
    timeout : float, optional
        The time in seconds to wait between checks. Default is 3 seconds.

    """
    # The Grid API might not immediately return enough client node IDs, so we
    # loop and wait until enough client nodes are available.
    while True:
        all_node_ids = grid.get_node_ids()
        connected_nodes = len(list(all_node_ids))
        log(DEBUG, f"Got {connected_nodes} client nodes: {all_node_ids}")
        if connected_nodes >= n_nodes:
            break
        time.sleep(timeout)


def construct_message_for_client(
    ids: tuple[int, int],
    gen_ins_function: Callable[[int, int | str], dict[str, ConfigRecordValues]],
    message_constants: tuple[str, str],
    server_state: ServerState,
) -> Message:
    """Construct a federated learning message for a specific client-node pair.

    This function creates a `Message` object to be sent to a client node during a
    federated learning round. It packages the current server round, sampled client IDs,
    client states, and server step information into a `RecordDict`, along with
    client-specific instructions generated by `gen_ins_function`. The resulting message
    is addressed to the specified node and client, and is ready for dispatch via the
    server's communication grid.

    Parameters
    ----------
    ids : tuple[int, int]
        A tuple containing the node ID and client ID for which the message is built.
    gen_ins_function : Callable[[int, int | str], dict[str, ConfigRecordValues]]
        A function that generates client-specific instruction configurations based on
        the current round and client ID.
    message_constants : tuple[str, str]
        A tuple specifying the message type (e.g., "fit", "evaluate") and a string
        identifier for the message content.
    server_state : ServerState
        The current state of the server, including round number, sampled clients, and
        client states.

    Returns
    -------
    Message
        The constructed message containing configuration and instructions for the
        client.

    Notes
    -----
    - The message includes both global server state and client-specific instructions.
    - The message is addressed to the specified node and grouped by the current round.

    """
    # Unpack the message constants
    message_type, msg_str = message_constants
    # Unpack the ids
    node_id, cid = ids
    ins = gen_ins_function(server_state.current_round, cid)
    client_state = server_state.client_states.get(cid, None)
    assert client_state is not None, f"Client state for {cid} not found"
    ins |= {
        SERVER_ROUND: server_state.current_round,
        CLIENT_ID: cid,
        CLIENT_STATE: str(asdict(client_state)),
        SERVER_STEPS_CUMULATIVE: server_state.server_steps_cumulative,
    }
    recorddict = RecordDict(records={msg_str: ConfigRecord(ins)})
    return Message(
        content=recorddict,
        message_type=message_type,
        dst_node_id=node_id,
        group_id=str(server_state.current_round),
        ttl=mt_dao_TTL,
    )
