"""Communication utilities for efficient model parameter transfer in federated training.

This module implements optimized communication methods for transferring large model
parameters between distributed components using different backends. It provides
functionality to offload parameters from recorddicts to remote storage and load them
back as needed, supporting three communication methods:

1. S3 object storage: Parameters are serialized, stored in S3, and replaced with
   references in the recorddict.
2. Shared Memory (shm): Parameters are stored in shared memory accessible by
   multiple processes on the same machine.
3. Ray object store: Parameters are stored in Ray's distributed object store
   with optimized conversions.

The module optimizes large parameter transfers by:
- Processing arrays in batches to balance parallelism with overhead
- Converting between precision formats (float16/float32) as appropriate
- Optimizing serialization operations
- Using vectorized operations where possible

These optimizations help reduce memory usage, network transfer times, and overall
communication overhead in distributed federated learning workflows.

Functions:
    offload_recorddict_parameters_to_remote: Move parameters from recorddict to remote
        storage
    load_recorddict_parameters_from_remote: Retrieve parameters from remote storage into
        recorddict
"""

import ast
import os
import time
from copy import deepcopy
from logging import DEBUG, ERROR
from multiprocessing.shared_memory import SharedMemory
from pathlib import Path
from queue import Queue
from tempfile import TemporaryDirectory

import numpy as np
from composer.loggers import RemoteUploaderDownloader
from composer.utils.file_helpers import (
    validate_given_remote_path,
)
from flwr.common import (
    Code,
    ConfigRecord,
    Message,
    NDArrays,
    Parameters,
    RecordDict,
    log,
    ndarrays_to_parameters,
    parameters_to_ndarrays,
)
from flwr.common.recorddict_compat import (
    _extract_status_from_recorddict,  # noqa: PLC2701
    arrayrecord_to_parameters,
    parameters_to_arrayrecord,
)

import ray
from mt_dao.conf.base_schema import CommStack
from mt_dao.constants import (
    BIN_SUFFIX,
    COMM_STACK,
    EMPTY,
    ENDPOINT_ID,
    NUMPY_ARRAY_SUFFIX,
    PARAMETERS,
)
from mt_dao.file_utils import (
    download_file_from_s3,
    dump_model_parameters_to_file,
    load_model_states_from_file,
    upload_file_to_s3,
)
from mt_dao.server.comms_ray_batch_variants import (
    RayOffloadStrategyName,
    RayOffloadStrategyNotImplementedError,
    load_from_ray_parameters_with_strategy,
    offload_to_ray,
    offload_to_ray_legacy,
)
from mt_dao.server.s3_helper import (
    extract_s3_comm_config_from_configrecord,
)
from mt_dao.shm.utils import (
    ModelParametersMetadata,
    get_parameters_shm,
    is_shm_existing,
)
from mt_dao.utils import (
    set_parameters_in_place,
)

mt_dao_RAY_OFFLOAD_STRATEGY = str(
    os.getenv("mt_dao_RAY_OFFLOAD_STRATEGY", RayOffloadStrategyName.STANDARD),
)  # Configurable via environment variable


