"""Utility functions for initialization task on main server loop in flwr next."""

import copy
import os
import random
import time
from logging import DEBUG, INFO
from typing import Any

import numpy as np
from composer.loggers import RemoteUploaderDownloader
from flwr.common import (
    NDArrays,
    log,
)
from omegaconf import OmegaConf

from mt_dao.clients.configs import CentralizedConfig
from mt_dao.clients.llm_client_functions import (
    get_trainer_object,
)
from mt_dao.conf.base_schema import BaseConfig, StrategyName
from mt_dao.constants import (
    CURRENT_SERVER_STATE_FILENAME,
)
from mt_dao.file_utils import create_remote_up_down
from mt_dao.initialization_utils.helper import (
    extract_layer_names_and_experts,
    retrieve_checkpoint_states,
    retrieve_model_states_and_types,
)
from mt_dao.server.client_sampler import ClientSampler
from mt_dao.server.param_scheduler_dispatcher import dispatch_model_state_scheduler
from mt_dao.server.s3_utils import (
    copy_old_checkpoints_to_new_run,
    get_server_checkpoint,
    interpret_resume_round,
    upload_server_checkpoint_model_states,
    upload_server_checkpoint_server_state,
)
from mt_dao.server.utils import (
    ServerState,
)
from mt_dao.strategy_functional.dispatcher import dispatch_strategy_state_keys
from mt_dao.utils import (
    ClientState,
    ModelStateNames,
    get_random_model_states_from_config,
)
from mt_dao.wandb_history import WandbHistory


def initialize_server_state(  # noqa: C901
    cfg: BaseConfig,
    mt_dao_save_path: str,
) -> ServerState:
    """Initialize the server state for federated learning.

    This function prepares the server state by either resuming from a past checkpoint
    or creating a fresh state. It handles the creation of the RemoteUploaderDownloader
    for checkpointing operations, manages state restoration logic, and initializes the
    federated learning strategy.

    Args:
    ----
    cfg : BaseConfig
        The configuration object containing settings for federated learning and system
        behavior, including checkpoint and S3 communication parameters.
    mt_dao_save_path : str
        The local path the server uses to store and load checkpoints.

    Returns:
    -------
    ServerState
        The server state object, either resumed from a checkpoint or newly initialized,
        with the configured strategy.

    Raises:
    ------
    ValueError
        If checkpointing is needed but not enabled, or if a RemoteUploaderDownloader
        is required but not available.

    """
    # Create RemoteUploaderDownloader
    remote_up_down: RemoteUploaderDownloader | None = None

    def generate_remote_up_down_object() -> RemoteUploaderDownloader:
        """Generate a remote downloader object to communicate with S3.

        Returns
        -------
        RemoteUploaderDownloader
            Object to manage I/O with remote S3 bucket.

        """
        return create_remote_up_down(
            bucket_name=cfg.s3_comm_config.bucket_name,
            prefix=f"{cfg.run_uuid}/server",
            run_uuid=cfg.run_uuid,
            num_attempts=cfg.s3_comm_config.num_attempts,
            client_config=OmegaConf.to_container(
                cfg.s3_comm_config.backend_kwargs.client_config,
            ),  # type: ignore[reportArgumentType, arg-type]
        )

    if cfg.mt_dao.local_init and cfg.mt_dao.remote_init:
        msg = "We cannot have local and remote initialisation simulatenously."
        raise ValueError(msg)
    if cfg.mt_dao.remote_init:
        log(INFO, "Creating remote up/down for initialising from remote state.")
        remote_up_down = generate_remote_up_down_object()
    elif cfg.mt_dao.local_init:
        log(INFO, "Init. from locally available checkpoints.")
    # Resume experiment from a previously saved checkpoint
    if cfg.mt_dao.resume_round is not None:
        if (cfg.mt_dao.local_init is None) or (cfg.mt_dao.remote_init is None):  # type: ignore[reportUnnecessaryComparison]
            msg = "Cannot resume if local or remote initialisation is not enabled."
            raise ValueError(msg)
        server_state = resume_from_round(
            cfg=cfg,
            remote_up_down=remote_up_down,
            mt_dao_save_path=mt_dao_save_path,
        )
    else:
        server_state = initialize_federated_learning(
            cfg=cfg,
            remote_up_down=remote_up_down,
            mt_dao_save_path=mt_dao_save_path,
        )
        log(
            DEBUG,
            "Sampled %s Client IDs: %s",
            len(server_state.sampled_clients),
            server_state.sampled_clients,
        )
    # Loop over the PRNG to get to the correct round
    for _ in range(server_state.start_round):
        server_state.previously_sampled_clients = server_state.sampled_clients
        server_state.sampled_clients = server_state.client_sampler.sample_clients(
            rng=server_state.rng,
        )
    if cfg.mt_dao.restore_cent.checkpoint_dir:
        log(
            INFO,
            "Restoring from a centralized checkpoint at %s",
            cfg.mt_dao.restore_cent.checkpoint_dir,
        )
        cent_model_state, cent_layer_names = restore_run_from_centralised_checkpoint(
            centralised_checkpoint_dir=cfg.mt_dao.restore_cent.checkpoint_dir,
            cent_run_uuid=cfg.mt_dao.restore_cent.run_uuid,
            cfg=copy.deepcopy(cfg),
            preserve_optimizer_states=cfg.mt_dao.restore_cent.preserve_optimizer_states,
        )
        server_state.model_states = cent_model_state
        server_state.layer_names_and_types = cent_layer_names
        if cfg.mt_dao.restore_cent.run_batches is not None:
            server_state.server_steps_cumulative = cfg.mt_dao.restore_cent.run_batches
            log(
                INFO,
                "Setting server steps to %s",
                server_state.server_steps_cumulative,
            )

    # TODO(anon): Make this useful
    # # Exclude from server's model states and layer names and types those types which
    # # are not going to be scheduled for synchronization
    # server_state.model_states, server_state.layer_names_and_types = (
    #     reconcile_model_state_with_scheduler(
    #         server_state.model_states, server_state.layer_names_and_types, cfg))

    # Save the checkpoint to S3 Object Store (w/ model parameters)
    server_state.current_round = server_state.start_round
    server_state.perf_counter = server_state.time_offset
    # NOTE: Determine if we need the remote up/down to persist.
    if cfg.mt_dao.checkpointing.remote_checkpoint or cfg.mt_dao.comm_stack.s3:
        if not server_state.remote_up_down:
            # NOTE: Regenerate the remote-up-down in case we didn't have it for
            # initialisation, but need it for checkpointing.
            server_state.remote_up_down = generate_remote_up_down_object()
    else:
        server_state.remote_up_down = None
    if (
        cfg.mt_dao.checkpointing.local_checkpoint
        or cfg.mt_dao.checkpointing.remote_checkpoint
    ):
        upload_server_checkpoint_server_state(server_state=server_state)
        upload_server_checkpoint_model_states(
            server_state=server_state,
        )

    return server_state


