"""Ray batching variants for parameter offloading/loading.

This module provides two alternative strategies to reduce the number of Ray tasks and
Ray object references created when offloading large lists of model parameter ndarrays
from Flower RecordDict structures:

Strategies
----------
1. batched_simple (per-array ObjectRefs, but fewer Ray remote tasks)
   - Greedy size-based batching of consecutive tensors up to a target cumulative byte
     size (``target_batch_bytes``).
   - One Ray remote task per batch converts each ndarray to float32 (if not already)
     and calls ``ray.put`` for every converted tensor inside the worker, returning a
     list of ``(original_index, ObjectRef)`` pairs.
   - The driver reorders the refs by original index and stores a serialized list of
     ObjectRefs in the RecordDict, plus a ``ray_strategy`` marker.

2. packed (contiguous packed buffer per batch)
   - Same batching logic, but each Ray task converts arrays to float32 then packs the
     flattened data into a single contiguous float32 buffer.
   - Metadata (original index, shape, offset, size) is produced per tensor.
   - Only one ObjectRef per batch is retained (the packed buffer). Metadata is
     serialized and stored in the RecordDict for reconstruction.

Loader
------
The provided ``load_from_ray_parameters_with_strategy`` function detects the strategy
via the ``ray_strategy`` key in the config record:
   * If missing -> legacy behavior (treat each ObjectRef as an individual float32
     tensor; convert to float32).
   * ``batched_simple`` -> same handling as legacy since each ObjectRef is a single
     float32 tensor.
   * ``packed`` -> unpacks buffers using stored metadata and restores original order.

Ordering
--------
Ordering is always preserved either by directly storing refs in original index order
(batched_simple) or by recording the original index in metadata (packed).

Usage Integration
-----------------
These functions are additive and do not modify existing logic. To integrate:
1. During offload, choose one of the new functions instead of the legacy
   ``offload_to_ray``.
2. During load, call ``load_from_ray_parameters_with_strategy`` instead of the legacy
   loader.

RecordDict Keys Added
---------------------
OBJECT_REFS          : Serialized list of ObjectRefs (array or packed buffer refs).
ray_strategy         : One of ("batched_simple", "packed"). Omitted for legacy.
ray_packed_meta      : Serialized metadata (only for packed strategy).

All other semantics (emptying parameters via ``parameters_records``) are kept.
"""

from __future__ import annotations

import os
import sys
from enum import Enum, auto
from logging import DEBUG
from typing import TYPE_CHECKING

import numpy as np
from flwr.common import ConfigRecord, Parameters, RecordDict, log
from flwr.common.recorddict_compat import parameters_to_arrayrecord

import ray
from mt_dao.constants import COMM_STACK, EMPTY, OBJECT_REFS, PARAMETERS

if TYPE_CHECKING:
    from queue import Queue

# Strategy markers
RAY_STRATEGY_KEY = "ray_strategy"
RAY_PACKED_META_KEY = "ray_packed_meta"

# Constants / defaults
_DEFAULT_TARGET_BATCH_BYTES = 64 * 1024 * 1024  # 64 MiB
_DEFAULT_MAX_BATCHES = 20
_DEFAULT_MIN_BATCHES_TO_MERGE = 2
_DEFAULT_LEGACY_RAY_WAIT_BATCH = 32

# Overwrite from environmental variables
mt_dao_RAY_TARGET_BATCH_BYTES = int(
    os.getenv("mt_dao_RAY_TARGET_BATCH_BYTES", str(_DEFAULT_TARGET_BATCH_BYTES)),
)
mt_dao_RAY_MAX_BATCHES = int(
    os.getenv("mt_dao_RAY_MAX_BATCHES", str(_DEFAULT_MAX_BATCHES)),
)
mt_dao_RAY_MIN_BATCHES_TO_MERGE = int(
    os.getenv("mt_dao_RAY_MIN_BATCHES_TO_MERGE", str(_DEFAULT_MIN_BATCHES_TO_MERGE)),
)
mt_dao_RAY_LEGACY_WAIT_BATCH = int(
    os.getenv("mt_dao_RAY_LEGACY_WAIT_BATCH", str(_DEFAULT_LEGACY_RAY_WAIT_BATCH)),
)