def offload_to_s3(
    parameters: NDArrays,
    remote_uploader_downloader: RemoteUploaderDownloader,
    recorddict: RecordDict,
    msg_str: str,
) -> RecordDict:
    """Offload model parameters to S3 object storage.

    This function serializes model parameters and uploads them to S3 object storage,
    replacing the original parameters in the recorddict with empty placeholders. It
    creates a temporary local file, uploads it to S3, and then verifies the upload was
    successful before returning. This approach reduces memory usage during transmission
    and enables cross-node parameter sharing through persistent cloud storage.

    Parameters
    ----------
    parameters : NDArrays
        The model parameters to offload to S3 storage.
    remote_uploader_downloader : RemoteUploaderDownloader
        The uploader/downloader instance for interacting with S3.
    recorddict : RecordDict
        The recorddict containing the message configuration and parameters.
    msg_str : str
        A string identifier used to prefix keys in the recorddict.

    Returns
    -------
    RecordDict
        The modified recorddict with emptied parameters and S3 reference information.

    Notes
    -----
    This function optimizes data offloading by:
    1. Using temporary local files to minimize memory overhead during serialization
    2. Extracting endpoint and file info from the recorddict for consistent naming
    3. Employing polling to verify successful upload completion
    4. Emptying the recorddict parameters after upload to reduce memory usage

    """
    # Create a temporary directory for storing the downloaded parameters
    with TemporaryDirectory() as temp_dir:
        # Extract endpoint id from the content of the message
        endpoint_id, file_name, folder_name = extract_s3_comm_config_from_configrecord(
            recorddict.config_records[f"{msg_str}.{COMM_STACK}"],
        )
        # Set the file names
        local_file_name = (
            Path(temp_dir) / f"{endpoint_id}_{file_name}{NUMPY_ARRAY_SUFFIX}"
        )
        dump_model_parameters_to_file(
            local_file_name,
            parameters,
        )
        # Upload the parameters to S3 Object Store
        upload_file_to_s3(
            remote_up_down=remote_uploader_downloader,
            remote_file_name=f"{folder_name}/{endpoint_id}/{file_name}{NUMPY_ARRAY_SUFFIX}",
            local_file_name=local_file_name,
        )
        # Empty the recorddict parameters
        recorddict.parameters_records[f"{msg_str}.{PARAMETERS}"] = (
            parameters_to_arrayrecord(
                Parameters(tensors=[], tensor_type=EMPTY),
                keep_input=False,
            )
        )

        # Check whether the server has uploaded the parameters
        remote_file_name_no_ext = (
            f"s3://{remote_uploader_downloader.remote_bucket_name}/"
            f"{remote_uploader_downloader.backend_kwargs['prefix']}/"
            f"{folder_name}/{endpoint_id}/{file_name}"
        )
        while not (
            validate_given_remote_path(
                remote_file_name_no_ext + BIN_SUFFIX,
            )
            or validate_given_remote_path(remote_file_name_no_ext + NUMPY_ARRAY_SUFFIX)
        ):
            time.sleep(0.5)
        log(
            DEBUG,
            "Node %s parameters have been pushed to the S3",
            endpoint_id,
        )
    return recorddict


def offload_to_shm(
    parameters: NDArrays,
    recorddict: RecordDict,
    msg_str: str,
) -> tuple[RecordDict, list[str]]:
    """Offload model parameters to shared memory.

    This function stores model parameters in shared memory that can be accessed by
    multiple processes on the same machine. It creates or reuses shared memory segments,
    writes the parameters to them, and replaces the recorddict's parameters with
    metadata references to the shared memory. This approach optimizes memory usage by
    allowing different processes to access the same physical memory.

    Parameters
    ----------
    parameters : NDArrays
        The model parameters to offload to shared memory.
    recorddict : RecordDict
        The recorddict containing the message configuration and parameters.
    msg_str : str
        A string identifier used to prefix keys in the recorddict.

    Returns
    -------
    RecordDict
        The modified recorddict with emptied parameters and references to the shared
        memory.

    Notes
    -----
    This function optimizes data offloading by:
    1. Storing parameters in shared memory to avoid duplicate memory allocation
    2. Handling cases where existing shared memory is too small by recreating it
    3. Emptying the recorddict parameters to reduce memory usage after transfer
    4. Storing essential metadata in the recorddict for later reconstruction

    """
    # Get parameters metadata
    parameters_metadata = ModelParametersMetadata.from_ndarrays(parameters)
    # Create the parameters shared memory
    shm_name = str(
        recorddict.config_records[f"{msg_str}.{COMM_STACK}"][ENDPOINT_ID],
    )
    try:
        shm_parameters, _shm_parameters_sh = get_parameters_shm(
            parameters_metadata=parameters_metadata,
            create=not is_shm_existing(shm_name),
            name=shm_name,
        )
    except TypeError:
        log(DEBUG, "offload_to_shm: Shared memory too small, replacing it.")
        try:
            shm = SharedMemory(name=shm_name)
            shm.close()
            shm.unlink()
        except Exception as e:  # noqa: BLE001
            if "[Errno 2] No such file or directory" not in str(e):
                log(
                    ERROR,
                    "Removing Shared Memory %s failed because of %s",
                    shm_name,
                    e,
                )
        shm_parameters, _shm_parameters_sh = get_parameters_shm(
            parameters_metadata=parameters_metadata,
            create=True,
            name=shm_name,
        )
    # Set the parameters in the shared memory
    set_parameters_in_place(
        old_parameters=shm_parameters,
        new_parameters=parameters,
        parameter_pos=range(len(parameters)),
    )
    # Empty the recorddict parameters
    recorddict.parameters_records[f"{msg_str}.{PARAMETERS}"] = (
        parameters_to_arrayrecord(
            Parameters(tensors=[], tensor_type=EMPTY),
            keep_input=False,
        )
    )
    # Serialize ModelParametersMetadata and set it in the recorddict
    parameters_metadata_dict = parameters_metadata.__dict__
    parameters_metadata_dict["dtypes"] = [
        str(v) for v in parameters_metadata_dict["dtypes"]
    ]
    recorddict.config_records[f"{msg_str}.parameters_metadata"] = ConfigRecord(
        {"metadata": str(parameters_metadata_dict)},
    )
    return recorddict, [shm_name]


