import datetime
import logging
import os
import shutil
import time

import numpy
import pygloo

import ray
from ray._private import ray_constants
from ray.util.collective.collective_group import gloo_util
from ray.util.collective.collective_group.base_collective_group import BaseGroup
from ray.util.collective.const import get_store_name
from ray.util.collective.types import (
    AllGatherOptions,
    AllReduceOptions,
    Backend,
    BarrierOptions,
    BroadcastOptions,
    RecvOptions,
    ReduceOptions,
    ReduceScatterOptions,
    SendOptions,
)

logger = logging.getLogger(__name__)


class Rendezvous:
    """A rendezvous class for different actor/task processes to meet.

    To initialize an GLOO collective communication group, different
    actors/tasks spawned in Ray in a collective group needs to meet
    each other to synchronize the GLOOUniqueID. This class guarantees
    they meet via the GLOOUniqueIDStore, initialized on the rank=0
    process.

    Args:
        group_name: the unique user-specified group name.
    """

    def __init__(self, group_name, context, store_type, device_type):
        self._group_name = group_name
        self._context = context
        redis_address = ray._private.worker._global_node.redis_address
        (self._redis_ip_address, self._redis_port) = (
            redis_address.split(":") if store_type == "redis" else (None, None)
        )
        self._process_ip_address = ray.util.get_node_ip_address()
        logger.debug(
            "Redis address: {}, port: {}, this actor address: {}.".format(
                self._redis_ip_address, self._redis_port, self._process_ip_address
            )
        )
        self._store_type = store_type
        self._device_type = device_type
        self._store = None
        self._device = None
        self.create_store(store_type)
        self.create_device(device_type)

    def create_store(self, store_type):
        if store_type == "ray_internal_kv":
            ray_internal_kv_store = gloo_util.RayInternalKvStore(self._group_name)
            self._store = pygloo.rendezvous.CustomStore(ray_internal_kv_store)
        elif store_type == "redis":
            redisStore = pygloo.rendezvous.RedisStore(
                self._redis_ip_address, int(self._redis_port)
            )
            redis_password = ray._private.worker._global_node.redis_password
            if redis_password is None or len(redis_password) == 0:
                redis_password = ray_constants.REDIS_DEFAULT_PASSWORD
            redisStore.authorize(redis_password)
            self._store = redisStore
        elif store_type == "file":
            store_name = get_store_name(self._group_name)
            store_path = gloo_util.get_gloo_store_path(store_name)
            if self._context.rank == 0:
                if not os.path.exists(store_path):
                    os.makedirs(store_path)
                elif os.listdir(store_path) and os.listdir(store_path):
                    shutil.rmtree(store_path)
                    os.makedirs(store_path)
            else:
                while not os.path.exists(store_path):
                    time.sleep(0.1)
            # Note: multi-machines needs a shared NFS.
            fileStore = pygloo.rendezvous.FileStore(store_path)
            self._store = pygloo.rendezvous.PrefixStore(self._group_name, fileStore)
        elif store_type == "hash":
            raise NotImplementedError("No implementation for hash store.")
        else:
            raise RuntimeError("Unrecognized store type: {}.".format(store_type))

    def create_device(self, device_type):
        if device_type == "tcp":
            attr = pygloo.transport.tcp.attr(self._process_ip_address)
            self._device = pygloo.transport.tcp.CreateDevice(attr)
        elif device_type == "uv":
            raise NotImplementedError("No implementation for uv.")

    def meet(self, timeout_s=180):
        """Meet at the named actor store.

        Args:
            timeout_s: timeout in seconds.

        Return:
            None
        """
        if timeout_s <= 0:
            raise ValueError(
                "The 'timeout' argument must be positive. "
                "Got '{}'.".format(timeout_s)
            )

        timeout_delta = datetime.timedelta(seconds=timeout_s)
        elapsed = datetime.timedelta(seconds=0)
        start_time = datetime.datetime.now()
        q, s = None, None

        if self._store_type == "redis" or self._store_type == "ray_internal_kv":
            while elapsed < timeout_delta:
                try:
                    # I don't quite understand why we need gloo queue actor.
                    q = ray.get_actor("gloo_queue")
                    s = ray.get_actor(f"gloo_{self._group_name}_signal")
                    break
                except ValueError:
                    if self._context.rank == 0:
                        if not q:
                            ray.remote(gloo_util.glooQueue).options(
                                name="gloo_queue", lifetime="detached"
                            ).remote(1000)
                        if not s:
                            gloo_util.SignalActor.options(
                                name=f"gloo_{self._group_name}_signal",
                                lifetime="detached",
                            ).remote(self._context.size)
                    else:
                        time.sleep(0.1)
                elapsed = datetime.datetime.now() - start_time
            if not q:
                raise RuntimeError("Unable to get gloo_queue.")
            if self._context.rank == 0:
                ray.get(q.put_nowait.remote(self._group_name))
            while ray.get(q.index.remote(self._group_name)):
                time.sleep(0.1)

            self._context.connectFullMesh(self._store, self._device)
            ray.get(s.send.remote(self._context.rank))
            if self._context.rank == 0:
                ray.get(s.wait.remote())
                keys = []
                keys += [f"rank_{i}" for i in range(self._context.size)]
                keys += [f"{i}" for i in range(self._context.size)]
                self._store.delKeys(keys)
                group_name = ray.get(q.get_nowait.remote())
                assert group_name == self._group_name
                ray.kill(s)

    @property
    def store_type(self):
        return self._store_type

    @property
    def store(self):
        return self._store

    @property
    def device_type(self):
        return self._device_type

    @property
    def device(self):
        return self._device

    def destroy(self):
        """GC the store and device used by this rendevzous."""
        self._device = None