class RayOffloadStrategyName(str, Enum):  # noqa: UP042
    """Strategy type."""

    @staticmethod
    def _generate_next_value_(
        name: str,
        start: int,  # noqa: ARG004
        count: int,  # noqa: ARG004
        last_values: list[int],  # noqa: ARG004
    ) -> str:
        """Generate the next value.

        Replacement for StrEnum to support python 3.10

        Parameters
        ----------
        name: str
            Name of the strategy
        start: int
            Start value
        count: int
            Count value
        last_values: list[int]
            List of last values

        Returns
        -------
        str
            Lowercase name

        """
        return name.lower()

    BATCHED_SIMPLE = auto()
    PACKED = auto()
    STANDARD = auto()


class RayOffloadStrategyNotImplementedError(Exception):
    """Custom exception for unsupported strategies."""

    def __init__(self, strategy_name: str) -> None:
        """Initialize the exception with the strategy name.

        Parameters
        ----------
        strategy_name : str
            The name of the strategy that is not implemented or supported.

        """
        super().__init__(
            f"RayOffloadStrategy '{strategy_name}' is not implemented or supported.",
        )
        self.strategy_name = strategy_name


# ----------------------------- Remote Helper Functions -----------------------------


@ray.remote
def batch_convert_and_put_float32(array: np.ndarray) -> ray.ObjectRef:
    """Convert a single ndarray to float32 and put it in the Ray object store.

    This helper is used by the legacy offload path to spawn one Ray task per
    tensor. The input array is converted (or viewed, if already float32) and
    stored via ``ray.put``.

    Parameters
    ----------
    array : np.ndarray
        The ndarray to convert to ``float32`` and upload to the Ray object store.

    Returns
    -------
    ray.ObjectRef
        The Ray object reference to the converted array in the object store.

    """
    return ray.put(array.astype(np.float32))


@ray.remote
def pack_and_convert_batch(
    batch: list[tuple[int, np.ndarray]],
) -> tuple[list[dict], np.ndarray]:
    """Convert a batch of tensors to float32 and pack into a single buffer.

    Each element of ``batch`` is a tuple ``(original_index, array)``. Arrays are
    converted to ``float32`` (without copy when possible), flattened, and
    concatenated into one contiguous 1D buffer. Per-array metadata is produced to
    enable reconstructing the original tensors on the receiver.

    Parameters
    ----------
    batch : list[tuple[int, np.ndarray]]
        Consecutive tensors along with their original indices.

    Returns
    -------
    tuple[list[dict], np.ndarray]
        A pair ``(metadata, packed_buffer)`` where ``metadata`` is a list of
        dictionaries with keys:
        - ``idx`` (int): original tensor index
        - ``shape`` (tuple[int, ...]): original tensor shape
        - ``dtype`` (str): always ``"float32"``
        - ``offset_elems`` (int): starting element offset into the packed buffer
        - ``size_elems`` (int): number of elements for this tensor
        and ``packed_buffer`` is a ``float32`` ndarray containing all concatenated
        tensor data.

    """
    # Initialize conversion and packing variables
    converted: list[tuple[int, np.ndarray]] = []
    total_elems = 0

    # Loop over arrays to convert and calculate total size
    for idx, arr in batch:
        arr32 = arr if arr.dtype == np.float32 else arr.astype(np.float32, copy=False)
        converted.append((idx, arr32))
        total_elems += arr32.size

    # Vectorized sizes and offsets
    sizes = np.fromiter(
        (arr32.size for _, arr32 in converted),
        dtype=np.int64,
        count=len(converted),
    )
    offsets = np.zeros_like(sizes)
    if len(sizes) > 1:
        np.cumsum(sizes[:-1], out=offsets[1:])

    # Concatenate all data in C
    packed = np.concatenate(
        [arr32.ravel() for _, arr32 in converted],
    )

    # Build metadata (list comprehension)
    metadata_list = [
        {
            "idx": idx,
            "shape": arr32.shape,
            "dtype": "float32",
            "offset_elems": int(off),
            "size_elems": int(sz),
        }
        for (idx, arr32), off, sz in zip(converted, offsets, sizes, strict=True)
    ]

    # Return metadata and packed buffer
    return metadata_list, packed