def restore_run_from_centralised_checkpoint(
    centralised_checkpoint_dir: str,
    cfg: BaseConfig,
    cent_run_uuid: str,
    *,
    preserve_optimizer_states: bool,
) -> tuple[NDArrays, tuple[tuple[str, ModelStateNames], ...]]:
    """Retrieve model parameters from a previous centralized training run.

    This function finds a checkpoint from a centralized training run stored in S3
    or a local directory. In the case the checkpoint is a .pt file, the checkpoint needs
    to match the desired number of training batches. It loads that checkpoint into a
    temporary trainer, extracts both the model parameters and optimizer states, if
    asked to be retained, and returns them in the format needed for federated learning
    initialization. In the case of checkpoints being .npz or .safetensors files,
    the checkpoints are downloaded manually, and loaded manually as requested by th
    above specifications.

    Parameters
    ----------
    centralised_checkpoint_dir : str
        The directory where the centraised checkpoint is located.
        This should be the parent directory of the file itself.
        In the case of the .safetensors, it should be the directory
        above the "current_server_parameters" level.
    cfg : BaseConfig
        The configuration object containing information about the centralized run to
        restore, including the run UUID and desired number of batches.
    cent_run_uuid : str
        The run_uuid of the centralised checkpoint. This is only needed if the
        loading from an s3 bucket, where the checkpoints are not .pt files.
    preserve_optimizer_states : bool
        Whether to preserve the saved optimizer states.

    Returns
    -------
    tuple[NDArrays, tuple[tuple[str, ModelStateNames], ...]]
        A tuple containing:
        - model_states: The model parameters and optimizer states as NDArrays
        - layer_names_and_types: A tuple mapping parameter names to their state types

    """
    cfg = copy.deepcopy(cfg)
    llm_config = cfg.llm_config
    os.environ["APPOINTED_CUDA_DEVICE"] = str(None)
    llm_config.load_path = None  # NOTE: Set so that the Trainer does not load
    llm_config.save_folder = None  # NOTE: Set so that the Trainer does not save
    llm_config.device_train_microbatch_size = 1
    desired_steps = cfg.mt_dao.restore_cent.run_batches
    checkpoint_information = retrieve_checkpoint_states(
        checkpoint_dir=centralised_checkpoint_dir,
        s3_comm_config=cfg.s3_comm_config,
        run_uuid=cent_run_uuid,
        desired_steps=desired_steps,
    )
    (local_checkpoint_directory, tmp_dir, remote_or_local, checkpoint_suffixes) = (
        checkpoint_information
    )
    # Creating ClientConfig object
    client_config = CentralizedConfig(
        allow_unigram_metrics_failures=cfg.fl.allow_unigram_metrics_failures,
        resize_vocab=cfg.fl.resize_vocab,
        split_eval=cfg.centralized.split_eval,
        use_unigram_metrics=cfg.fl.use_unigram_metrics,
        s3_comm_config=cfg.s3_comm_config,
        cid=None,
    )
    # Get the temporary trainer object
    trainer, *_ = get_trainer_object(
        llm_config,
        client_config=client_config,
        no_data_loading=True,
    )
    param_names, n_experts = extract_layer_names_and_experts(trainer=trainer)
    model_states, layer_names_and_types = retrieve_model_states_and_types(
        local_checkpoint_directory=local_checkpoint_directory,
        checkpoint_suffixes=checkpoint_suffixes,
        preserve_optimizer_states=preserve_optimizer_states,
        desired_steps=desired_steps,
        param_names=param_names,
        n_experts=n_experts,
    )
    if remote_or_local and tmp_dir:
        log(DEBUG, "Cleaning up the temporary directory.")
        tmp_dir.cleanup()
    assert len(layer_names_and_types) > 0
    assert len(model_states) > 0
    assert len(layer_names_and_types) == len(model_states)
    return model_states, tuple(layer_names_and_types)


