"""Server state management for federated learning."""

import ast
import copy
import inspect
import random
import time
from collections.abc import Callable
from dataclasses import asdict, dataclass
from logging import DEBUG, WARNING
from queue import Queue
from typing import Any

from composer.loggers import RemoteUploaderDownloader
from flwr.common import (
    NDArrays,
    log,
)

import ray
from mt_dao.constants import (
    CLIENT_STATES,
    DEFAULT_MAX_SHARD_SIZE,
    HISTORY,
    KEY_LIGHTWEIGHT_SERVER_STATE,
    LOCAL_STEPS_CUMULATIVE,
    SERVER_ROUND,
    SERVER_STEPS_CUMULATIVE,
    TIME_OFFSET,
)
from mt_dao.server.client_sampler import ClientSampler
from mt_dao.utils import (
    ClientState,
    ModelStateNames,
)
from mt_dao.wandb_history import WandbHistory


class ServerStateDictError(Exception):
    """Exception raised when the server state dictionary is not lightweight."""


@dataclass
class ServerState:
    """Server state for federated learning.

    This class holds the server state, including the model states, history,
    and other relevant information for the federated learning process.
    """

    model_states: NDArrays
    layer_names_and_types: tuple[tuple[str, ModelStateNames], ...]
    transmission_scheduler: Callable[
        [str | int, int],
        list[ModelStateNames] | dict[ModelStateNames, list[str]],
    ]
    history: WandbHistory
    time_offset: float
    server_steps_cumulative: int
    client_states: dict[str | int, ClientState]
    sampled_clients: list[int]
    previously_sampled_clients: list[int]
    start_round: int
    total_rounds: int
    rng: random.Random
    current_round: int
    perf_counter: float
    local_checkpoint_path: str
    client_sampler: ClientSampler
    momentum_vector: NDArrays | None
    second_momentum_vector: NDArrays | None
    state_keys: tuple[str, ...]
    aggregation_mask_scheduler: Callable[
        [str | int, int],
        list[ModelStateNames] | dict[ModelStateNames, list[str]],
    ]
    remote_up_down: RemoteUploaderDownloader | None = None
    params_garbage_queue: Queue[ray.ObjectRef | str] | None = None
    max_shard_size: str = DEFAULT_MAX_SHARD_SIZE

    # NOTE: The max shard size is set as 5GB as per conventions in the HuggingFace
    # transformers library, in their save_pretrained() function.


def create_server_state_lightweight_dict(
    server_state: ServerState,
) -> dict[str, Any]:
    """Create a lightweight dictionary representation of the server state.

    This function extracts essential information from a ServerState object and creates
    a reduced dictionary containing only the key components needed for checkpointing
    and state recovery. The lightweight dictionary excludes large model parameters and
    other memory-intensive components, focusing on elements like the current round,
    history, elapsed time, client states, and cumulative server steps.

    Parameters
    ----------
    server_state : ServerState
        The complete server state object containing all federated learning state
        information, including model parameters, client states, round information,
        and other tracking metadata.

    Returns
    -------
    dict[str, Any]
        A lightweight dictionary containing only essential server state components:
        - server_round: The current round number
        - history: The WandbHistory object tracking metrics
        - time_offset: The current elapsed time in the federated learning process
        - client_states: String representation of client states dictionary
        - server_steps_cumulative: The total number of steps taken by the server

    Notes
    -----
    The lightweight dictionary is used for efficient serialization during checkpointing
    and when transmitting server state information. The client_states are converted to
    a string representation to ensure compatibility with various serialization methods.

    """
    return {
        SERVER_ROUND: server_state.current_round,
        HISTORY: server_state.history,
        # NOTE: Using the server_state.perf_counter to get the time offset
        TIME_OFFSET: (
            server_state.time_offset + time.perf_counter() - server_state.perf_counter
        ),
        CLIENT_STATES: str(
            {k: asdict(v) for k, v in server_state.client_states.items()},
        ),
        SERVER_STEPS_CUMULATIVE: server_state.server_steps_cumulative,
    }