# ------------------------------ Batching Helper Logic ------------------------------


def build_size_batches(
    arrays: list[np.ndarray],
    *,
    target_batch_bytes: int = mt_dao_RAY_TARGET_BATCH_BYTES,
    min_batches_to_merge: int = mt_dao_RAY_MIN_BATCHES_TO_MERGE,
    max_batches: int = mt_dao_RAY_MAX_BATCHES,
) -> list[list[tuple[int, np.ndarray]]]:
    """Group consecutive arrays into size-limited batches.

    Arrays are traversed in order and greedily accumulated into a batch until the
    cumulative byte-size would exceed ``target_batch_bytes``. If the resulting
    number of batches exceeds ``max_batches``, the two smallest batches are merged
    repeatedly (at least ``min_batches_to_merge`` must be available) until the
    limit is met. Original ordering is preserved via indices stored alongside each
    array.

    Parameters
    ----------
    arrays : list[np.ndarray]
        The arrays to partition into batches.
    target_batch_bytes : int, optional
        Target cumulative size in bytes for each batch, by default
        ``mt_dao_RAY_TARGET_BATCH_BYTES``.
    min_batches_to_merge : int, optional
        Minimum number of batches required to attempt a merge step, by default
        ``mt_dao_RAY_MIN_BATCHES_TO_MERGE``.
    max_batches : int, optional
        Upper bound on the number of batches, by default ``mt_dao_RAY_MAX_BATCHES``.

    Returns
    -------
    list[list[tuple[int, np.ndarray]]]
        A list of batches; each batch is a list of ``(original_index, array)``
        pairs.

    """
    # Initialize batching variables
    batches: list[list[tuple[int, np.ndarray]]] = []
    current: list[tuple[int, np.ndarray]] = []
    cur_bytes = 0

    # Loop over the arrays to pack in batches
    for idx, arr in enumerate(arrays):
        # Get the current array's byte size
        arr_bytes = arr.nbytes
        # If the current batch is not empty and adding this array would exceed the
        # target size, start a new batch
        if current and cur_bytes + arr_bytes > target_batch_bytes:
            # Append the current batch
            batches.append(current)
            # Clear the bytes counter and buffer
            current = []
            cur_bytes = 0
        # Append the current array to the batch and increment the bytes counter
        current.append((idx, arr))
        cur_bytes += arr_bytes

    # Append any remaining arrays as the last batch
    if current:
        batches.append(current)

    # Correct for too many batches by merging smallest ones
    if len(batches) > max_batches:
        # Merge smallest batches until we are within the limit
        while len(batches) > max_batches:
            # Exit condition if too few batches to merge
            if len(batches) < min_batches_to_merge:
                break
            # Estimate batch sizes
            sizes = [sum(a.nbytes for _i, a in b) for b in batches]
            # Find two smallest batches
            s1, s2 = sorted(range(len(sizes)), key=lambda k: sizes[k])[:2]
            # Swap to ensure s1 < s2 for safe deletion
            if s1 > s2:
                s1, s2 = s2, s1
            # Merge s2 into s1 and delete s2
            batches[s1].extend(batches[s2])
            # Delete the second batch
            del batches[s2]

    # Return the list of batches
    return batches


# ------------------------------ Offload Variants ----------------------------------