def _get_n_experts_from_base_config(
    cfg: BaseConfig,
) -> int | None:
    n_experts: int | None = None
    ffn_config: dict[str, Any] = {}
    # Get model config from the base config
    model_config = cfg.llm_config.get("model", {})
    if model_config:
        # Get ffn_config from the model config
        ffn_config = model_config.get("ffn_config", {})
    if ffn_config:
        # If ffn_config is not empty, get the number of experts
        n_experts = ffn_config.get("ff_n_experts", None)
    if n_experts is not None:
        return int(n_experts)
    return n_experts


def resolve_pretrained_model_path(
    cfg: BaseConfig,
    initial_model_states: NDArrays,
    layer_names_and_types: tuple[tuple[str, ModelStateNames], ...],
) -> NDArrays:
    """Resolve the path to the pretrained model and load its states.

    This function checks if a pretrained model path is specified in the configuration.
    If so, it downloads the model from the specified path, transforms it into the
    appropriate format, and returns the model states. If no path is specified,
    it returns the initial model states.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object containing the pretrained model path.
    initial_model_states : NDArrays
        The initial model states to return if no pretrained model path is specified.
    layer_names_and_types: tuple[tuple[str, ModelStateNames], ...],
        The layer names and corresponding types for the corresponding model.
    remote_up_down : RemoteUploaderDownloader, optional
        The remote uploader/downloader instance to use for downloading the model
        if the path is a remote S3 path. If None, it will not download from S3.

    Returns
    -------
    NDArrays
        The model states loaded from the pretrained model path or the initial model
        states if no path is specified.

    Raises
    ------
    ValueError
        If the pretrained model path is specified but the model states cannot be loaded.

    """
    # Return initial model states if no pretrained model path is specified
    if not cfg.pretrained_model_path:
        return initial_model_states
    log(
        DEBUG,
        "FL server is loading pretrained model from %s",
        cfg.pretrained_model_path,
    )
    layer_names = []
    for layer_name, _ in layer_names_and_types:
        if layer_name not in layer_names:
            layer_names.append(layer_name)
    checkpoint_information = retrieve_checkpoint_states(
        checkpoint_dir=cfg.pretrained_model_path,
        s3_comm_config=cfg.s3_comm_config,
        run_uuid=cfg.run_uuid,
        desired_steps=-1,
    )
    (local_checkpoint_directory, tmp_dir, remote_or_local, checkpoint_suffixes) = (
        checkpoint_information
    )
    n_experts = _get_n_experts_from_base_config(cfg)
    pretrained_model_states, _ = retrieve_model_states_and_types(
        local_checkpoint_directory=local_checkpoint_directory,
        checkpoint_suffixes=checkpoint_suffixes,
        preserve_optimizer_states=True,  # NOTE: We are preserving the optimizer.
        desired_steps=-1,  # NOTE: We choose the last checkpoint if of .pt type in a dir
        n_experts=n_experts,
        param_names=layer_names,
    )
    if remote_or_local and tmp_dir:
        log(DEBUG, "Cleaning up the temporary directory.")
        tmp_dir.cleanup()
    if len(pretrained_model_states) != len(initial_model_states):
        msg = (
            f"Expected {len(initial_model_states)} parameters, "
            f"but got {len(pretrained_model_states)} from the pretrained checkpoint"
            f" at {cfg.pretrained_model_path}."
        )
        raise ValueError(msg)
    shape_match = all(
        a.shape == b.shape
        for a, b in zip(pretrained_model_states, initial_model_states, strict=False)
    )
    if not shape_match:
        msg = "Not all the parameters in the expected and loaded model match."
        raise ValueError(msg)
    return pretrained_model_states


