"""Provides the internal functions used by the LLM client.

This module contains utility functions that handle local training and evaluation for
Large Language Model (LLM) clients in a federated learning setup. It sets up and
configures Trainers, loads checkpoints, applies parameters, runs training or evaluation,
gathers metrics, and cleans up resources.
"""

import copy
import time
from logging import DEBUG, ERROR
from typing import Any

import torch

# NOTE: We need this if we want to compile the model because the attention
# implementation in the MPT code is dispatched using a dictionary that raises:
# `AssertionError: Dict types must use ConstDictVariable.`
import torch._dynamo
from composer import Trainer
from composer.utils import dist
from flwr.common.logger import log
from flwr.common.record import ConfigRecord
from flwr.common.typing import NDArrays, Scalar
from llmfoundry.utils.config_utils import TrainConfig
from omegaconf import DictConfig

from mt_dao.clients.configs import EvaluateConfig, FitConfig
from mt_dao.clients.llm_config_functions import (
    check_skip_init,
    get_train_config,
    set_initial_config_from_fit_config,
)
from mt_dao.clients.trainer_utils import (
    correct_time_before_fit,
    get_trainer_mutables_from_config,
    get_trainer_object,
    load_trainer_checkpoint,
    set_mutables_trainer,
    trainer_clean_up,
)
from mt_dao.clients.utils import (
    get_client_state_struct,
    post_process_client_result,
    streaming_shms_clean_up,
)
from mt_dao.composer_trainer_utils.composer_trainer_setters import (
    set_trainer_states_for_eval,
    set_trainer_states_for_fit,
)
from mt_dao.masks_utils import ModelStateNames

torch._dynamo.config.suppress_errors = True  # type: ignore[reportAttributeAccessIssue]  # noqa: SLF001