class GLOOGroup(BaseGroup):
    def __init__(
        self,
        world_size,
        rank,
        group_name,
        store_type="ray_internal_kv",
        device_type="tcp",
    ):
        """Init an GLOO collective group.

        Args:
            world_size: The number of processes.
            rank: The id of process
            group_name: The unique user-specified group name.
            store_type: The store type. Optional: "redis",
                              "file", "hash".
            device_type: The device type to transport.
                               Optional: "tcp", "uv".
        """
        super(GLOOGroup, self).__init__(world_size, rank, group_name)
        self._gloo_context = gloo_util.create_gloo_context(self.rank, self.world_size)
        self._rendezvous = Rendezvous(
            self.group_name, self._gloo_context, store_type, device_type
        )
        self._rendezvous.meet()

    def destroy_group(self):
        """Destroy the group and release GLOO communicators."""
        self._rendezvous.destroy()

        if self._gloo_context is not None:
            pygloo.barrier(self._gloo_context)
            # destroy the communicator
            self._gloo_context = None

        if self.rank == 0 and self._rendezvous.store_type == "file":
            store_name = get_store_name(self._group_name)
            store_path = gloo_util.get_gloo_store_path(store_name)
            if os.path.exists(store_path):
                shutil.rmtree(store_path)
        super(GLOOGroup, self).destroy_group()

    @classmethod
    def backend(cls):
        return Backend.GLOO

    def allreduce(self, tensors, allreduce_options=AllReduceOptions()):
        """AllReduce a list of tensors following options.

        Args:
            tensor: the tensor to be reduced, each tensor locates on CPU
            allreduce_options:

        Returns:
            None
        """

        def collective_fn(input_tensor, output_tensor, context):
            pygloo.allreduce(
                context,
                gloo_util.get_tensor_ptr(input_tensor),
                gloo_util.get_tensor_ptr(output_tensor),
                gloo_util.get_tensor_n_elements(input_tensor),
                gloo_util.get_gloo_tensor_dtype(input_tensor),
                gloo_util.get_gloo_reduce_op(allreduce_options.reduceOp),
            )

        self._collective(tensors, tensors, collective_fn)

    def barrier(self, barrier_options=BarrierOptions()):
        """Blocks until all processes reach this barrier.

        Args:
            barrier_options: barrier options.

        Returns:
            None
        """
        barrier_tensor = numpy.array([1])
        self.allreduce([barrier_tensor])

    def reduce(self, tensors, reduce_options=ReduceOptions()):
        """Reduce tensors following options.

        Args:
            tensors: the list of tensors to be reduced,
                            this list only have one tensor.
            reduce_options: reduce options.

        Returns:
            None
        """
        root_rank = reduce_options.root_rank

        def collective_fn(input_tensor, output_tensor, context):
            pygloo.reduce(
                context,
                gloo_util.get_tensor_ptr(input_tensor),
                gloo_util.get_tensor_ptr(output_tensor),
                gloo_util.get_tensor_n_elements(input_tensor),
                gloo_util.get_gloo_tensor_dtype(input_tensor),
                gloo_util.get_gloo_reduce_op(reduce_options.reduceOp),
                root_rank,
            )

        self._collective(tensors, tensors, collective_fn)

    def broadcast(self, tensors, broadcast_options=BroadcastOptions()):
        """Broadcast tensors to all other processes following options.

        Args:
            tensors: tensors to be broadcast or received.
            broadcast_options: broadcast options.

        Returns:
            None
        """
        root_rank = broadcast_options.root_rank

        def collective_fn(input_tensor, output_tensor, context):
            pygloo.broadcast(
                context,
                gloo_util.get_tensor_ptr(input_tensor),
                gloo_util.get_tensor_ptr(output_tensor),
                gloo_util.get_tensor_n_elements(input_tensor),
                gloo_util.get_gloo_tensor_dtype(input_tensor),
                root_rank,
            )

        self._collective(tensors, tensors, collective_fn)

    def allgather(self, tensor_lists, tensors, allgather_options=AllGatherOptions()):
        """Allgather tensors on CPU into a list of tensors.

        Args:
            tensor_lists (List[List[Tensor]]): allgathered tensors.
            tensors: the list of tensors to allgather across the group.
                     Each tensor must locate on CPU.
            allgather_options: allgather options.

        Returns:
            None
        """

        def collective_fn(input_tensor, output_tensor, context):
            pygloo.allgather(
                context,
                gloo_util.get_tensor_ptr(input_tensor),
                gloo_util.get_tensor_ptr(output_tensor),
                gloo_util.get_tensor_n_elements(input_tensor),
                gloo_util.get_gloo_tensor_dtype(input_tensor),
            )

        _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
        output_flattened = [
            _flatten_for_scatter_gather(tensor_list, copy=False)
            for tensor_list in tensor_lists
        ]

        def postprocess_fn():
            for i, tensor_list in enumerate(tensor_lists):
                for j, tensor in enumerate(tensor_list):
                    gloo_util.copy_tensor(tensor, output_flattened[i][j])

        self._collective(
            tensors, output_flattened, collective_fn, postprocess_fn=postprocess_fn
        )

    def reducescatter(
        self, tensors, tensor_lists, reducescatter_options=ReduceScatterOptions()
    ):
        """Reduce the scatter a list of tensors across the group.

        Args:
            tensors: the output tensors (could be unspecified), each
                            located on CPU.
            tensor_lists (List[List]): the list of tensors to be reduced then
                                       scattered.
            reducescatter_options: reduce-scatter options.

        Returns:
            None
        """

        def collective_fn(input_tensor, output_tensor, context):
            size = gloo_util.get_tensor_n_elements(input_tensor)
            world_size = self._gloo_context.size
            pygloo.reduce_scatter(
                context,
                gloo_util.get_tensor_ptr(input_tensor),
                gloo_util.get_tensor_ptr(output_tensor),
                size,
                [size // world_size for _ in range(world_size)],
                gloo_util.get_gloo_tensor_dtype(output_tensor),
                gloo_util.get_gloo_reduce_op(reducescatter_options.reduceOp),
            )

        _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
        input_flattened = [
            _flatten_for_scatter_gather(tensor_list, copy=False)
            for tensor_list in tensor_lists
        ]

        def preprocess_fn():
            for i, tensor_list in enumerate(tensor_lists):
                for j, tensor in enumerate(tensor_list):
                    gloo_util.copy_tensor(input_flattened[i][j], tensor)

        self._collective(
            input_flattened, tensors, collective_fn, preprocess_fn=preprocess_fn
        )

    def send(self, tensors, send_options=SendOptions()):
        """Send a tensor to a destination rank in the group.

        Args:
            tensors: the tensor to send.
            send_options: send options.

        Returns:
            None
        """

        def p2p_fn(tensor, context, peer):
            pygloo.send(
                context,
                gloo_util.get_tensor_ptr(tensor),
                gloo_util.get_tensor_n_elements(tensor),
                gloo_util.get_gloo_tensor_dtype(tensor),
                peer,
            )

        self._point2point(tensors, p2p_fn, send_options.dst_rank)

    def recv(self, tensors, recv_options=RecvOptions()):
        """Receive a tensor from a source rank in the group.

        Args:
            tensors: the received tensor.
            recv_options: Receive options.

        Returns:
            None
        """

        def p2p_fn(tensor, context, peer):
            pygloo.recv(
                context,
                gloo_util.get_tensor_ptr(tensor),
                gloo_util.get_tensor_n_elements(tensor),
                gloo_util.get_gloo_tensor_dtype(tensor),
                peer,
            )

        self._point2point(tensors, p2p_fn, recv_options.src_rank)

    def _collective(
        self,
        input_tensors,
        output_tensors,
        collective_fn,
        preprocess_fn=None,
        postprocess_fn=None,
    ):
        """A method to encapsulate all collective calls.

        Args:
            input_tensors: the list of the input tensors.
            output_tensors: the list of the output tensors.
            collective_fn: the collective function call.
            preprocess_fn: preprocess procedures before collective calls.
            postprocess_fn: postprocess procedures after collective calls.

        Returns:
            None
        """
        _check_cpu_tensors(input_tensors)
        _check_cpu_tensors(output_tensors)

        if preprocess_fn:
            preprocess_fn()
        collective_fn(input_tensors[0], output_tensors[0], self._gloo_context)
        if postprocess_fn:
            postprocess_fn()

    def _point2point(self, tensors, p2p_fn, peer_rank: int):
        """A method to encapsulate all peer-to-peer calls (i.e., send/recv).

        Args:
            tensors: the tensor to send or receive.
            p2p_fn: the p2p function call.
            peer_rank: the rank of the peer process.

        Returns:
            None
        """
        _check_cpu_tensors(tensors)

        p2p_fn(tensors[0], self._gloo_context, peer_rank)


def _check_cpu_tensors(tensors):
    """Check only have one tensor and located on CPU."""
    if not tensors or not isinstance(tensors, list):
        raise RuntimeError("'tensors' must be a nonempty list.")
    if len(tensors) != 1:
        raise RuntimeError(
            "Gloo only accept one tensor in the tensor list."
            " Got {} != 1.".format(len(tensors))
        )
    d = gloo_util.get_tensor_device(tensors[0])
    if d != "cpu":
        raise RuntimeError("Gloo only accept cpu tensor . Got {}.".format(d))


def _flatten_for_scatter_gather(tensor_list, copy=False):
    """Flatten the tensor for gather/scatter operations.

    Args:
        tensor_list: the list of tensors to be scattered/gathered.
        copy: whether the copy the tensors in tensor_list into the buffer.

    Returns:
        The flattened tensor buffer.
    """
    if not tensor_list:
        raise RuntimeError("Received an empty list.")

    t = tensor_list[0]
    # note we need a numpy dtype here.
    dtype = gloo_util.get_numpy_tensor_dtype(t)
    buffer_shape = [len(tensor_list)] + gloo_util.get_tensor_shape(t)

    buffer = numpy.empty(buffer_shape, dtype=dtype)
    if copy:
        for i, tensor in enumerate(tensor_list):
            gloo_util.copy_tensor(buffer[i], tensor)
    return buffer


def _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists):
    """Check the compatibility between tensor input and tensor list input."""
    if not tensors or not isinstance(tensors, list):
        raise RuntimeError("The first argument 'tensors' expects a list of tensors.")

    if len(tensors) != 1:
        raise RuntimeError(
            "Gloo only accept one tensor in the first argument 'tensors'."
            " Got {} != 1.".format(len(tensors))
        )

    if not tensor_lists or not isinstance(tensor_lists, list):
        raise RuntimeError(
            "The second argument 'tensor_lists' expects a list of tensor list."
        )

    if len(tensor_lists) != 1:
        raise RuntimeError(
            "Gloo only accept one tensor list "
            "in the second argument 'tensor_lists'."
            " Got {} != 1.".format(len(tensor_lists))
        )

    dtype = gloo_util.get_gloo_tensor_dtype(tensors[0])
    shape = gloo_util.get_tensor_shape(tensors[0])

    # check all tensors in `tensor_lists` match.
    for t in tensor_lists[0]:
        # check dtype
        dt = gloo_util.get_gloo_tensor_dtype(t)
        if dt != dtype:
            raise RuntimeError(
                "All tensor operands to scatter/gather must "
                "have the same dtype. Got '{}' and '{}'.".format(dt, dtype)
            )
        s = gloo_util.get_tensor_shape(t)
        if s != shape:
            raise RuntimeError(
                "All tensor operands to scatter/gather must "
                "have the same shape. Got '{}' and '{}'.".format(s, shape)
            )