def read_server_state_lightweight_dict(
    server_state_dict: dict[str, Any],
    n_total_clients: int,
    n_local_steps: int,
) -> tuple[WandbHistory, int, float, int, dict[str | int, ClientState]]:
    """Parse a lightweight server state dictionary into its component parts.

    This function extracts and reconstructs the essential components of a server state
    from a lightweight dictionary representation. It retrieves the current round,
    history object, client states, time offset, and server steps count. If certain
    components are missing from the dictionary, the function generates appropriate
    default values to maintain compatibility with previous versions.

    Parameters
    ----------
    server_state_dict : dict[str, Any]
        A lightweight dictionary containing the serialized server state components.
        Expected to contain keys like SERVER_ROUND, HISTORY, CLIENT_STATES,
        TIME_OFFSET, and SERVER_STEPS_CUMULATIVE.
    n_total_clients : int
        The total number of clients in the federated learning process. Used to
        generate default client states if they are missing from the dictionary.
    n_local_steps : int
        The number of local steps per round. Used in conjunction with start_round
        to calculate default values for client states if needed.

    Returns
    -------
    tuple[WandbHistory, int, float, int, dict[str | int, ClientState]]
        A tuple containing:
        - history: The WandbHistory object tracking metrics
        - start_round: The current round number
        - time_offset: The elapsed time in the federated learning process
        - server_steps_cumulative: The total number of steps taken by the server
        - client_states: Dictionary mapping client IDs to their respective states

    Notes
    -----
    The function provides backward compatibility by generating default values when
    certain components are missing. For example, if client states are not found in
    the dictionary, it creates dummy client states with appropriate cumulative step
    counts based on the current round and local steps configuration.

    """
    # Get the server round
    start_round = int(server_state_dict.pop(SERVER_ROUND))
    # Get the history
    history: WandbHistory = server_state_dict.pop(HISTORY)
    # Read the server state dictionary
    if CLIENT_STATES in server_state_dict:
        saved_client_states: dict[str | int, dict[str, Any]] = ast.literal_eval(
            server_state_dict.pop(CLIENT_STATES),
        )
    else:
        # NOTE: This `LOCAL_STEPS_CUMULATIVE` is just used for backlogging and not
        # by any logic during training so we put a zero for now.
        log(
            WARNING,
            "No client state found in the checkpoint. "
            "We will put dummy values (%s) for the %s clients.",
            int(n_local_steps * start_round),
            n_total_clients,
        )
        saved_client_states = {
            cid: {LOCAL_STEPS_CUMULATIVE: n_local_steps * start_round}
            for cid in range(n_total_clients)
        }
    # NOTE: We maintain partial compatibility across ClientState implementations
    # by only loading the attributes we actually need
    # this decision should be revisited at a late time

    # Automatically deduce the existing args from the ClientState class
    client_states_args = inspect.signature(ClientState.__init__).parameters.keys()

    # Filter and create client_states dictionary in one line
    client_states = {
        k: ClientState(**{attr: v[attr] for attr in v if attr in client_states_args})
        for k, v in saved_client_states.items()
    }
    # Get the time offset
    time_offset = 0.0
    if TIME_OFFSET in server_state_dict:
        time_offset = server_state_dict.pop(TIME_OFFSET)
    # Get the server steps cumulative
    if SERVER_STEPS_CUMULATIVE in server_state_dict:
        server_steps_cumulative = server_state_dict.pop(SERVER_STEPS_CUMULATIVE)
    else:
        # Make it back compatible with the previous versions
        server_steps_cumulative = max(
            *[
                _client_states.local_steps_cumulative
                for _client_states in client_states.values()
            ],
            0,
        )
    # Check if there are any keys left uninterpreted
    if server_state_dict:
        log(
            WARNING,
            "Server state keys %s are not consistent",
            server_state_dict.keys(),
        )
    # Return the server state components
    return (
        history,
        start_round,
        time_offset,
        server_steps_cumulative,
        client_states,
    )


def check_server_state_lightweight_dict(
    server_state_dict: dict[str, Any],
    *,
    raise_if_not_lightweight: bool = False,
) -> bool:
    """Check if a server state dictionary contains only lightweight keys.

    This function verifies that a server state dictionary only contains the expected
    keys for a lightweight representation. It validates that all keys in the provided
    dictionary are included in the predefined set of lightweight server state keys
    (KEY_LIGHTWEIGHT_SERVER_STATE) and that there are no additional keys present.

    Parameters
    ----------
    server_state_dict : dict[str, Any]
        The server state dictionary to check for lightweight compliance. This should
        be a dictionary containing server state information, typically created by
        the `create_server_state_lightweight_dict` function.
    raise_if_not_lightweight : bool, optional
        If True and the dictionary is found not to be lightweight, raises a
        ServerStateDictError exception. Default is False, which only returns the
        boolean result without raising an exception.

    Returns
    -------
    bool
        True if the server state dictionary contains only the expected lightweight
        keys, False otherwise.

    Raises
    ------
    ServerStateDictError
        If raise_if_not_lightweight is True and the server state dictionary contains
        unexpected keys or is missing expected keys.

    Notes
    -----
    A lightweight server state dictionary is expected to contain only the specific
    keys defined in KEY_LIGHTWEIGHT_SERVER_STATE. This is important for efficient
    serialization and storage of server state information during checkpointing in
    the federated learning process.

    """
    # Initialize the return value
    is_lightweight = True
    # Copy to avoid modifying the original dictionary
    dummy_server_state = copy.deepcopy(server_state_dict)
    # Check if the server state is lightweight
    for key in server_state_dict:
        if key not in KEY_LIGHTWEIGHT_SERVER_STATE:
            log(
                DEBUG,
                "Server state key %s is not consistent",
                key,
            )
            is_lightweight = False
        dummy_server_state.pop(key)
    # Check if there are any keys left
    if dummy_server_state:
        log(
            DEBUG,
            "Server state keys %s are not consistent",
            dummy_server_state.keys(),
        )
        is_lightweight = False
    # Raise an error if the server state is not lightweight
    if not is_lightweight and raise_if_not_lightweight:
        raise ServerStateDictError
    return is_lightweight
