import sys
import time

import numpy as np
import pandas as pd
import pytest

import ray


def test_uninitialized():
    with pytest.raises(RuntimeError):
        ray.experimental.get_object_locations([])


def test_get_locations_empty_list(ray_start_regular):
    locations = ray.experimental.get_object_locations([])
    assert len(locations) == 0


def test_get_locations(ray_start_regular):
    node_id = ray.get_runtime_context().get_node_id()
    sizes = [100, 1000]
    obj_refs = [ray.put(np.zeros(s, dtype=np.uint8)) for s in sizes]
    ray.wait(obj_refs)
    locations = ray.experimental.get_object_locations(obj_refs)
    assert len(locations) == 2
    for idx, obj_ref in enumerate(obj_refs):
        location = locations[obj_ref]
        assert location["object_size"] > sizes[idx]
        assert location["node_ids"] == [node_id]


def test_get_locations_inlined(ray_start_regular):
    node_id = ray.get_runtime_context().get_node_id()
    obj_refs = [ray.put("123")]
    ray.wait(obj_refs)
    locations = ray.experimental.get_object_locations(obj_refs)
    for idx, obj_ref in enumerate(obj_refs):
        location = locations[obj_ref]
        assert location["node_ids"] == [node_id]
        assert location["object_size"] > 0


def test_spilled_locations(ray_start_cluster_enabled):
    cluster = ray_start_cluster_enabled
    cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    ray.init(cluster.address)
    cluster.wait_for_nodes()

    node_id = ray.get_runtime_context().get_node_id()

    @ray.remote
    def task():
        arr = np.random.rand(5 * 1024 * 1024)  # 40 MB
        refs = []
        refs.extend([ray.put(arr) for _ in range(2)])
        ray.get(ray.put(arr))
        ray.get(ray.put(arr))
        return refs

    object_refs = ray.get(task.remote())
    ray.wait(object_refs)
    locations = ray.experimental.get_object_locations(object_refs)
    for obj_ref in object_refs:
        location = locations[obj_ref]
        assert location["node_ids"] == [node_id]
        assert location["object_size"] > 0


def test_get_locations_multi_nodes(ray_start_cluster_enabled):
    cluster = ray_start_cluster_enabled
    # head node
    cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    ray.init(cluster.address)
    # add 1 worker node
    cluster.add_node(
        num_cpus=0, resources={"custom": 1}, object_store_memory=75 * 1024 * 1024
    )
    cluster.wait_for_nodes()

    all_node_ids = list(map(lambda node: node["NodeID"], ray.nodes()))
    driver_node_id = ray.get_runtime_context().get_node_id()
    all_node_ids.remove(driver_node_id)
    worker_node_id = all_node_ids[0]

    @ray.remote(num_cpus=0, resources={"custom": 1})
    def create_object():
        return np.random.rand(1 * 1024 * 1024)

    @ray.remote
    def task():
        return [create_object.remote()]

    object_refs = ray.get(task.remote())
    ray.wait(object_refs)
    locations = ray.experimental.get_object_locations(object_refs)
    for obj_ref in object_refs:
        location = locations[obj_ref]
        assert set(location["node_ids"]) == {driver_node_id, worker_node_id}
        assert location["object_size"] > 0