def offload_recorddict_parameters_to_remote(
    parameters: NDArrays,
    remote_uploader_downloader: RemoteUploaderDownloader | None,
    outgoing_recorddict: RecordDict,
    msg_str: str,
    comm_stack: CommStack,
) -> tuple[RecordDict, list[ray.ObjectRef] | list[str]]:
    """Offload model parameters to remote storage based on communication configuration.

    This function serves as a dispatcher that moves model parameters from a recorddict
    to the appropriate remote storage based on the communication stack configuration. It
    supports offloading to S3 object storage, shared memory, or Ray's distributed object
    store. The function delegates to specialized handlers for each storage type and
    returns the modified recorddict with appropriate parameter references replacing the
    original large parameter data.

    Parameters
    ----------
    parameters : NDArrays
        The model parameters to offload to remote storage.
    remote_uploader_downloader : RemoteUploaderDownloader | None
        The uploader/downloader instance for interacting with S3, or None if S3 is not
        used.
    outgoing_recorddict : RecordDict
        The recorddict containing the message configuration and parameters to be
        modified.
    msg_str : str
        A string identifier used to prefix keys in the recorddict.
    comm_stack : CommStack
        Configuration specifying which communication method(s) to use.

    Returns
    -------
    tuple[RecordDict, list[ray.ObjectRef] | list[str]]
        A tuple containing the modified recorddict with emptied parameters and
        references to the remote storage, and a list of Ray object references if Ray
        communication is used (or None otherwise).

    Raises
    ------
    RayOffloadStrategyNotImplementedError
        If the specified offload strategy is not implemented.

    Notes
    -----
    This function optimizes parameter offloading by:
    1. Selecting the most appropriate storage backend based on configuration
    2. Using specialized handlers optimized for each storage type
    3. Emptying the recorddict parameters after offloading to reduce memory usage
    4. Preserving Ray object references for later garbage collection when needed

    """
    # Initialize params_garbage_list
    params_garbage_list: list[ray.ObjectRef] | list[str] = []

    # Dispatch to the appropriate handler based on communication protocol
    if comm_stack.s3 and remote_uploader_downloader is not None:
        outgoing_recorddict = offload_to_s3(
            parameters,
            remote_uploader_downloader,
            outgoing_recorddict,
            msg_str,
        )
    elif comm_stack.shm:
        outgoing_recorddict, params_garbage_list = offload_to_shm(
            parameters,
            outgoing_recorddict,
            msg_str,
        )
    elif comm_stack.ray:
        match mt_dao_RAY_OFFLOAD_STRATEGY:
            case RayOffloadStrategyName.STANDARD:
                outgoing_recorddict, params_garbage_list = offload_to_ray_legacy(
                    parameters,
                    outgoing_recorddict,
                    msg_str,
                )
            case RayOffloadStrategyName.PACKED:
                outgoing_recorddict, params_garbage_list = offload_to_ray(
                    parameters,
                    outgoing_recorddict,
                    msg_str,
                    is_packed=True,
                )
            case RayOffloadStrategyName.BATCHED_SIMPLE:
                outgoing_recorddict, params_garbage_list = offload_to_ray(
                    parameters,
                    outgoing_recorddict,
                    msg_str,
                )
            case _:
                raise RayOffloadStrategyNotImplementedError(mt_dao_RAY_OFFLOAD_STRATEGY)
    else:
        outgoing_recorddict.parameters_records[f"{msg_str}.{PARAMETERS}"] = (
            parameters_to_arrayrecord(
                ndarrays_to_parameters(parameters),
                keep_input=True,
            )
        )

    # Update the message content with the modified recorddict
    return outgoing_recorddict, params_garbage_list