def initialize_federated_learning(
    cfg: BaseConfig,
    remote_up_down: RemoteUploaderDownloader | None,
    mt_dao_save_path: str,
) -> ServerState:
    """Initialize the state for a federated learning process.

    This function sets up the initial state for a new federated learning process:
    bookkeeping variables, client states, global model parameters, and optionally
    saving the initial checkpoint to an S3 Object Store if configured. It prepares the
    server for the federated learning process by setting the starting round, time
    offset, cumulative server steps, and initializing the momentum vector for
    optimization algorithms.

    Args:
    ----
    cfg : BaseConfig
        The configuration object containing settings for federated learning and system
        behavior.
    remote_up_down : RemoteUploaderDownloader | None
        An optional uploader/downloader object for interacting with remote storage,
        required if checkpointing or S3 communication is enabled.
    mt_dao_save_path : str
        The local path the server uses to store and load checkpoints.

    Returns:
    -------
    ServerState
        The initial server state object.

    """
    # Instantiate a PRNG
    rng = random.Random(cfg.seed)  # noqa: S311
    # Initialize the bookkeeping variables
    start_round: int = 0
    time_offset: float = 0.0
    server_steps_cumulative: int = 0
    history = WandbHistory(use_wandb=cfg.use_wandb)
    # Initialize client_state_dict
    client_states: dict[str | int, ClientState] = {
        cid: ClientState(0) for cid in range(cfg.fl.n_total_clients)
    }
    # Initialize the model parameters and layer names mapping
    initial_model_states, layer_names_and_types = get_random_model_states_from_config(
        cfg,
    )
    # Resolve the path to the pretrained model if specified
    initial_model_states = resolve_pretrained_model_path(
        cfg=cfg,
        initial_model_states=initial_model_states,
        layer_names_and_types=layer_names_and_types,
    )

    momentum_vector: NDArrays = []
    second_momentum_vector: NDArrays = []
    # NOTE: We should unify the state with the strategy
    # Only focusing on saving space right now
    match cfg.fl.strategy_name.lower():
        case StrategyName.FEDMOM | StrategyName.NESTEROV:
            momentum_vector = [np.zeros_like(x) for x in initial_model_states]
        case StrategyName.FEDADAM | StrategyName.FEDYOGI:
            momentum_vector = [np.zeros_like(x) for x in initial_model_states]
            second_momentum_vector = [np.zeros_like(x) for x in initial_model_states]

    return ServerState(
        model_states=initial_model_states,
        layer_names_and_types=layer_names_and_types,
        transmission_scheduler=dispatch_model_state_scheduler(cfg),
        aggregation_mask_scheduler=dispatch_model_state_scheduler(cfg),
        history=history,
        time_offset=time_offset,
        server_steps_cumulative=server_steps_cumulative,
        client_states=client_states,
        momentum_vector=momentum_vector,
        second_momentum_vector=second_momentum_vector,
        state_keys=(CURRENT_SERVER_STATE_FILENAME, *dispatch_strategy_state_keys(cfg)),
        sampled_clients=[],
        previously_sampled_clients=[],
        rng=rng,
        start_round=start_round,
        current_round=start_round,
        total_rounds=cfg.fl.n_rounds,
        perf_counter=time.perf_counter(),
        remote_up_down=remote_up_down,
        local_checkpoint_path=mt_dao_save_path,
        client_sampler=ClientSampler(
            total_number_of_clients=cfg.fl.n_total_clients,
            number_of_clients_per_round=cfg.fl.n_clients_per_round,
            dropout_ratio=cfg.fl.dropout_ratio,
            dropout_function_name=cfg.fl.dropout_function_name,
        ),
        max_shard_size=cfg.mt_dao.max_shard_size,
    )


