"""Utilities for handling messages and context in the node manager app."""

import ast
import copy
from logging import DEBUG
from typing import cast

from flwr.common import (
    ConfigRecord,
    RecordDict,
    log,
)

from mt_dao.constants import (
    AGGREGATION_MASK,
    CLIENT_ID,
    CLIENT_STATE,
    MASKS,
    NUM_SAMPLES,
    mt_dao_ARRAY_RECORD,
    mt_dao_CLIENT_PARAMETERS,
    mt_dao_SERVER_PARAMETERS,
    SERVER_ROUND,
    SERVER_STEPS_CUMULATIVE,
    TRAINING_RESULTS,
    TRANSMISSION_MASK,
)
from mt_dao.flwr_utils import mt_daoRecordDict
from mt_dao.masks_utils import generate_empty_mask
from mt_dao.node_manager.data_types import (
    GatherParametersState,
    SetParametersState,
    TrainEvalInstructions,
    TrainEvalState,
)
from mt_dao.node_manager.node_manager_app import NodeManagerApp
from mt_dao.shm.utils import compress_with_strict
from mt_dao.utils import ClientState, set_parameters_in_place


def transform_record_dict_to_mt_dao(
    record_dict: RecordDict | mt_daoRecordDict,
) -> mt_daoRecordDict:
    """Transform a RecordDict to a mt_daoRecordDict.

    If the record_dict is already a mt_daoRecordDict, it is returned as is.
    If it has `mt_dao_array_records`, it is cast to mt_daoRecordDict.
    If it does not have `mt_dao_array_records`, a new mt_daoRecordDict is created.

    Parameters
    ----------
    record_dict : RecordDict | mt_daoRecordDict
        The record dict to transform.

    Returns
    -------
    mt_daoRecordDict
        The transformed mt_daoRecordDict.

    """
    if isinstance(record_dict, mt_daoRecordDict):
        return record_dict
    if hasattr(record_dict, mt_dao_ARRAY_RECORD):
        # If the record_dict already has mt_dao_array_records, we can just cast it
        return cast("mt_daoRecordDict", record_dict)
    record_dict = mt_daoRecordDict(
        dict(record_dict.items()),
    )
    # Initialize mt_dao_array_records if not present
    if not hasattr(record_dict, mt_dao_ARRAY_RECORD):
        log(
            DEBUG,
            "Initializing mt_dao_array_records in mt_daoRecordDict",
        )
        record_dict.mt_dao_array_records = {}
    return record_dict


def extract_train_eval_state(
    state: mt_daoRecordDict,
) -> TrainEvalState:
    """Extract the record dict for training and evaluation from the node manager.

    Parameters
    ----------
    state : mt_daoRecordDict
        The record dict containing the state of the node manager.

    Returns
    -------
    TrainEvalState
        The extracted context for training.

    """
    # Get transmission mask from the record dict
    state_masks_config = state.config_records.get(
        MASKS,
        None,
    )
    assert state_masks_config is not None, "Masks config must be set"
    state_transmission_mask_str = state_masks_config.get(TRANSMISSION_MASK, None)
    assert (
        state_transmission_mask_str is not None
    ), "extract_train_eval_state: transmission mask must be set in state_masks_config"
    state_transmission_mask = ast.literal_eval(
        str(state_transmission_mask_str),
    )
    # Get the parameters from the record dict
    state_parameters = state.mt_dao_array_records.get(
        mt_dao_SERVER_PARAMETERS,
        None,
    )
    assert state_parameters is not None, "Parameters must be set"

    return TrainEvalState(
        parameters=state_parameters,
        transmission_mask=state_transmission_mask,
    )