def load_from_s3_parameters(
    remote_uploader_downloader: RemoteUploaderDownloader,
    recorddict: RecordDict,
    msg_str: str,
) -> tuple[RecordDict, NDArrays]:
    """Load model parameters from S3 object storage.

    This function retrieves model parameters that were previously offloaded to S3
    storage and loads them back into memory. It extracts reference information from the
    recorddict, verifies the parameters are available in S3, downloads them to a
    temporary local file, and then loads them into memory. This approach enables
    efficient cross-node parameter sharing through persistent cloud storage.

    Parameters
    ----------
    remote_uploader_downloader : RemoteUploaderDownloader
        The uploader/downloader instance for interacting with S3.
    recorddict : RecordDict
        The recorddict containing references to the parameters in S3 storage.
    msg_str : str
        A string identifier used to prefix keys in the recorddict.

    Returns
    -------
    tuple[RecordDict, NDArrays]
        A tuple containing the original recorddict and the retrieved model parameters.

    Notes
    -----
    This function optimizes parameter loading by:
    1. Using temporary local files to minimize memory overhead during deserialization
    2. Polling S3 to ensure parameters are available before attempting download
    3. Supporting multiple file formats (.bin and .npz) for flexibility
    4. Choosing the appropriate file format based on what's available in S3

    """
    # Create a temporary directory for storing the downloaded parameters
    with TemporaryDirectory() as temp_dir:
        # Extract endpoint id from the content of the message
        endpoint_id, file_name, folder_name = extract_s3_comm_config_from_configrecord(
            recorddict.config_records[f"{msg_str}.{COMM_STACK}"],
        )
        # Check whether the server has uploaded the parameters
        remote_file_name_no_ext = (
            f"s3://{remote_uploader_downloader.remote_bucket_name}/"
            f"{remote_uploader_downloader.backend_kwargs['prefix']}/"
            f"{folder_name}/{endpoint_id}/{file_name}"
        )
        while not (
            validate_given_remote_path(
                remote_file_name_no_ext + BIN_SUFFIX,
            )
            or validate_given_remote_path(remote_file_name_no_ext + NUMPY_ARRAY_SUFFIX)
        ):
            time.sleep(0.5)
        # Set the file names depending on the extension found
        local_file_name = (
            Path(temp_dir) / f"tmp-{endpoint_id}{BIN_SUFFIX}"
            if validate_given_remote_path(remote_file_name_no_ext + BIN_SUFFIX)
            else Path(temp_dir) / f"tmp-{endpoint_id}{NUMPY_ARRAY_SUFFIX}"
        )
        download_file_from_s3(
            remote_up_down=remote_uploader_downloader,
            remote_file_name=(
                f"{folder_name}/{endpoint_id}/{file_name}{BIN_SUFFIX}"
                if validate_given_remote_path(remote_file_name_no_ext + BIN_SUFFIX)
                else f"{folder_name}/{endpoint_id}/{file_name}{NUMPY_ARRAY_SUFFIX}"
            ),
            local_file_name=local_file_name,
        )
        log(
            DEBUG,
            "Node %s parameters have been read from disk and assigned to the Message",
            endpoint_id,
        )
        return recorddict, load_model_states_from_file(local_file_name)[1]


def load_from_shm_parameters(
    recorddict: RecordDict,
    msg_str: str,
    params_garbage_queue: Queue[ray.ObjectRef | str] | None = None,
) -> tuple[RecordDict, NDArrays]:
    """Load model parameters from shared memory.

    This function retrieves model parameters that were previously stored in shared
    memory segments and loads them back into memory. It extracts metadata from the
    recorddict, accesses the shared memory segment, and creates numpy arrays backed by
    the shared memory buffers. For safety, it makes a deep copy of the parameters to
    avoid potential memory corruption when shared memory references are lost. It also
    handles error cases such as shared memory being too small.

    Parameters
    ----------
    recorddict : RecordDict
        The recorddict containing references to parameters in shared memory and their
        metadata.
    msg_str : str
        A string identifier used to prefix keys in the recorddict.
    params_garbage_queue : Queue[ray.ObjectRef | str] | None, optional
        Queue for Ray or Shm garbage collection.

    Returns
    -------
    tuple[RecordDict, NDArrays]
        A tuple containing the original recorddict and the retrieved model parameters.

    Notes
    -----
    This function optimizes parameter loading by:
    1. Using shared memory to minimize copying between processes on the same machine
    2. Reconstructing parameters using metadata stored in the recorddict
    3. Handling cases where shared memory needs to be recreated with the right size
    4. Making a deep copy to avoid issues with dangling references to shared buffers

    """
    # Get parameters metadata
    parameters_metadata_dict = ast.literal_eval(
        str(
            recorddict.config_records[f"{msg_str}.parameters_metadata"]["metadata"],
        ),
    )
    parameters_metadata_dict["dtypes"] = [
        np.dtype(v) for v in parameters_metadata_dict["dtypes"]
    ]
    # Create the parameters shared memory
    shm_name = str(
        recorddict.config_records[f"{msg_str}.{COMM_STACK}"][ENDPOINT_ID],
    )
    try:
        shm_parameters, _shm_parameters_sh = get_parameters_shm(
            parameters_metadata=ModelParametersMetadata(
                **parameters_metadata_dict,
            ),
            create=False,
            name=shm_name,
        )
    except TypeError:
        log(DEBUG, "load_from_shm_parameters: Shared memory too small, replacing it.")
        try:
            shm = SharedMemory(name=shm_name)
            shm.close()
            shm.unlink()
        except Exception as e:  # noqa: BLE001
            if "[Errno 2] No such file or directory" not in str(e):
                log(
                    ERROR,
                    "Removing Shared Memory %s failed because of %s",
                    shm_name,
                    e,
                )
        shm_parameters, _shm_parameters_sh = get_parameters_shm(
            parameters_metadata=ModelParametersMetadata(
                **parameters_metadata_dict,
            ),
            create=True,
            name=shm_name,
        )
    # NOTE: We need to copy since we will lose the reference to the shared memory once
    # this function returns, making invalid all the buffers backing the ndarrays
    ret_tuple = recorddict, deepcopy(shm_parameters)
    # NOTE: Add the shared memory name to the garbage queue if it exists
    if params_garbage_queue is not None:
        params_garbage_queue.put(shm_name)
    return ret_tuple