def offload_to_ray_legacy(
    parameters: list[np.ndarray],
    recorddict: RecordDict,
    msg_str: str,
) -> tuple[RecordDict, list[ray.ObjectRef]]:
    """Offload parameters to Ray using the legacy per-array strategy.

    Spawns one Ray task per tensor to convert to ``float32`` and ``ray.put`` it.
    Waits for completions in small batches to reduce driver pressure, preserves
    original ordering, stores serialized ObjectRefs under ``OBJECT_REFS`` and sets
    the ``ray_strategy`` marker to ``STANDARD``. The parameters payload is cleared
    from the ``RecordDict``.

    Parameters
    ----------
    parameters : list[np.ndarray]
        The list of tensors to offload.
    recorddict : flwr.common.RecordDict
        Mutable RecordDict to be updated with ObjectRefs and markers.
    msg_str : str
        Prefix used for namespacing in ``recorddict``.

    Returns
    -------
    tuple[RecordDict, list[ray.ObjectRef]]
        The updated ``RecordDict`` and the list of ObjectRefs (in original order).

    """
    # Create futures for each item in the ndarrays list
    futures = [batch_convert_and_put_float32.remote(ndarray) for ndarray in parameters]
    # Collect ObjectRefs into a dictionary to keep the original ordering
    ordering = {ref: i for i, ref in enumerate(futures)}
    # Build a mutable sequence of remaining refs to process
    remaining = list(futures)
    # Initialize the output dictionary
    ray_object_refs: dict[int, ray.ObjectRef] = {}

    # Loop until all refs are processed
    while remaining:

        # Wait for a batch of refs to be ready
        ready, remaining = ray.wait(
            remaining,
            num_returns=min(mt_dao_RAY_LEGACY_WAIT_BATCH, len(remaining)),
        )

        # Get the ready ObjectRefs
        ready_ray_object_refs = ray.get(ready)

        # Update the output dictionary in the original order
        ray_object_refs.update(
            {
                ordering[r]: ref
                for r, ref in zip(ready, ready_ray_object_refs, strict=True)
            },
        )

    # Put serialized ObjectRef in recorddict using vectorized operation
    list_of_ray_object_refs = [ray_object_refs[i] for i in range(len(ray_object_refs))]
    recorddict.config_records[f"{msg_str}.{COMM_STACK}"].update(
        ConfigRecord(
            {
                OBJECT_REFS: ray.cloudpickle.dumps(list_of_ray_object_refs),  # type: ignore[reportAttributeAccessIssue]
                RAY_STRATEGY_KEY: RayOffloadStrategyName.STANDARD,
            },
        ),
    )
    # Empty the recorddict parameters
    recorddict.parameters_records[f"{msg_str}.{PARAMETERS}"] = (
        parameters_to_arrayrecord(
            Parameters(tensors=[], tensor_type=EMPTY),
            keep_input=False,
        )
    )
    return recorddict, list_of_ray_object_refs