def extract_gather_parameters_state(
    state: mt_daoRecordDict,
) -> GatherParametersState:
    """Extract the record dict for gathering parameters from the node manager.

    Parameters
    ----------
    state : mt_daoRecordDict
        The record dict containing the state of the node manager.

    Returns
    -------
    GatherParametersState
        The extracted context for gathering parameters.

    """
    # Get relevant objects from the record dict
    masks_cfg = state.config_records.get(MASKS)
    assert masks_cfg is not None, "Masks config must be set"
    assert masks_cfg.get(AGGREGATION_MASK) is not None, "Aggregation mask must be set"
    aggregation_mask = copy.deepcopy(ast.literal_eval(str(masks_cfg[AGGREGATION_MASK])))

    # Get the parameters from the record dict
    parameters = state.mt_dao_array_records.get(
        mt_dao_CLIENT_PARAMETERS,
        None,
    )
    assert parameters is not None, "Parameters must be set"

    # Get the training results from the record dict
    training_results = state.config_records.get(
        TRAINING_RESULTS,
    )
    assert training_results is not None, "Training results must be set"
    num_samples = training_results.get(NUM_SAMPLES)
    assert num_samples is not None, "Training results do not contain num_samples"
    return GatherParametersState(
        parameters=parameters,
        aggregation_mask=aggregation_mask,
        num_samples=cast("int", num_samples),
    )


def insert_gather_parameters_state(
    state: mt_daoRecordDict,
    gather_parameters_state: GatherParametersState,
) -> None:
    """Insert the gather parameters context into the node manager's context.

    Parameters
    ----------
    state : mt_daoRecordDict
        The record dict containing the state of the node manager.
    gather_parameters_state : GatherParametersState
        The context containing the parameters to gather.

    """
    # Set the parameters in the mt_dao_array_records
    state.mt_dao_array_records[mt_dao_CLIENT_PARAMETERS] = (
        gather_parameters_state.parameters
    )
    # Set the number of samples in the config_records
    if TRAINING_RESULTS not in state.config_records:
        state.config_records[TRAINING_RESULTS] = ConfigRecord(
            {
                NUM_SAMPLES: gather_parameters_state.num_samples,
            },
        )
    else:
        state.config_records[TRAINING_RESULTS][
            NUM_SAMPLES
        ] = gather_parameters_state.num_samples
    # Set the aggregation mask in the config_records
    if MASKS not in state.config_records:
        state.config_records[MASKS] = ConfigRecord(
            {AGGREGATION_MASK: repr(gather_parameters_state.aggregation_mask)},
        )
    else:
        state.config_records[MASKS][AGGREGATION_MASK] = repr(
            gather_parameters_state.aggregation_mask,
        )


def extract_set_parameters_state(
    state: mt_daoRecordDict,
    node_manager_app: NodeManagerApp,
) -> SetParametersState:
    """Extract the record dict for setting parameters in the node manager.

    Parameters
    ----------
    state : mt_daoRecordDict
        The record dict containing the state of the node manager.
    node_manager_app : NodeManagerApp
        The node manager application instance.

    Returns
    -------
    SetParametersState
        The extracted context for setting parameters.

    """
    # Get the transmission mask from the record dict
    state_masks_config = state.config_records.get(MASKS, None)
    state_transmission_mask_str = (
        state_masks_config.get(TRANSMISSION_MASK, None) if state_masks_config else None
    )
    state_transmission_mask: tuple[tuple[bool, ...], list[str], list[str]] | None = None
    if state_transmission_mask_str is not None:
        state_transmission_mask = ast.literal_eval(
            str(state_transmission_mask_str),
        )
    # Get the parameters from the record dict
    state_parameters = state.mt_dao_array_records.get(
        mt_dao_SERVER_PARAMETERS,
        None,
    )
    if state_parameters is None:
        state_parameters = node_manager_app.parameters_metadata.create_ndarrays()

    return SetParametersState(
        parameters=state_parameters,
        transmission_mask=state_transmission_mask
        or generate_empty_mask(
            layer_names_and_types=node_manager_app.layer_names_and_types,
        ),
    )


def insert_set_parameters_state(
    state: mt_daoRecordDict,
    set_parameters_state: SetParametersState,
) -> None:
    """Insert the set parameters context into the node manager's context.

    Parameters
    ----------
    state : mt_daoRecordDict
        The record dict containing the state of the node manager.
    set_parameters_state : SetParametersState
        The context containing the parameters to set.

    """
    # Set the transmission mask in the config_records
    if MASKS not in state.config_records:
        state.config_records[MASKS] = ConfigRecord(
            {TRANSMISSION_MASK: repr(set_parameters_state.transmission_mask)},
        )
    else:
        state.config_records[MASKS][TRANSMISSION_MASK] = repr(
            set_parameters_state.transmission_mask,
        )
    # Set the parameters in the mt_dao_array_records
    state.mt_dao_array_records[mt_dao_SERVER_PARAMETERS] = (
        set_parameters_state.parameters
    )