def load_recorddict_parameters_from_remote(
    remote_uploader_downloader: RemoteUploaderDownloader | None,
    incoming_message: Message,
    msg_str: str,
    comm_stack: CommStack,
    params_garbage_queue: Queue[ray.ObjectRef | str] | None = None,
) -> tuple[Message, NDArrays]:
    """Load model parameters from remote storage into a message recorddict.

    This function retrieves model parameters that were previously offloaded to remote
    storage (S3, shared memory, or Ray's object store) and incorporates them back into
    the message recorddict. It first extracts basic parameter information from the
    recorddict, checks if the task was successful, and then dispatches to the
    appropriate handler based on the configured communication protocol. This approach
    allows efficient handling of large parameter sets by using specialized storage
    methods.

    Parameters
    ----------
    remote_uploader_downloader : RemoteUploaderDownloader | None
        The uploader/downloader instance for interacting with S3, or None if S3 is not
        used.
    incoming_message : Message
        The incoming message containing a recorddict with parameter references.
    msg_str : str
        A string identifier used to prefix keys in the recorddict.
    comm_stack : CommStack
        Configuration specifying which communication method(s) to use.
    params_garbage_queue : Queue[ray.ObjectRef | str] | None, optional
        Queue for Ray or Shm garbage collection, or None if not using Ray.

    Returns
    -------
    tuple[Message, NDArrays]
        A tuple containing the updated message with modified recorddict and the loaded
        model parameters.

    Notes
    -----
    This function optimizes parameter loading by:
    1. Extracting minimal parameter information before deciding how to load
    2. Only performing actual loading operations when needed based on status code
    3. Using specialized handlers for each storage method (S3, shared memory, Ray)
    4. Preserving the original message structure while updating its content

    """
    # Extract the content of the incoming message
    recorddict = incoming_message.content
    # Get the NDArrays from the parameters
    ndarrays: NDArrays = []
    array_record = recorddict.parameters_records.pop(f"{msg_str}.{PARAMETERS}", None)
    if array_record is not None:
        ndarrays = parameters_to_ndarrays(
            arrayrecord_to_parameters(
                record=array_record,
                keep_input=False,
            ),
        )

    # Check if the task was successful
    status = _extract_status_from_recorddict(msg_str, recorddict)
    if status.code != Code.OK:
        log(
            ERROR,
            "Task failed with code %s and message: %s",
            status.code,
            status.message,
        )
        # No translation performed as we assume the task failed
        return incoming_message, []

    # Dispatch to the appropriate handler based on communication protocol
    if comm_stack.s3 and remote_uploader_downloader is not None:
        recorddict, ndarrays = load_from_s3_parameters(
            remote_uploader_downloader,
            recorddict,
            msg_str,
        )
    elif comm_stack.shm:
        recorddict, ndarrays = load_from_shm_parameters(
            recorddict,
            msg_str,
            params_garbage_queue,
        )
    elif comm_stack.ray:
        recorddict, ndarrays = load_from_ray_parameters_with_strategy(
            recorddict,
            msg_str,
            params_garbage_queue,
        )

    # Update the message content with the modified recorddict
    incoming_message.content = recorddict
    return incoming_message, ndarrays