def offload_to_ray(
    parameters: list[np.ndarray],
    recorddict: RecordDict,
    msg_str: str,
    *,
    is_packed: bool = False,
) -> tuple[RecordDict, list[ray.ObjectRef]]:
    """Offload parameters to Ray using batched or packed variants.

    Tensors are grouped with :func:`build_size_batches`. For each batch, a single
    remote task converts arrays to ``float32`` and returns packed data plus
    metadata. The packed buffers are uploaded with ``ray.put`` and both the buffer
    refs and metadata are serialized into ``recorddict``.

    When ``is_packed`` is ``True``, a single batch is forced to minimize the number
    of Ray objects (one packed buffer). When ``False``, multiple batches can be
    created subject to size limits.

    Parameters
    ----------
    parameters : list[np.ndarray]
        The list of tensors to offload.
    recorddict : flwr.common.RecordDict
        Mutable RecordDict to be updated with ObjectRefs and markers.
    msg_str : str
        Prefix used for namespacing in ``recorddict``.
    is_packed : bool, optional
        If ``True``, force a single packed buffer (``PACKED`` strategy); otherwise
        allow multiple batches (``BATCHED_SIMPLE``), by default ``False``.

    Returns
    -------
    tuple[RecordDict, list[ray.ObjectRef]]
        The updated ``RecordDict`` and the list of buffer ObjectRefs (one per
        batch).

    """
    # Create the one batch of all arrays
    batches = build_size_batches(
        arrays=list(parameters),
        # NOTE: We force one single batch in the packed strategy to avoid
        # multiple packed buffers which would increase the number of ObjectRefs
        # and Ray tasks, defeating the purpose of this strategy.
        target_batch_bytes=sys.maxsize if is_packed else mt_dao_RAY_TARGET_BATCH_BYTES,
        min_batches_to_merge=(
            sys.maxsize if is_packed else mt_dao_RAY_MIN_BATCHES_TO_MERGE
        ),
        max_batches=1 if is_packed else mt_dao_RAY_MAX_BATCHES,
    )
    assert not is_packed or len(batches) == 1, f"Packed strategy should create exactly one batch, not {len(batches)}."
    # Create one Ray task per batch to convert and pack
    futures = [pack_and_convert_batch.remote(batch) for batch in batches]

    # Gather results and put packed buffers
    packed_buffers: list[ray.ObjectRef] = []
    packed_metadata: list[list[dict]] = []
    for meta_list, packed in ray.get(futures):
        packed_metadata.append(meta_list)
        packed_buffers.append(ray.put(packed))

    # Put serialized ObjectRef and metadata in recorddict
    recorddict.config_records[f"{msg_str}.{COMM_STACK}"].update(
        ConfigRecord(
            {
                OBJECT_REFS: ray.cloudpickle.dumps(packed_buffers),  # type: ignore[attr-defined]
                RAY_STRATEGY_KEY: (
                    RayOffloadStrategyName.PACKED
                    if is_packed
                    else RayOffloadStrategyName.BATCHED_SIMPLE
                ),
                RAY_PACKED_META_KEY: ray.cloudpickle.dumps(packed_metadata),  # type: ignore[attr-defined]
            },
        ),
    )
    # Empty the recorddict parameters
    recorddict.parameters_records[f"{msg_str}.{PARAMETERS}"] = (
        parameters_to_arrayrecord(
            Parameters(tensors=[], tensor_type=EMPTY),
            keep_input=False,
        )
    )
    return recorddict, packed_buffers


# ------------------------------ Load (Strategy Aware) ------------------------------


def resolve_meta(
    ready_meta: list[list[dict]],
    arrays_fp32: list[np.ndarray],
    tensor_dict: dict[int, np.ndarray],
) -> None:
    """Reconstruct tensors from packed buffers using metadata.

    Parameters
    ----------
    ready_meta : list[list[dict]]
        Metadata lists aligned to ``arrays_fp32``. Each inner dict must contain
        ``idx``, ``shape``, ``offset_elems``, and ``size_elems``.
    arrays_fp32 : list[np.ndarray]
        The packed buffers (each a flat ``float32`` ndarray) corresponding to the
        metadata entries.
    tensor_dict : dict[int, np.ndarray]
        Output mapping from original tensor index to reconstructed ndarray. Filled
        in-place.

    """
    # Unpack each buffer using the corresponding metadata
    for meta_list, buf_fp32 in zip(
        ready_meta,
        arrays_fp32,
        strict=True,
    ):
        # Extract each tensor using the metadata
        for meta in meta_list:
            idx = meta["idx"]
            off = meta["offset_elems"]
            size = meta["size_elems"]
            shape = tuple(meta["shape"])
            # Slice and reshape
            tensor_dict[idx] = buf_fp32[off : off + size].reshape(shape)


