"""Module for various utility functions used in mt_dao clients."""

import atexit
import copy
import gc
import random
from dataclasses import asdict
from logging import DEBUG
from typing import TYPE_CHECKING, cast

import numpy as np
import streaming
from composer import Trainer
from composer.utils import dist, reproducibility
from flwr.common.logger import log
from flwr.common.typing import NDArrays, Scalar
from omegaconf import DictConfig
from streaming.base.shared.memory import SharedMemory, shared_memory_list

from mt_dao.clients.configs import FitConfig
from mt_dao.composer_trainer_utils.composer_trainer_getters import (
    calculate_trainer_post_training_l2_norms,
    get_trainer_model_states,
    update_client_state,
)
from mt_dao.composer_trainer_utils.shared_utils import filter_expert_parameters
from mt_dao.conf.base_schema import BaseConfig
from mt_dao.constants import CLIENT_STATE_ACCUMULATOR
from mt_dao.utils import (
    ClientState,
    get_random_model_states_from_config,
)

if TYPE_CHECKING:

    from mt_dao.conf.base_schema import BaseConfig


def get_client_state_struct(fit_config: FitConfig) -> ClientState:
    """Retrieve the state of a specific client.

    Parameters
    ----------
    fit_config : FitConfig
        The configuration containing the client states and server steps.

    Returns
    -------
    ClientState: The state structure of the client.

    Raises
    ------
    ValueError: If the client states or server steps are not provided in the
        configuration

    """
    # Retrieve the clients' states
    if fit_config.client_state is None:
        msg = "Client states must be provided."
        raise ValueError(msg)
    if fit_config.server_steps_cumulative is None:
        msg = "Server steps must be provided."
        raise ValueError(msg)

    # Extract current client's state
    return fit_config.client_state


def post_process_client_result(  # noqa: PLR0913, PLR0917
    train_metrics: dict[str, Scalar],
    client_state_struct: ClientState,
    llm_config: DictConfig,
    trainer: Trainer,
    payload: NDArrays,
    fit_config: FitConfig,
) -> tuple[NDArrays, int]:
    """Compute post-training metrics on the client result.

    This function retrieves the model parameters after training, calculates the number
    of samples trained, updates the client state, and collects various training metrics.

    Parameters
    ----------
    train_metrics : dict[str, Scalar]
        A dictionary to store training metrics.
    client_state_struct : ClientState
        The state structure of the client.
    llm_config : DictConfig
        The configuration dictionary for the local model.
    trainer : Trainer
        The trainer instance used for training.
    payload : NDArrays
        The server payload, to be used for computting metrics.
    fit_config : FitConfig
        The configuration containing the client states and server steps.

    Returns
    -------
    tuple[NDArrays, int]
        A tuple containing:
        - The model parameters after training.
        - The number of samples trained.

    """
    # Obtain the model states from the trainer
    model_states = get_trainer_model_states(
        trainer=trainer,
        fit_config=fit_config,
        train_metrics=train_metrics,
    )

    # Only rank 0 collects metrics related to the pseudo gradients
    if dist.get_global_rank() == 0:
        calculate_trainer_post_training_l2_norms(
            model_states=model_states,
            train_metrics=train_metrics,
            fit_config=fit_config,
            payload=payload,
        )

    # Update client state
    update_client_state(
        client_state=client_state_struct,
        trainer=trainer,
        fit_config=fit_config,
    )
    # Add the client state to the training metrics
    train_metrics |= {
        CLIENT_STATE_ACCUMULATOR: str(
            {fit_config.cid: asdict(client_state_struct)},
        ),
    }
    # Retrieve training metrics
    train_metrics |= {
        k: v.detach().cpu().item()  # type: ignore[attr-defined]
        for k, v in trainer.state.train_metric_values.items()
    }

    # Retrieve model states as a full payload
    model_states_list: NDArrays = []
    for model_state in model_states.values():
        model_states_list.extend(model_state)

    # Filter from the mask experts that this client is not supposed to train
    # NOTE: We need now to change the real aggregation masks by modifying in place the
    # mask so that the Worker-NodeManagerApp-ServerApp pipeline is consistent with the
    # global model
    if fit_config.experts_ids is not None:
        # Get the aggregation mask
        (mask_aggregation, layer_names_aggregation, layer_types_aggregation) = (
            fit_config.aggregation_mask
        )
        # Assign the filtered mask to the fit_config
        fit_config.aggregation_mask = filter_expert_parameters(
            mask_aggregation,
            layer_names_aggregation,
            layer_types_aggregation,
            fit_config.experts_ids,
            change_names=False,
        )

    log(
        DEBUG,
        "Client payload size is %s for rank-%s",
        len(model_states_list),
        dist.get_global_rank(),
    )
    return model_states_list, client_state_struct.steps_done * int(
        llm_config["global_train_batch_size"],
    )


