"""Node Manager Application for Federated Learning.

This module defines the NodeManagerApp class, which is responsible for managing the
training and evaluation tasks on a node in a federated learning setup. The class
handles the creation and management of worker processes, the distribution of tasks
to these workers, and the aggregation of results. It also manages the configuration
and shared memory required for the tasks.

"""

import contextlib
import gc
import os
import pickle  # noqa: S403
import socket
import time
import uuid
from logging import DEBUG
from typing import TYPE_CHECKING, Any, cast

import cloudpickle
import torch
from composer.loggers import RemoteUploaderDownloader
from composer.utils.misc import get_free_tcp_port
from flwr.client import ClientApp
from flwr.client.typing import ClientFnExt, Mod
from flwr.common.logger import log
from flwr.common.typing import NDArrays, Scalar
from multiprocess import (
    Queue,  # type: ignore[reportAttributeAccessIssue]
    set_start_method,  # type: ignore[reportAttributeAccessIssue]
)
from omegaconf import DictConfig, OmegaConf

import ray
from mt_dao.constants import (
    CONFIG_FILE_YAML,
    EVAL,
    FIT,
    MASTER_PORT,
    mt_dao_SAVE_PATH,
    RUN_UUID,
)
from mt_dao.masks_utils import (
    generate_full_mask,
)
from mt_dao.node_manager.data_types import TrainEvalInstructions
from mt_dao.shm.constants import (
    NM_CONFIG_SHM,
    NM_PARAMETERS_SHM,
    W_EVAL_LOSS_SHM,
    W_METRICS_SHM,
    W_N_SAMPLES_SHM,
)
from mt_dao.shm.utils import (
    ModelParametersMetadata,
    close_all_shms,
    get_config_shm,
    get_dict_configrecord_shm,
    get_eval_loss_shm,
    get_num_samples_shm,
    get_parameters_shm,
    is_shm_existing,
    remove_shm_from_resource_tracker,
    set_dict_configrecord_shm,
)
from mt_dao.utils import (
    get_n_cuda_devices,
    get_random_model_states_from_config,
)
from mt_dao.worker.worker import (
    Worker,
    create_new_worker,
    get_training_results_from_worker,
    start_worker,
)

if TYPE_CHECKING:
    from multiprocessing.queues import Queue as QueueType
    from multiprocessing.shared_memory import SharedMemory

    from mt_dao.conf.base_schema import BaseConfig
    from mt_dao.worker.utils import WorkerResultMessage

set_start_method("spawn", force=True)
pickle.Pickler = cloudpickle.Pickler  # type: ignore[misc]