def llm_fit(  # noqa: PLR0914, PLR0915
    external_objects: tuple[
        Trainer | None,
        FitConfig | EvaluateConfig | None,
        TrainConfig | None,
    ],
    payload: NDArrays,
    config: ConfigRecord,
    llm_config: DictConfig,
) -> tuple[
    NDArrays,
    int,
    dict[str, Scalar] | dict[Any, Any],
    tuple[tuple[bool, ...], list[str], list[ModelStateNames]],
    tuple[Trainer | None, FitConfig | EvaluateConfig | None, TrainConfig | None],
]:
    """Perform local training on the LLM client.

    This function sets up the model and trainer environment, handles checkpoint
    loading if available, sets model parameters, performs local training, and
    returns the updated model parameters, the number of trained samples, training
    metrics, and the configuration objects used for the training process.

    The function uses a skip initialization optimization that avoids recreating the
    trainer's mutable attributes and reloading checkpoints when the same client is
    used across consecutive rounds.

    Parameters
    ----------
    external_objects : tuple[
        Trainer | None,
        FitConfig | EvaluateConfig | None,
        TrainConfig | None
    ]
        A tuple containing:
        - An optional external Trainer object. If provided, it will be used instead
          of creating a new Trainer.
        - An optional external FitConfig or EvaluateConfig from a previous round.
        - An optional external TrainConfig from a previous round.
        These objects are used to optimize initialization when the same client is used
        in consecutive rounds.
    payload : NDArrays
        The model parameters to be loaded into the trainer before training.
    config : ConfigRecord
        Configuration information for federated training.
    llm_config : DictConfig
        The local LLM configuration specifying trainer and model details.

    Returns
    -------
    tuple[
        NDArrays,
        int,
        dict[str, Scalar] | dict[Any, Any],
        tuple[tuple[bool, ...], list[str], list[ModelStateNames]],
        tuple[Trainer | None, FitConfig | EvaluateConfig | None, TrainConfig | None]
    ]
        A tuple containing:
        - The updated model parameters after training.
        - The number of samples used during training.
        - A dictionary of training metrics.
        - The aggregation mask used for the training round, which includes:
          - A tuple of booleans indicating which parameters were updated.
          - A list of layer names that were updated.
          - A list of layer types corresponding to the updated parameters.
        - A tuple containing the final Trainer object, FitConfig, and TrainConfig for
          potential reuse in future rounds.

    """
    # Interpret external objects
    external_trainer, ext_fit_config, ext_train_cfg = external_objects
    try:
        fit_config = FitConfig(**config)  # type: ignore[reportArgumentType,arg-type]

        assert fit_config.transmission_mask is not None, (
            "llm_fit: transmission mask must be set"
        )

        client_state_struct, start_time = (
            get_client_state_struct(fit_config),
            time.perf_counter(),
        )
        train_metrics: dict[str, Scalar] = {}

        skip_iteration, checkpoint_exists, server_steps_cumulative = (
            set_initial_config_from_fit_config(fit_config, llm_config)
        )
        dummy_config = copy.deepcopy(llm_config)

        if external_trainer is not None:
            # NOTE: This is a workaorund that pairs well with our fork for the
            # `streaming` library. Why is it necessary? In short, our modified
            # `StreamingDataset` class makes its `self.__iter__` function returning an
            # actual iterator instead of an object that re-samples with replacement
            # every time the iterator exits and its called again. As a result, the
            # `Trainer._train_loop` function spins the iterator to get to the
            # `batch_idx` (received from an `enumerate` that starts from 0) equal to the
            # desired timestamp, thus skipping batches. We can prevent the spinning from
            # happening by setting the `dataset_resumption` flag to True.
            external_trainer.state.dataset_resumption = {"train": True}
            log(DEBUG, "External trainer object exists.")
            train_cfg, device, logged_cfg, icl_tasks_config_dict = get_train_config(
                cfg=llm_config,
                cid=fit_config.cid,
            )
            # Performs the heavy initialization only if strictly necessary
            if not check_skip_init(
                ext_fit_config,
                ext_train_cfg,
                fit_config,
                train_cfg,
            ):
                # NOTE: If we need to re-create data loaders, then we must let the
                # Trainer object handle the spinning up to the correct timestamp.
                external_trainer.state.dataset_resumption.pop("train")
                # Assign the device microbatch size from the previous execution
                assert ext_train_cfg is not None
                train_cfg.device_train_microbatch_size = (
                    ext_train_cfg.device_train_microbatch_size
                )
                trainer_mutable_attributes = get_trainer_mutables_from_config(
                    trainer=external_trainer,
                    train_cfg=train_cfg,
                    client_config=fit_config,
                    icl_tasks_config_dict=icl_tasks_config_dict,
                    device=device,
                    logged_cfg=logged_cfg,
                )
                set_mutables_trainer(
                    external_trainer,
                    trainer_mutable_attributes,
                    client_config=fit_config,
                )
                if checkpoint_exists:
                    log(DEBUG, "Checkpoint exists.")
                    # TODO(<anon>): Add to the ignore keys all the keys related to
                    # model parameters
                    load_trainer_checkpoint(external_trainer, train_cfg)
        else:
            log(DEBUG, "External trainer object doesn't exit.")
            (
                external_trainer,
                train_cfg,
                _,
            ) = get_trainer_object(
                cfg=llm_config,
                client_config=fit_config,
            )

        log(DEBUG, "Trainer object obtained.")
        train_metrics["client/fit_init_time"] = time.perf_counter() - start_time

        if not skip_iteration and fit_config.n_local_steps > 0:
            correct_time_before_fit(
                trainer=external_trainer,
                fit_config=fit_config,
                server_steps_cumulative=server_steps_cumulative,
                client_state_struct=client_state_struct,
            )

            start_time = time.perf_counter()
            set_trainer_states_for_fit(
                payload=payload,
                trainer=external_trainer,
                configs=(fit_config, dummy_config),
                client_state_struct=client_state_struct,
                train_metrics=train_metrics,
            )
            # NOTE: We removed here the parameter checkers for efficiency reasons.
            # Re-add them if needed
            train_metrics["client/fit_set_model_states_time"] = (
                time.perf_counter() - start_time
            )

            if train_cfg.eval_first:
                # TODO(<anon>): Modify metric names in MosaicML
                # to allow logging the first eval
                start_time = time.perf_counter()
                external_trainer.eval()
                train_metrics["client/fit_pre_eval_time"] = (
                    time.perf_counter() - start_time
                )
                train_metrics.update(
                    {
                        f"PrePersonalization{k}": v.detach().cpu().item()  # type: ignore[attr-defined]
                        for k, v in external_trainer.state.eval_metric_values.items()
                    },
                )

            try:
                start_time = time.perf_counter()
                external_trainer.fit(
                    duration=(
                        0 if skip_iteration else str(fit_config.n_local_steps) + "ba"
                    ),
                )
                # If automicrobatching is enabled, try to get the final microbatch size
                # and set it for the next steps
                if external_trainer.state.auto_microbatching:
                    external_trainer.state.auto_microbatching = False
                    final_device_train_microbatch_size = (
                        external_trainer.state.device_train_microbatch_size
                    )
                    assert isinstance(final_device_train_microbatch_size, int)
                    train_cfg.device_train_microbatch_size = (
                        final_device_train_microbatch_size
                    )
                train_metrics["client/fit_time"] = time.perf_counter() - start_time
            except Exception as e:
                log(ERROR, "llm_fit::trainer.fit", exc_info=e, stack_info=True)
                raise

        trained_model_states, n_samples_trained = post_process_client_result(
            train_metrics=train_metrics,
            trainer=external_trainer,
            payload=payload,
            client_state_struct=client_state_struct,
            llm_config=llm_config,
            fit_config=fit_config,
        )

    except Exception as e:
        log(ERROR, "Error in llm_fit function", exc_info=e, stack_info=True)
        raise
    else:
        return (
            trained_model_states,
            n_samples_trained or 1,
            train_metrics,
            fit_config.aggregation_mask,
            (external_trainer, fit_config, train_cfg),
        )