def test_location_pending(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    ray.init(cluster.address)
    cluster.wait_for_nodes()

    @ray.remote
    def task():
        # sleep for 1 hour so the object will be pending
        time.sleep(3600)
        return 1

    object_ref = task.remote()
    locations = ray.experimental.get_object_locations([object_ref])
    assert len(locations) == 1
    location = locations[object_ref]
    assert location["node_ids"] == []
    assert location["object_size"] is None

    local_locations = ray.experimental.get_local_object_locations([object_ref])
    assert len(local_locations) == 1
    local_location = local_locations[object_ref]
    assert local_location["node_ids"] == []
    assert local_location["object_size"] is None


# Tests for `get_local_object_locations`. We use matrix test:
# - callee can be regular ray task, or streaming generator;
# - caller can be in the same node (single node cluster), or different node.
#
# ... so we have 4 tests.
#
# Each test has the caller to produce Object(s) that consumes big memory but has a small
# sys.getsizeof. The caller then asserts the object size from the API
# `ray.experimental.get_local_object_locations` is > the actual memory consumed.

BIG_OBJ_SIZE = 3 * 1024 * 1024  # 3 MiB


class BigObject:
    def __init__(self):
        self.data = np.zeros((BIG_OBJ_SIZE,), dtype=np.uint8)


@ray.remote
def gen_big_object(block_size):
    return pd.DataFrame([{"data": BigObject()} for _ in range(block_size)])


@ray.remote
def gen_big_objects(block_size, block_count):
    for _ in range(block_count):
        big_object = ray.get(gen_big_object.remote(block_size))
        yield big_object


def assert_object_size_gt(obj_ref: ray.ObjectRef, size: int):
    d = ray.experimental.get_local_object_locations([obj_ref])
    assert d is not None
    assert len(d) == 1
    assert d[obj_ref]["object_size"] > size


def test_get_local_locations(ray_start_regular):
    """
    caller and callee are in the same node.
    callee is a regular ray task.
    """
    obj_ref = gen_big_object.remote(3)
    ray.wait([obj_ref])
    # The dataframe consists of 3 MiB of NumPy NDArrays.
    assert_object_size_gt(obj_ref, BIG_OBJ_SIZE)


def test_get_local_locations_generator(ray_start_regular):
    """
    caller and callee are in the same node.
    callee is a streaming generator.
    """
    for obj_ref in gen_big_objects.remote(3, 10):
        # No need to ray.wait, the object ref must have been ready before it's yielded.
        # The dataframe consists of 3 MiB of NumPy NDArrays.
        assert_object_size_gt(obj_ref, BIG_OBJ_SIZE)


def test_get_local_locations_multi_nodes(ray_start_cluster):
    """
    caller and callee are in different nodes.
    callee is a regular ray task.
    """
    cluster = ray_start_cluster
    # head node
    head_node = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    head_node_id = head_node.node_id
    ray.init(cluster.address)
    # add 1 worker node
    worker_node = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    worker_node_id = worker_node.node_id
    cluster.wait_for_nodes()

    @ray.remote
    def caller():
        obj_ref = gen_big_object.options(
            label_selector={ray._raylet.RAY_NODE_ID_KEY: worker_node_id}
        ).remote(3)
        ray.wait([obj_ref])
        # The dataframe consists of 3 MiB of NumPy NDArrays.
        assert_object_size_gt(obj_ref, BIG_OBJ_SIZE)

    ray.get(
        caller.options(
            label_selector={ray._raylet.RAY_NODE_ID_KEY: head_node_id}
        ).remote()
    )


def test_get_local_locations_generator_multi_nodes(ray_start_cluster):
    """
    caller and callee are in different nodes.
    callee is a streaming generator.
    """
    cluster = ray_start_cluster
    # head node
    head_node = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    head_node_id = head_node.node_id
    ray.init(cluster.address)
    # add 1 worker node
    worker_node = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    worker_node_id = worker_node.node_id
    cluster.wait_for_nodes()

    @ray.remote
    def caller():
        gen = gen_big_objects.options(
            label_selector={ray._raylet.RAY_NODE_ID_KEY: worker_node_id}
        ).remote(3, 10)
        for obj_ref in gen:
            # No need to ray.wait, the object ref must have been ready before it's
            # yielded.
            assert_object_size_gt(obj_ref, BIG_OBJ_SIZE)

    ray.get(
        caller.options(
            label_selector={ray._raylet.RAY_NODE_ID_KEY: head_node_id}
        ).remote()
    )


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))