"""Reference counting tests that require their own custom fixture.

The other reference counting tests use a shared Ray instance across the test module
to reduce overheads & overall test runtime.
"""
# coding: utf-8
import logging
import platform
import random
import sys
import time

import numpy as np
import pytest

import ray
import ray.cluster_utils
from ray._common.test_utils import (
    SignalActor,
    fetch_prometheus_metrics,
    wait_for_condition,
)
from ray._private.internal_api import memory_summary

logger = logging.getLogger(__name__)


def _fill_object_store_and_get(obj, succeed=True, object_MiB=20, num_objects=5):
    for _ in range(num_objects):
        ray.put(np.zeros(object_MiB * 1024 * 1024, dtype=np.uint8))

    if type(obj) is bytes:
        obj = ray.ObjectRef(obj)

    if succeed:
        wait_for_condition(
            lambda: ray._private.worker.global_worker.core_worker.object_exists(obj)
        )
    else:
        wait_for_condition(
            lambda: not ray._private.worker.global_worker.core_worker.object_exists(obj)
        )


@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
def test_object_unpin(ray_start_cluster):
    nodes = []
    cluster = ray_start_cluster
    head_node = cluster.add_node(
        num_cpus=0,
        object_store_memory=100 * 1024 * 1024,
        _system_config={
            "subscriber_timeout_ms": 100,
            "health_check_initial_delay_ms": 0,
            "health_check_period_ms": 1000,
            "health_check_failure_threshold": 5,
        },
    )
    ray.init(address=cluster.address)

    # Add worker nodes.
    for i in range(2):
        nodes.append(
            cluster.add_node(
                num_cpus=1,
                resources={f"node_{i}": 1},
                object_store_memory=100 * 1024 * 1024,
            )
        )
    cluster.wait_for_nodes()

    one_mb_array = np.ones(1 * 1024 * 1024, dtype=np.uint8)
    ten_mb_array = np.ones(10 * 1024 * 1024, dtype=np.uint8)

    @ray.remote
    class ObjectsHolder:
        def __init__(self):
            self.ten_mb_objs = []
            self.one_mb_objs = []

        def put_10_mb(self):
            self.ten_mb_objs.append(ray.put(ten_mb_array))

        def put_1_mb(self):
            self.one_mb_objs.append(ray.put(one_mb_array))

        def pop_10_mb(self):
            if len(self.ten_mb_objs) == 0:
                return False
            self.ten_mb_objs.pop()
            return True

        def pop_1_mb(self):
            if len(self.one_mb_objs) == 0:
                return False
            self.one_mb_objs.pop()
            return True

    # Head node contains 11MB of data.
    one_mb_arrays = []
    ten_mb_arrays = []

    one_mb_arrays.append(ray.put(one_mb_array))
    ten_mb_arrays.append(ray.put(ten_mb_array))

    def check_memory(mb):
        return f"Plasma memory usage {mb} MiB" in memory_summary(
            address=head_node.address, stats_only=True
        )

    def wait_until_node_dead(node):
        for n in ray.nodes():
            if n["ObjectStoreSocketName"] == node.address_info["object_store_address"]:
                return not n["Alive"]
        return False

    wait_for_condition(lambda: check_memory(11))

    # Pop one mb array and see if it works.
    one_mb_arrays.pop()
    wait_for_condition(lambda: check_memory(10))

    # Pop 10 MB.
    ten_mb_arrays.pop()
    wait_for_condition(lambda: check_memory(0))

    # Put 11 MB for each actor.
    # actor 1: 1MB + 10MB
    # actor 2: 1MB + 10MB
    actor_on_node_1 = ObjectsHolder.options(resources={"node_0": 1}).remote()
    actor_on_node_2 = ObjectsHolder.options(resources={"node_1": 1}).remote()
    ray.get(actor_on_node_1.put_1_mb.remote())
    ray.get(actor_on_node_1.put_10_mb.remote())
    ray.get(actor_on_node_2.put_1_mb.remote())
    ray.get(actor_on_node_2.put_10_mb.remote())
    wait_for_condition(lambda: check_memory(22))

    # actor 1: 10MB
    # actor 2: 1MB
    ray.get(actor_on_node_1.pop_1_mb.remote())
    ray.get(actor_on_node_2.pop_10_mb.remote())
    wait_for_condition(lambda: check_memory(11))

    # The second node is dead, and actor 2 is dead.
    cluster.remove_node(nodes[1], allow_graceful=False)
    wait_for_condition(lambda: wait_until_node_dead(nodes[1]))
    wait_for_condition(lambda: check_memory(10))

    # The first actor is dead, so object should be GC'ed.
    ray.kill(actor_on_node_1)
    wait_for_condition(lambda: check_memory(0))