def extract_fit_eval_ins(
    msg_content: RecordDict,
    state: mt_daoRecordDict,
    msg_str: str,
) -> TrainEvalInstructions:
    """Extract the training or evaluation instructions from the message content.

    Parameters
    ----------
    msg_content : RecordDict
        The content of the message containing the instructions.
    state : mt_daoRecordDict
        The current state of the node manager.
    msg_str : str
        The string identifier for the message type (e.g., "fit_ins", "eval_ins

    Returns
    -------
    TrainEvalInstructions
        The extracted training or evaluation instructions.

    """
    # Get the instructions from the message content
    ins = msg_content.get(msg_str, None)
    assert ins is not None, f"Instructions for {msg_str} must be in the message content"
    assert isinstance(
        ins,
        ConfigRecord,
    ), f"Expected ConfigRecord for {msg_str}, got {type(ins)}"
    # NOTE: This shouldn't be popped as it used later in the app
    server_round = ins.get(SERVER_ROUND, None)
    assert server_round is not None, "Server round must be set in the instructions"
    assert not isinstance(server_round, list), "Server round must not be a list"
    client_id = ins.pop(CLIENT_ID, None)
    assert client_id is not None, "Client ID must be set in the instructions"
    assert not isinstance(client_id, list), "Client ID must not be a list"
    # NOTE: This shouldn't be popped as it used later in the app
    client_state_dict = ins.get(CLIENT_STATE, None)
    assert client_state_dict is not None, "Client state must be set in the instructions"
    assert isinstance(client_state_dict, str), "Client state must be a string"
    client_state = ClientState(**ast.literal_eval(client_state_dict))
    # NOTE: This shouldn't be popped as it used later in the app
    server_step_cumulative = ins.get(SERVER_STEPS_CUMULATIVE, None)
    assert (
        server_step_cumulative is not None
    ), "Server step cumulative must be set in the instructions"
    assert not isinstance(
        server_step_cumulative,
        list,
    ), "Server step cumulative must not be a list"

    # Get the transmission mask from the record dict
    state_masks_config = state.config_records.get(MASKS, None)
    assert state_masks_config is not None, "Masks config must be set in the state"
    state_transmission_mask_str = state_masks_config.get(TRANSMISSION_MASK, None)
    assert (
        state_transmission_mask_str is not None
    ), "extract_fit_eval_ins: transmission mask must be set in state_masks_config"
    # Add the trasmission mask to the instructions as string
    ins[TRANSMISSION_MASK] = str(state_transmission_mask_str)

    return TrainEvalInstructions(
        ins=ins,
        server_round=int(server_round),
        client_id=int(client_id),
        client_state=client_state,
        server_step_cumulative=int(server_step_cumulative),
    )


def get_and_set_parameters_and_transmission_mask(
    state: mt_daoRecordDict,
    node_manager_app: NodeManagerApp,
) -> None:
    """Get the parameters and transmission mask and set them in the NodeManagerApp.

    Parameters
    ----------
    state : mt_daoRecordDict
        The record dict containing the state of the node manager.
    node_manager_app : NodeManagerApp
        The node manager application instance.

    """
    train_state = extract_train_eval_state(state)
    # Set the transmission mask in the app
    node_manager_app.transmission_mask = train_state.transmission_mask
    # Set the parameters in the app
    node_manager_app.create_parameters_shm()
    assert node_manager_app.round_parameters is not None, "Round parameters must be set"
    transmission_parameters = list(
        compress_with_strict(
            data=train_state.parameters,
            selectors=train_state.transmission_mask[0],
            strict=True,
        ),
    )
    set_parameters_in_place(
        list(
            compress_with_strict(
                data=node_manager_app.round_parameters,
                selectors=train_state.transmission_mask[0],
                strict=True,
            ),
        ),
        transmission_parameters,
        parameter_pos=range(len(transmission_parameters)),
    )