def llm_eval(
    external_trainer: Trainer | None,
    payload: NDArrays,
    config: ConfigRecord,
    llm_config: DictConfig,
) -> tuple[float, int, dict[str, Scalar], Trainer]:
    """Perform local evaluation on the LLM client.

    This function sets up the trainer environment for evaluation, applies model
    parameters, runs the evaluation, collects evaluation metrics, and cleans up
    resources.

    Parameters
    ----------
    external_trainer : Trainer | None
        An optional external Trainer object. If provided, it will be used instead
        of creating a new Trainer.
    payload : NDArrays
        The model parameters to be used for evaluation.
    config : ConfigRecord
        Configuration information for federated evaluation.
    llm_config : DictConfig
        The local LLM configuration specifying trainer and model details.

    Returns
    -------
    tuple[float, int, dict[str, Scalar], Trainer]
        A tuple containing:
        - The (unused) loss value set to 0.0 in this function.
        - The number of samples evaluated.
        - A dictionary of evaluation metrics.
        - The final Trainer object.

    """
    start_time = time.perf_counter()
    num_samples = 0
    eval_metrics: dict[str, Scalar] = {}

    llm_config.autoresume = False  # type: ignore[union-attr]
    llm_config.save_folder = None  # type: ignore[union-attr]
    llm_config.load_path = None  # type: ignore[union-attr]
    llm_config.loggers = None  # type: ignore[union-attr]

    client_eval_config = EvaluateConfig(**dict(config))  # type: ignore[reportArgumentType,arg-type]

    (
        external_trainer,
        _,
        _,
    ) = get_trainer_object(
        cfg=llm_config,
        client_config=client_eval_config,
    )

    eval_metrics["client/eval_init_time"] = time.perf_counter() - start_time

    start_time = time.perf_counter()
    set_trainer_states_for_eval(
        payload=payload,
        trainer=external_trainer,
        eval_config=client_eval_config,
    )
    eval_metrics["client/eval_set_parameters_time"] = time.perf_counter() - start_time

    start_time = time.perf_counter()
    external_trainer.eval()
    eval_metrics["client/eval_time"] = time.perf_counter() - start_time
    start_time = time.perf_counter()
    if dist.get_global_rank() == 0:
        num_samples = (
            external_trainer.state.eval_timestamp._sample.value  # noqa: SLF001
        )
        eval_metrics.update(
            {
                "Val" + k: v.detach().cpu().item()  # type: ignore[attr-defined]
                for k, v in external_trainer.state.eval_metric_values.items()
            },
        )
        eval_metrics["client/eval_metrics_collection_time"] = (
            time.perf_counter() - start_time
        )

    start_time = time.perf_counter()
    trainer_clean_up(trainer=external_trainer)
    streaming_shms_clean_up()
    if dist.get_global_rank() == 0:
        eval_metrics["client/eval_trainer_closing_time"] = (
            time.perf_counter() - start_time
        )

    return 0.0, num_samples, eval_metrics, external_trainer