@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
def test_object_unpin_stress(ray_start_cluster):
    nodes = []
    cluster = ray_start_cluster
    cluster.add_node(
        num_cpus=1, resources={"head": 1}, object_store_memory=1000 * 1024 * 1024
    )
    ray.init(address=cluster.address)

    # Add worker nodes.
    for i in range(2):
        nodes.append(
            cluster.add_node(
                num_cpus=1,
                resources={f"node_{i}": 1},
                object_store_memory=1000 * 1024 * 1024,
            )
        )
    cluster.wait_for_nodes()

    one_mb_array = np.ones(1 * 1024 * 1024, dtype=np.uint8)
    ten_mb_array = np.ones(10 * 1024 * 1024, dtype=np.uint8)

    @ray.remote
    class ObjectsHolder:
        def __init__(self):
            self.ten_mb_objs = []
            self.one_mb_objs = []

        def put_10_mb(self):
            self.ten_mb_objs.append(ray.put(ten_mb_array))

        def put_1_mb(self):
            self.one_mb_objs.append(ray.put(one_mb_array))

        def pop_10_mb(self):
            if len(self.ten_mb_objs) == 0:
                return False
            self.ten_mb_objs.pop()
            return True

        def pop_1_mb(self):
            if len(self.one_mb_objs) == 0:
                return False
            self.one_mb_objs.pop()
            return True

        def get_obj_size(self):
            return len(self.ten_mb_objs) * 10 + len(self.one_mb_objs)

    actor_on_node_1 = ObjectsHolder.options(resources={"node_0": 1}).remote()
    actor_on_node_2 = ObjectsHolder.options(resources={"node_1": 1}).remote()
    actor_on_head_node = ObjectsHolder.options(resources={"head": 1}).remote()

    ray.get(actor_on_node_1.get_obj_size.remote())
    ray.get(actor_on_node_2.get_obj_size.remote())
    ray.get(actor_on_head_node.get_obj_size.remote())

    def random_ops(actors):
        r = random.random()
        for actor in actors:
            if r <= 0.25:
                actor.put_10_mb.remote()
            elif r <= 0.5:
                actor.put_1_mb.remote()
            elif r <= 0.75:
                actor.pop_10_mb.remote()
            else:
                actor.pop_1_mb.remote()

    total_iter = 15
    for _ in range(total_iter):
        random_ops([actor_on_node_1, actor_on_node_2, actor_on_head_node])

    # Simulate node dead.
    cluster.remove_node(nodes[1])
    for _ in range(total_iter):
        random_ops([actor_on_node_1, actor_on_head_node])

    total_size = sum(
        [
            ray.get(actor_on_node_1.get_obj_size.remote()),
            ray.get(actor_on_head_node.get_obj_size.remote()),
        ]
    )

    wait_for_condition(
        lambda: (
            (f"Plasma memory usage {total_size} MiB") in memory_summary(stats_only=True)
        )
    )


@pytest.mark.parametrize("inline_args", [True, False])
def test_inlined_nested_refs(ray_start_cluster, inline_args):
    cluster = ray_start_cluster
    config = {}
    if not inline_args:
        config["max_direct_call_object_size"] = 0
    cluster.add_node(
        num_cpus=2, object_store_memory=100 * 1024 * 1024, _system_config=config
    )
    ray.init(address=cluster.address)

    @ray.remote
    class Actor:
        def __init__(self):
            return

        def nested(self):
            return ray.put("x")

    @ray.remote
    def nested_nested(a):
        return a.nested.remote()

    @ray.remote
    def foo(ref):
        time.sleep(1)
        return ray.get(ref)

    a = Actor.remote()
    nested_nested_ref = nested_nested.remote(a)
    # We get nested_ref's value directly from its owner.
    nested_ref = ray.get(nested_nested_ref)

    del nested_nested_ref
    x = foo.remote(nested_ref)
    del nested_ref
    ray.get(x)


# https://github.com/ray-project/ray/issues/17553
@pytest.mark.parametrize("inline_args", [True, False])
def test_return_nested_ids(shutdown_only, inline_args):
    config = dict()
    if inline_args:
        config["max_direct_call_object_size"] = 100 * 1024 * 1024
    else:
        config["max_direct_call_object_size"] = 0
    ray.init(object_store_memory=100 * 1024 * 1024, _system_config=config)

    class Nested:
        def __init__(self, blocks):
            self._blocks = blocks

    @ray.remote
    def echo(fn):
        return fn()

    @ray.remote
    def create_nested():
        refs = [ray.put(np.random.random(1024 * 1024)) for _ in range(10)]
        return Nested(refs)

    @ray.remote
    def test():
        ref = create_nested.remote()
        result1 = ray.get(ref)
        del ref
        result = echo.remote(lambda: result1)  # noqa
        del result1

        time.sleep(5)
        block = ray.get(result)._blocks[0]
        print(ray.get(block))

    ray.get(test.remote())