class NodeManagerApp(ClientApp):
    """NodeManagerApp is responsible for managing the training and evaluation on a node.

    This class handles the creation and management of worker processes, the distribution
    of tasks to these workers, and the aggregation of results. It also manages the
    configuration and shared memory required for the tasks.

    Attributes
    ----------
        cfg : BaseConfig
            The configuration object loaded from the config file.
        task_queue : QueueType
            The queue for tasks assigned to the workers.
        result_queue : QueueType
            The queue for results produced by the workers.
        node : NodeProperties
            The properties of the node, including hardware accelerators.
        properties : dict[str, ConfigRecordValues]
            A dictionary containing the node properties.
        node_manager_uuid : str
            The unique identifier for the node manager.
        refresh_period : int | None
            The refresh period for the node manager.
        remote_up_down : RemoteUploaderDownloader | None
            The remote uploader/downloader for handling S3 communication.
        parameters_metadata : ModelParametersMetadata
            Metadata for the model parameters.
        round_parameters : NDArrays | None
            The model parameters for the current round.
        round_parameters_sh : SharedMemory | None
            Shared memory object for the round parameters.
        workers_dict : dict[int, Worker]
            A dictionary of worker processes.
        fl_instructions_config : TypedDict[str, ConfigRecord] | None
            The configuration for the federated learning instructions.
        fl_instructions_config_sh : SharedMemory | None
            Shared memory object for the federated learning instructions.

    """

    def __init__(
        self,
        client_fn: ClientFnExt | None = None,  # Only for backward compatibility
        mods: list[Mod] | None = None,
    ) -> None:
        """Initialize the NodeManagerApp instance.

        This method sets up the NodeManagerApp by loading the configuration, setting up
        queues, getting node properties, creating workers, and initializing shared
        memories.

        Parameters
        ----------
        client_fn : ClientFnExt | None, optional
            A client function for backward compatibility. Default is None.
        mods : list[Mod] | None, optional
            A list of modifications to apply. Default is None.

        Raises
        ------
        ValueError
            If the environmental variable mt_dao_SAVE_PATH is not set.
        TypeError
            If the LLM configuration is not a DictConfig.

        """
        super().__init__(client_fn=client_fn, mods=mods)
        # Get the environmental variable for the dump folder
        mt_dao_save_path = os.environ.get(mt_dao_SAVE_PATH, "")
        # Raise an error if the environmental variable is not set
        if not mt_dao_save_path:
            msg = "The environmental variable mt_dao_SAVE_PATH is not set."
            raise ValueError(msg)
        # Load the configuration from the config file
        self.cfg = cast(
            "BaseConfig",
            OmegaConf.load(mt_dao_save_path + f"/{CONFIG_FILE_YAML}"),
        )
        # Resolve the config and set it to be editable in place
        OmegaConf.resolve(self.cfg)
        OmegaConf.set_struct(self.cfg, value=False)
        # Set up Queues
        self.task_queue: QueueType = Queue()
        # One result_queue for all GPUs
        self.result_queue: QueueType[WorkerResultMessage] = Queue()
        self.node_manager_uuid = self.cfg.run_uuid + "-" + str(uuid.uuid4())
        self.refresh_period = self.cfg.mt_dao.refresh_period
        self.remote_up_down: RemoteUploaderDownloader | None = None
        self.create_remote_up_down()
        # Call the monkey-patch for the resource-register
        remove_shm_from_resource_tracker()
        # Extract the LLM part of the config
        llm_config = self.cfg.llm_config
        if not isinstance(llm_config, DictConfig):
            msg = "The LLM configuration is not a DictConfig."
            raise TypeError(msg)
        # Get initial model parameters
        parameters, self.layer_names_and_types = get_random_model_states_from_config(
            self.cfg,
        )
        # TODO(anon): Make this useful
        # parameters, self.layer_names_and_types = reconcile_model_state_with_scheduler(
        #     parameters,
        #     self.layer_names_and_types,
        #     self.cfg)
        parameters_metadata = ModelParametersMetadata.from_ndarrays(parameters)
        del parameters
        # Parameter metadata
        self.parameters_metadata = parameters_metadata
        # Shared memory for round parameters
        self.round_parameters: NDArrays | None = None
        self.round_parameters_sh: SharedMemory | None = None
        # Used to store the transmission mask
        # which workers need for the shared memories
        # and clients need for the parameters
        self.transmission_mask: tuple[tuple[bool, ...], list[str], list[str]] = (
            generate_full_mask(
                layer_names_and_types=self.layer_names_and_types,
            )
        )
        self.aggregation_mask: tuple[tuple[bool, ...], list[str], list[str]] = (
            generate_full_mask(
                layer_names_and_types=self.layer_names_and_types,
            )
        )
        self.create_parameters_shm()
        # Create workers
        self.workers_dict: dict[int, Worker] = {}
        # Master port for PyTorch Distributed
        self.nm_master_port: int | None = None
        # Initialize Ray
        if self.cfg.mt_dao.comm_stack.ray:
            ray.init(
                "auto",
            )
        # Training results
        self.worker_parameters: NDArrays = []
        self.worker_metrics: dict[str, Scalar] = {}
        self.worker_num_samples: int = 0
        self.worker_loss: float = 0.0
        self._worker_shms: tuple[SharedMemory, ...] = ()

    def create_parameters_shm(self) -> None:
        """Create shared memory for round parameters.

        This method creates or retrieves shared memory for storing round parameters
        using the parameters metadata. It assigns the resulting parameters and shared
        memory object to instance variables.
        """
        self.round_parameters, self.round_parameters_sh = get_parameters_shm(
            parameters_metadata=self.parameters_metadata,
            create=not is_shm_existing(self.node_manager_uuid + NM_PARAMETERS_SHM),
            name=self.node_manager_uuid + NM_PARAMETERS_SHM,
        )

    def create_and_start_workers(self) -> None:
        """Create and start worker processes.

        This method creates worker processes based on the number of CUDA devices or CPU
        concurrency and starts them.
        """
        workers_to_start: list[int] = []
        for i in range(
            (get_n_cuda_devices()),
        ):
            # Get the worker if already in the dictionary
            if i in self.workers_dict:
                log(
                    DEBUG,
                    "NodeManagerApp %s on node %s: worker %s already exists."
                    " Checking health...",
                    self.node_manager_uuid,
                    socket.gethostname(),
                    i,
                )
                if self.workers_dict[i].is_alive():
                    log(
                        DEBUG,
                        "NodeManagerApp %s on node %s: worker %s is alive.",
                        self.node_manager_uuid,
                        socket.gethostname(),
                        i,
                    )
                    continue
            worker = create_new_worker(
                config=self.cfg,
                task_queue=self.task_queue,
                result_queue=self.result_queue,
                node_manager_uuid=self.node_manager_uuid,
                run_uuid=self.cfg.run_uuid,
                parameters_metadata=self.parameters_metadata,
                worker_rank=i,
            )
            self.workers_dict[i] = worker
            log(
                DEBUG,
                "Created worker with rank %s",
                i,
            )
            workers_to_start.append(i)
        # Start the workers
        for i in workers_to_start:
            start_worker(self.workers_dict[i])
        log(
            DEBUG,
            "NodeManagerApp %s on node %s: all workers started.",
            self.node_manager_uuid,
            socket.gethostname(),
        )

    def check_workers_health(self) -> None:
        """Check the health of worker processes.

        This method checks if the worker processes are alive. If a worker is found to be
        dead, it restarts the worker.
        """
        for rank, worker in self.workers_dict.items():
            if not worker.is_alive():
                log(
                    DEBUG,
                    "NodeManagerApp %s on node %s: worker %s is dead. Restarting it...",
                    self.node_manager_uuid,
                    socket.gethostname(),
                    rank,
                )
                close_all_shms(worker.worker_uuid)
                self.workers_dict[rank] = create_new_worker(
                    config=self.cfg,
                    task_queue=self.task_queue,
                    result_queue=self.result_queue,
                    node_manager_uuid=self.node_manager_uuid,
                    run_uuid=self.cfg.run_uuid,
                    parameters_metadata=self.parameters_metadata,
                    worker_rank=rank,
                )
                start_worker(self.workers_dict[rank])

    def close_workers(self) -> None:
        """Perform a soft shutdown of all worker processes.

        This method initiates a soft shutdown for each worker process, waits until they
        are terminated, and then performs garbage collection and clears the CUDA cache.
        """
        # Wait until the worker is dead
        for worker in self.workers_dict.values():
            worker.soft_shutdown()
            while worker.is_alive():
                time.sleep(0.1)
                worker.terminate()
        log(
            DEBUG,
            "NodeManagerApp %s on node %s: workers are dead.",
            self.node_manager_uuid,
            socket.gethostname(),
        )
        gc.collect()
        torch.cuda.empty_cache()

    def create_remote_up_down(self) -> None:
        """Create and initialize the remote uploader/downloader for S3 communication.

        This method sets up the remote uploader/downloader using the S3 configuration
        provided in the configuration file.
        """
        if self.cfg.mt_dao.comm_stack.s3:
            bucket_uri = f"s3://{self.cfg.s3_comm_config.bucket_name}"
            self.remote_up_down = RemoteUploaderDownloader(
                bucket_uri=bucket_uri,
                backend_kwargs={
                    "bucket": self.cfg.s3_comm_config.bucket_name,
                    "prefix": f"{self.cfg.run_uuid}/server",  # Don't touch
                    "region_name": None,  # Not necessary
                    "endpoint_url": None,  # Will be read from env var
                    "aws_access_key_id": None,  # Will be read from config file
                    "aws_secret_access_key": None,  # Will be read from config file
                    "aws_session_token": None,  # Will be automatically generated
                    "client_config": OmegaConf.to_container(
                        self.cfg.s3_comm_config.backend_kwargs.client_config,
                    ),  # And using defaults
                    "transfer_config": None,  # Using defaults
                },
                file_path_format_string="{remote_file_name}",  # Don't touch
                num_concurrent_uploads=1,
                upload_staging_folder=None,  # Don't touch, it's /tmp by default
                use_procs=True,  # Don't touch
                num_attempts=self.cfg.s3_comm_config.num_attempts,
            )
            self.remote_up_down.init(run_name=self.cfg.run_uuid)

    def fit(
        self,
        fit_ins: TrainEvalInstructions,
    ) -> tuple[NDArrays, int, dict[str, Scalar]]:
        """Fit the model using the provided training instructions.

        This method prepares the training instructions, updates the shared memory for
        the federated learning instructions, and sends the fit task to the workers. It
        waits for the result from the workers, checks if the training was successful,
        and retrieves the training results from the worker's shared memory. It also
        updates the worker metrics with the training time and returns the results.

        Parameters
        ----------
        fit_ins : TrainEvalInstructions
            The training instructions containing the model parameters and other
            configuration needed for training.

        Returns
        -------
        tuple[NDArrays, int, dict[str, Scalar]]
            A tuple containing the model parameters, number of samples, and training
            metrics.

        Raises
        ------
        RuntimeError
            If the training was not successful, i.e., if the number of samples is less
            than or equal to -1. This indicates that the workers failed to train the
            model.

        """
        start_time = time.perf_counter()
        # Here, workers are forced to collaborate with each other,
        # as such, we evaluate one client at a time
        self.check_workers_health()
        # Append NodeManager's config
        self.nm_master_port = self.nm_master_port or get_free_tcp_port()
        fit_ins.ins[MASTER_PORT] = self.nm_master_port
        # NOTE: Putting the node_manager_uuid in the config fails
        fit_ins.ins[RUN_UUID] = self.cfg.run_uuid
        # Update instruction config shared memory
        _fl_instructions_config, fl_instructions_config_sh = get_dict_configrecord_shm(
            config=fit_ins.ins,
            create=True,
            name=self.node_manager_uuid + NM_CONFIG_SHM,
        )
        set_dict_configrecord_shm(fit_ins.ins, fl_instructions_config_sh)
        # Send the collaborative task to the workers
        for _ in range(len(self.workers_dict)):
            self.task_queue.put((fit_ins.client_id, FIT))
        # Wait for the result
        worker_result = self.result_queue.get()
        # Check if the training was successful
        if worker_result.n_samples <= -1:
            # Close all workers to refresh the state
            self.close_workers()
            # TODO(anon): Decide what to do if the training was not successful
            msg = (
                f"NodeManager {self.node_manager_uuid} on node {socket.gethostname()}: "
                f"Workers failed to train the model. Restarting workers and retrying."
            )
            raise RuntimeError(msg)
        # Close the config shared memory so as not to waste resources
        with contextlib.suppress(FileNotFoundError):
            fl_instructions_config_sh.close()
            fl_instructions_config_sh.unlink()
        # Empty the tasks list in case anything has been left there
        while not self.task_queue.empty():
            self.task_queue.get()

        # Retrieve aggregation mask to read the shared memories
        self.aggregation_mask = worker_result.aggregation_mask
        # Get stuff from shared memories of the workers
        # NOTE: Keep a reference to the `*_shm` variables to prevent Seg Fault
        # NOTE: Only takes the metrics from one worker,
        (
            self.worker_parameters,
            self.worker_metrics,
            self.worker_num_samples,
            self._worker_shms,
        ) = get_training_results_from_worker(
            self.workers_dict[0],
        )
        # Adding node training time in the metrics
        self.worker_metrics.update(
            {
                "node_training_time_s": float(time.perf_counter() - start_time),
                "worker_training_time_s": worker_result.delta,
            },
        )
        # Return the results from the worker
        return (
            self.worker_parameters,
            self.worker_num_samples,
            self.worker_metrics,
        )

    def eval(
        self,
        eval_ins: TrainEvalInstructions,
    ) -> tuple[float, int, dict[Any, Any]]:
        """Evaluate the model using the provided evaluation instructions.

        This method prepares the evaluation instructions, updates the shared memory for
        the federated learning instructions, and sends the evaluation task to the
        workers. It waits for the result from the workers, checks if the evaluation was
        successful, and retrieves the evaluation results from the worker's shared
        memory. It also updates the worker metrics with the evaluation time and
        returns the results.

        Parameters
        ----------
        eval_ins : TrainEvalInstructions
            The evaluation instructions to be sent to the workers.

        Returns
        -------
        tuple[float, int, dict[Any, Any]]
            The evaluation loss, number of samples, and worker metrics.

        Raises
        ------
        RuntimeError
            If the evaluation was not successful, i.e., if the number of samples is less
            than or equal to -1. This indicates that the workers failed to evaluate the
            model.

        """
        start_time = time.perf_counter()
        # Here, workers are forced to collaborate with each other,
        # as such, we evaluate one client at a time
        self.check_workers_health()
        # Append NodeManager's config
        self.nm_master_port = self.nm_master_port or get_free_tcp_port()
        eval_ins.ins[MASTER_PORT] = self.nm_master_port
        # NOTE: Putting the node_manager_uuid in the config fails
        eval_ins.ins[RUN_UUID] = self.cfg.run_uuid
        # Update instruction config shared memory
        _fl_instructions_config, fl_instructions_config_sh = get_dict_configrecord_shm(
            config=eval_ins.ins,
            create=True,
            name=self.node_manager_uuid + NM_CONFIG_SHM,
        )
        set_dict_configrecord_shm(eval_ins.ins, fl_instructions_config_sh)
        # Send the collaborative task to the workers
        for _ in range(len(self.workers_dict)):
            self.task_queue.put((eval_ins.client_id, EVAL))
        # Wait for the result
        worker_result = self.result_queue.get()
        # Check if the evaluation was successful
        if worker_result.n_samples <= -1:
            # Close all workers to refresh the state
            self.close_workers()
            # TODO(anon): Decide what to do if the evaluation was not successful
            msg = (
                f"NodeManager {self.node_manager_uuid} on node {socket.gethostname()}: "
                "Workers failed to evaluate the model. Restarting workers and retrying."
            )
            raise RuntimeError(msg)
        # Get stuff from shared memories of the rank 0 worker
        # NOTE: Keep the `*_shm` variables to prevent Seg Fault
        self.worker_loss, w_eval_loss_shm = get_eval_loss_shm(
            name=self.workers_dict[0].worker_uuid + W_EVAL_LOSS_SHM,
        )
        self.worker_num_samples, w_num_samples_shm = get_num_samples_shm(
            name=self.workers_dict[0].worker_uuid + W_N_SAMPLES_SHM,
        )
        self.worker_metrics, w_metrics_shm = get_config_shm(
            config={},
            name=self.workers_dict[0].worker_uuid + W_METRICS_SHM,
        )
        self._worker_shms = (w_eval_loss_shm, w_num_samples_shm, w_metrics_shm)
        with contextlib.suppress(FileNotFoundError):
            # Close the config shared memory
            fl_instructions_config_sh.close()
            fl_instructions_config_sh.unlink()
        # Empty the tasks list
        while not self.task_queue.empty():
            self.task_queue.get()

        self.worker_metrics.update(
            {
                "node_eval_time_s": float(time.perf_counter() - start_time),
                "worker_eval_time_s": worker_result.delta,
            },
        )
        # Aggregation of eval samples
        log(
            DEBUG,
            "NodeManager %s: results have been processed. "
            "The time spent before collecting results was %s seconds.",
            self.node_manager_uuid,
            time.perf_counter() - start_time,
        )
        # Return results
        return (
            self.worker_loss,
            self.worker_num_samples,
            self.worker_metrics,
        )