def streaming_shms_clean_up() -> None:
    """Clean up leaking and stale shared memories.

    This function performs the following clean-up tasks:
    - Cleans up leaking shared memories from the shared_memory_list.
    - Un-registers the cleanup function from atexit.
    - Clears the shared_memory_list.
    - Cleans up stale shared memory using the streaming library.
    - Runs the garbage collector to free up memory.
    """
    # NOTE: Clean up leaking shared memories
    for shm in shared_memory_list:
        SharedMemory.cleanup(shm)
        atexit.unregister(SharedMemory.cleanup)
    shared_memory_list.clear()
    # Cleaning stale shared memory
    streaming.base.util.clean_stale_shared_memory()  # type: ignore[reportAttributeAccessIssue]
    # Clean-up garbage collector
    gc.collect()


def randomize_layers(  # noqa: PLR0913
    parameters: NDArrays,
    dummy_config: DictConfig,
    names: list[str],
    random_layers: list[str],
    cid: int = 0,
    *,
    server_round: int = 1,
    truly_random_init: bool,
) -> None:
    """Randomize specified layers of the model parameters.

    This function randomizes the specified layers of the model parameters based on the
    provided configuration. It can use a truly random initialization if specified.

    Parameters
    ----------
    parameters : NDArrays
        The list of model parameters to be randomized.
    dummy_config : DictConfig
        The configuration object used to create a dummy model for randomization.
    names : list[str]
        The list of parameter names corresponding to the model parameters.
    random_layers : list[str]
        The list of layer names to be randomized.
    truly_random_init : bool
        If True, uses a truly random initialization for the specified layers.
    cid : int, optional
        The client ID used for seeding the random initialization (default is 0).
    server_round : int, optional
        The current server round used for seeding the random initialization
        (default is 1).


    Example
    -------
    >>> from omegaconf import OmegaConf
    >>> parameters = [...]
    >>> dummy_config = OmegaConf.create({...})
    >>> names = ["layer1.weight", "layer1.bias", "layer2.weight"]
    >>> random_layers = ["layer1.weight"]
    >>> randomize_layers(
    >>>     parameters, dummy_config, names, random_layers,
    >>>     truly_random_init=True, cid=1, server_round=2
    >>> )

    """
    new_dummy_config = copy.deepcopy(dummy_config)
    if truly_random_init:
        new_seed = 51550
        for _ in range(server_round):
            new_seed = random.randint(0, 2**32 - 1) ^ cid  # noqa: S311
        new_dummy_config.global_seed = new_seed
        new_dummy_config.seed = new_seed
        reproducibility.seed_all(new_seed)
        log(DEBUG, f"Randomizing layers with seed {new_seed}")

    # Guarantee it is false for pre-trained models
    new_dummy_config.model.pretrained = False
    tmp_dummy_config: BaseConfig = cast(
        "BaseConfig",
        DictConfig(
            {
                "pretrained_model_path": None,
                "llm_config": new_dummy_config,
            },
        ),
    )
    log(DEBUG, f"Creating random model with this config: {tmp_dummy_config}")

    random_parameters, _ = get_random_model_states_from_config(tmp_dummy_config)

    indices = [names.index(key) for key in random_layers]
    for i in indices:
        np.copyto(parameters[i], random_parameters[i])

    log(DEBUG, f"Randomized layers: {random_layers} with indices: {indices}")
