"""Data types for the mt_dao node manager."""
from dataclasses import dataclass

from flwr.common import (
    ConfigRecord,
    NDArrays,
)

from mt_dao.utils import ClientState


@dataclass
class TrainEvalInstructions:
    """Data class for training and evaluation instructions."""

    ins: ConfigRecord
    server_round: int
    client_id: int
    client_state: ClientState
    server_step_cumulative: int


@dataclass
class TrainEvalState:
    """Context for training and evaluation in the node manager."""

    parameters: NDArrays
    transmission_mask: tuple[tuple[bool, ...], list[str], list[str]]


@dataclass
class GatherParametersState:
    """Context for gathering parameters from the node manager."""

    parameters: NDArrays
    aggregation_mask: tuple[tuple[bool, ...], list[str], list[str]]
    num_samples: int


@dataclass
class SetParametersState:
    """Context for setting parameters in the node manager."""

    parameters: NDArrays
    transmission_mask: tuple[tuple[bool, ...], list[str], list[str]]