def resume_from_round(
    cfg: BaseConfig,
    remote_up_down: RemoteUploaderDownloader | None,
    mt_dao_save_path: str,
) -> ServerState:
    """Resume from a previous round.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object.
    remote_up_down : RemoteUploaderDownloader
        The object to upload/download files from/to the S3 Object Store.
    mt_dao_save_path : str
        The local path the server uses to store and load checkpoints.

    Raises
    ------
    ValueError
        Cannot have restoration/resume features without specifying the
        initalisation location.

    Returns
    -------
    ServerState
        The server state object.

    """
    # Obtain server and strategy state keys
    state_keys = (CURRENT_SERVER_STATE_FILENAME, *dispatch_strategy_state_keys(cfg))
    # Interpret the resume round
    # NOTE: This is done such that we cater for both resuming and restoring
    # functionality. Particularly, we want the correct uuid to be set such that
    # when the interpret_resume_round function is called, we have set the correct
    # run to avoid duplicate work later.
    restore_or_resume_uuid = (
        cfg.run_uuid
        if cfg.mt_dao.restore_run_uuid is None
        else cfg.mt_dao.restore_run_uuid
    )
    if cfg.mt_dao.remote_init:
        run_uuid_path = (
            f"s3://{cfg.s3_comm_config.bucket_name}/{restore_or_resume_uuid}/"
        )
    elif cfg.mt_dao.local_init:
        # NOTE: by default, the mt_dao_save_path has the the run_uuid baked into it.
        #  As such, we have to replace it in the case we are restoring. If we are
        # resuming we will replace it with the exact same string.
        run_uuid_path = mt_dao_save_path.replace(cfg.run_uuid, restore_or_resume_uuid)
    else:
        msg = "Cannot resume from round without local/remote initialisation activated."
        raise ValueError(msg)
    cfg.mt_dao.resume_round = interpret_resume_round(
        resume_round=cfg.mt_dao.resume_round,
        run_uuid_path=run_uuid_path,
        # NOTE: This condition will ensure that if no checkpoints are found but the
        # `cfg.mt_dao.resume_round` equals -1, no error is raised and federated learning
        # will start from scratch. However, if the `cfg.mt_dao.resume_round` is lower
        # than -1, it will raise an error.
        raise_error=cfg.mt_dao.resume_round != -1,
        state_keys=state_keys,
    )
    # Import another experiment checkpoints for restoration
    if cfg.mt_dao.restore_run_uuid is not None:
        log(
            DEBUG,
            "Copy past checkpoints from run with UUID %s",
            cfg.mt_dao.restore_run_uuid,
        )
        copy_old_checkpoints_to_new_run(
            remote_up_down=remote_up_down,
            mt_dao_save_path=mt_dao_save_path,
            bucket_uri=f"s3://{cfg.s3_comm_config.bucket_name}",
            restore_run_data=(
                cfg.run_uuid,
                cfg.mt_dao.restore_run_uuid,
                cfg.mt_dao.resume_round,  # type: ignore[PGH003]
                cfg.mt_dao.resume_round * cfg.fl.n_local_steps,  # type: ignore[PGH003, operator]
            ),
            n_total_clients=cfg.fl.n_total_clients,
            copy_client_checkpoints=cfg.mt_dao.checkpointing.copy_client_checkpoints,
        )
    # NOTE: If the `interpret_resume_round` function returns
    # `cfg.mt_dao.resume_round == -1`, it means that it couldn't find the last
    # checkpoint. In fact, if it finds the last checkpoint for round R>0, it would have
    # replaced `cfg.mt_dao.resume_round` with R, thus `cfg.mt_dao.resume_round != -1`.
    assert cfg.mt_dao.resume_round is not None, (
        "Cannot resume run if `cfg.mt_dao.resume_round` is None"
    )
    if cfg.mt_dao.resume_round == -1:
        log(INFO, "There is no checkpoint to resume. Starting from scratch.")
        return initialize_federated_learning(
            cfg=cfg,
            remote_up_down=remote_up_down,
            mt_dao_save_path=mt_dao_save_path,
        )
    log(INFO, "A checkpoint was found. Resuming from it.")
    log(DEBUG, "Resume round %s", cfg.mt_dao.resume_round)
    return get_server_checkpoint(
        cfg=cfg,
        remote_up_down=remote_up_down,
        mt_dao_save_path=mt_dao_save_path,
    )