def load(
    obj_refs: list[ray.ObjectRef],
    ray_gc_queue: Queue[ray.ObjectRef | str] | None,
    packed_meta: list[list[dict]] | None = None,
) -> list[np.ndarray]:
    """Load arrays (or packed buffers) from Ray and return tensors in order.

    This function waits on Ray object references, retrieves the arrays, ensures
    ``float32`` dtype, and either returns them directly (no packing) or
    reconstructs individual tensors using provided ``packed_meta``.

    Parameters
    ----------
    obj_refs : list[ray.ObjectRef]
        References to arrays or packed buffers stored in the Ray object store.
    ray_gc_queue : Queue[ray.ObjectRef | str] | None
        Optional queue used to enqueue consumed refs for garbage collection.
    packed_meta : list[list[dict]] | None, optional
        If provided, metadata describing how to slice each packed buffer into the
        original tensors, by default ``None``.

    Returns
    -------
    list[np.ndarray]
        Tensors in their original ordering, always as ``float32`` ndarrays.

    """
    # Collect ObjectRefs into a dictionary to keep the original ordering
    ordering = {ref: i for i, ref in enumerate(obj_refs)}
    # Build a mutable sequence of remaining refs to process
    remaining = list(obj_refs)
    # Initialize the output dictionary
    tensor_dict: dict[int, np.ndarray] = {}
    # Loop until all refs are processed
    while remaining:

        # Wait for a batch of refs to be ready
        ready, remaining = ray.wait(
            remaining,
            num_returns=min(mt_dao_RAY_LEGACY_WAIT_BATCH, len(remaining)),
        )

        # Get the arrays locally and ensure float32 without spawning extra remote tasks
        arrays: list[np.ndarray] = ray.get(ready)
        arrays_fp32: list[np.ndarray] = [
            arr if arr.dtype == np.float32 else arr.astype(np.float32, copy=False)
            for arr in arrays
        ]

        # Add the resolved arrays to the output dictionary
        if packed_meta is not None:
            # Resolve using metadata
            resolve_meta(
                ready_meta=[packed_meta[ordering[r]] for r in ready],
                arrays_fp32=arrays_fp32,
                tensor_dict=tensor_dict,
            )
        else:
            # Update the output dictionary in the original order
            tensor_dict.update(
                {ordering[r]: arr for r, arr in zip(ready, arrays_fp32, strict=True)},
            )

        # Optionally enqueue for garbage collection
        if ray_gc_queue is not None:
            for r in ready:
                ray_gc_queue.put(r)

    # Return the list of arrays in the original order
    return [tensor_dict[i] for i in range(len(tensor_dict))]


def load_from_ray_parameters_with_strategy(
    recorddict: RecordDict,
    msg_str: str,
    ray_gc_queue: Queue[ray.ObjectRef | str] | None = None,
) -> tuple[RecordDict, list[np.ndarray]]:
    """Load Ray-offloaded parameters handling legacy/batched/packed strategies.

    Parameters
    ----------
    recorddict : RecordDict
        RecordDict containing strategy markers and serialized ObjectRefs.
    msg_str : str
        Prefix used when storing records.
    ray_gc_queue : Queue | None
        Optional queue for garbage collection of Ray ObjectRefs.

    Returns
    -------
    (RecordDict, list[np.ndarray])
        The (unchanged) RecordDict and list of float32 tensors (restored order).

    """
    # Collect the configuration of the current communication stack message
    conf = recorddict.config_records[f"{msg_str}.{COMM_STACK}"]
    # Determine strategy and load accordingly
    strategy_name = conf.get(RAY_STRATEGY_KEY, "")
    log(DEBUG, "Loading parameters with strategy: '%s'", strategy_name)
    strategy = RayOffloadStrategyName[str(strategy_name).upper()]
    # Collect serialized ObjectRefs
    obj_refs = ray.cloudpickle.loads(conf[OBJECT_REFS])  # type: ignore[attr-defined]
    # Dispatch to appropriate loader
    if strategy in {
        RayOffloadStrategyName.PACKED,
        RayOffloadStrategyName.BATCHED_SIMPLE,
    }:
        # Load the metadata for unpacking
        packed_meta = ray.cloudpickle.loads(conf[RAY_PACKED_META_KEY])  # type: ignore[attr-defined]
        tensors = load(
            obj_refs=obj_refs,
            ray_gc_queue=ray_gc_queue,
            packed_meta=packed_meta,
        )
    else:
        tensors = load(obj_refs=obj_refs, ray_gc_queue=ray_gc_queue)
    return recorddict, tensors


__all__ = [
    "RAY_PACKED_META_KEY",
    "RAY_STRATEGY_KEY",
    "load_from_ray_parameters_with_strategy",
    "offload_to_ray",
    "offload_to_ray_legacy",
]