def _check_refcounts(expected):
    actual = ray._private.worker.global_worker.core_worker.get_all_reference_counts()
    assert len(expected) == len(actual)
    for object_ref, (local, submitted) in expected.items():
        hex_id = object_ref.hex().encode("ascii")
        assert hex_id in actual
        assert local == actual[hex_id]["local"]
        assert submitted == actual[hex_id]["submitted"]


def test_out_of_band_serialized_object_ref(ray_start_regular):
    assert (
        len(ray._private.worker.global_worker.core_worker.get_all_reference_counts())
        == 0
    )
    obj_ref = ray.put("hello")
    _check_refcounts({obj_ref: (1, 0)})
    obj_ref_str = ray.cloudpickle.dumps(obj_ref)
    _check_refcounts({obj_ref: (2, 0)})
    del obj_ref
    assert (
        len(ray._private.worker.global_worker.core_worker.get_all_reference_counts())
        == 1
    )
    assert ray.get(ray.cloudpickle.loads(obj_ref_str)) == "hello"


def test_captured_object_ref(ray_start_regular):
    captured_id = ray.put(np.zeros(1024, dtype=np.uint8))

    @ray.remote
    def f(signal):
        ray.get(signal.wait.remote())
        ray.get(captured_id)  # noqa: F821

    signal = SignalActor.remote()
    obj_ref = f.remote(signal)

    # Delete local references.
    del f
    del captured_id

    # Test that the captured object ref is pinned despite having no local
    # references.
    ray.get(signal.send.remote())
    _fill_object_store_and_get(obj_ref)

    captured_id = ray.put(np.zeros(1024, dtype=np.uint8))

    @ray.remote
    class Actor:
        def get(self, signal):
            ray.get(signal.wait.remote())
            ray.get(captured_id)  # noqa: F821

    signal = SignalActor.remote()
    actor = Actor.remote()
    obj_ref = actor.get.remote(signal)

    # Delete local references.
    del Actor
    del captured_id

    # Test that the captured object ref is pinned despite having no local
    # references.
    ray.get(signal.send.remote())
    _fill_object_store_and_get(obj_ref)


def test_borrowed_id_failure_while_pulling(ray_start_cluster):
    """The driver creates an object and passes the ref to actor A via an actor
    task. That task passes the ref on to B, then A kills itself before
    finishing the task, so the task never reports the borrower B to the
    driver. The driver can therefore erase the ref while B is still pulling
    the object, and B's get must then fail promptly instead of hanging.
    """
    cluster = ray_start_cluster
    cluster.add_node(
        num_cpus=1,
        resources={"head_node": 1},
        object_store_memory=100 * 1024 * 1024,
        _system_config={
            "testing_asio_delay_us": (
                "ObjectManagerService.grpc_server.Pull=5000000000:5000000000"
            ),
            "metrics_report_interval_ms": 200,
        },
    )
    ray.init(address=cluster.address)
    cluster.add_node(
        num_cpus=1,
        resources={"worker_node": 1},
        object_store_memory=100 * 1024 * 1024,
    )
    cluster.wait_for_nodes()

    @ray.remote(resources={"head_node": 1})
    class A:
        def pass_ref(self, ref, b):
            ray.get(b.receive_ref.remote(ref))
            sys.exit(-1)

    @ray.remote(resources={"worker_node": 1})
    class B:
        def __init__(self):
            self.ref = None

        def receive_ref(self, ref):
            self.ref = ref[0]

        def resolve_ref(self):
            with pytest.raises(ray.exceptions.ObjectLostError):
                ray.get(self.ref)
            return True

        def ping(self):
            return

    a = A.remote()
    b = B.remote()
    ray.get(b.ping.remote())

    obj = ray.put(np.zeros(1024 * 1024, dtype=np.uint8))
    with pytest.raises(ray.exceptions.RayActorError):
        ray.get(a.pass_ref.remote([obj], b))

    resolved = b.resolve_ref.remote()

    def pull_in_flight():
        (worker_node,) = [n for n in ray.nodes() if "worker_node" in n["Resources"]]
        address = (
            f"{worker_node['NodeManagerAddress']}:{worker_node['MetricsExportPort']}"
        )
        samples = fetch_prometheus_metrics([address]).get(
            "ray_pull_manager_usage_bytes", []
        )
        return any(
            s.labels.get("Type") == "BeingPulled" and s.value > 0 for s in samples
        )

    # make sure to only del the last ref to let ref count go to 0 after actor B's raylet starts pulling that object.
    wait_for_condition(pull_in_flight, timeout=30)
    del obj
    assert ray.get(resolved, timeout=30)


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