# coding: utf-8

import gc
import logging
import sys
import time

import numpy as np
import pytest

from ray._common.test_utils import wait_for_condition
from ray._private.test_utils import client_test_enabled
from ray._private.worker import _wait_generators_bulk
from ray.exceptions import ObjectRefStreamEndOfStreamError, RayTaskError

if client_test_enabled():
    from ray.util.client import ray
else:
    import ray
    import ray.util.state

logger = logging.getLogger(__name__)


def test_wait(ray_start_regular):
    @ray.remote
    def f(delay):
        time.sleep(delay)
        return

    object_refs = [f.remote(0), f.remote(0), f.remote(0), f.remote(0)]
    ready_ids, remaining_ids = ray.wait(object_refs)
    assert len(ready_ids) == 1
    assert len(remaining_ids) == 3
    ready_ids, remaining_ids = ray.wait(object_refs, num_returns=4)
    assert set(ready_ids) == set(object_refs)
    assert remaining_ids == []

    object_refs = [f.remote(0), f.remote(5)]
    ready_ids, remaining_ids = ray.wait(object_refs, timeout=0.5, num_returns=2)
    assert len(ready_ids) == 1
    assert len(remaining_ids) == 1

    # Verify that calling wait with duplicate object refs throws an
    # exception.
    x = ray.put(1)
    with pytest.raises(Exception):
        ray.wait([x, x])

    # Make sure it is possible to call wait with an empty list.
    ready_ids, remaining_ids = ray.wait([])
    assert ready_ids == []
    assert remaining_ids == []

    # Test semantics of num_returns with no timeout.
    obj_refs = [ray.put(i) for i in range(10)]
    (found, rest) = ray.wait(obj_refs, num_returns=2)
    assert len(found) == 2
    assert len(rest) == 8

    # Verify that incorrect usage raises a TypeError.
    x = ray.put(1)
    with pytest.raises(TypeError):
        ray.wait(x)
    with pytest.raises(TypeError):
        ray.wait(1)
    with pytest.raises(TypeError):
        ray.wait([1])


def test_wait_timing(ray_start_2_cpus):
    @ray.remote
    def f():
        time.sleep(1)

    future = f.remote()

    start = time.time()
    ready, not_ready = ray.wait([future], timeout=0.2)
    assert 0.2 < time.time() - start < 0.3
    assert len(ready) == 0
    assert len(not_ready) == 1


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test_wait_always_fetch_local(monkeypatch, ray_start_cluster):
    monkeypatch.setenv("RAY_scheduler_report_pinned_bytes_only", "false")
    cluster = ray_start_cluster
    head_node = cluster.add_node(num_cpus=0, object_store_memory=300e6)
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(num_cpus=1, object_store_memory=300e6)

    @ray.remote(num_cpus=1)
    def return_large_object():
        # 100mb so will spill on worker, but not once on head
        return np.zeros(100 * 1024 * 1024, dtype=np.uint8)

    @ray.remote(num_cpus=0)
    def small_local_task():
        return 1

    put_on_head = {ray._raylet.RAY_NODE_ID_KEY: head_node.node_id}
    put_on_worker = {ray._raylet.RAY_NODE_ID_KEY: worker_node.node_id}
    x = small_local_task.options(label_selector=put_on_head).remote()
    y = return_large_object.options(label_selector=put_on_worker).remote()
    z = return_large_object.options(label_selector=put_on_worker).remote()

    # will return when tasks are done
    ray.wait([x, y, z], num_returns=3, fetch_local=False)
    assert (
        ray._private.state.available_resources_per_node()[head_node.node_id][
            "object_store_memory"
        ]
        > 250e6
    )

    # x should be immediately available locally, start fetching y and z
    ray.wait([x, y, z], num_returns=1, fetch_local=True)
    assert (
        ray._private.state.available_resources_per_node()[head_node.node_id][
            "object_store_memory"
        ]
        > 250e6
    )

    time.sleep(5)
    # y, z should be pulled here
    assert (
        ray._private.state.available_resources_per_node()[head_node.node_id][
            "object_store_memory"
        ]
        < 150e6
    )


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_fetch_local(monkeypatch, ray_start_cluster):
    monkeypatch.setenv("RAY_scheduler_report_pinned_bytes_only", "false")
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0, object_store_memory=500e6)
    ray.init(address=cluster.address)
    worker_node = cluster.add_node(num_cpus=2, object_store_memory=500e6)

    @ray.remote(num_cpus=1)
    def gen_large_objects():
        # 100mb so the object is stored in plasma.
        yield np.zeros(100 * 1024 * 1024, dtype=np.uint8)
        yield np.ones(100 * 1024 * 1024, dtype=np.uint8)

    put_on_worker = {ray._raylet.RAY_NODE_ID_KEY: worker_node.node_id}
    gen1 = gen_large_objects.options(label_selector=put_on_worker).remote()
    gen2 = gen_large_objects.options(label_selector=put_on_worker).remote()

    ready = _wait_generators_bulk(
        [(gen1, [True, False]), (gen2, [False, True])],
        num_return=2,
        timeout=10,
    )
    assert len(ready) == 2
    assert [gen for gen, _ in ready] == [gen1, gen2]
    assert all(len(refs) == 2 for _, refs in ready)

    assert np.all(ray.get(ready[0][1][0], timeout=0) == 0)
    assert np.all(ray.get(ready[1][1][1], timeout=0) == 1)


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_wait_for_at_most_num_return(ray_start_regular):
    @ray.remote
    def gen(base, delays):
        for i, delay in enumerate(delays):
            time.sleep(delay)
            yield base + i

    gen1 = gen.remote(10, [0, 0, 0])
    gen2 = gen.remote(20, [0, 5])

    ready = _wait_generators_bulk(
        [(gen1, [True, False]), (gen2, [False, True])],
        num_return=1,
        timeout=2,
    )

    assert len(ready) == 1
    ready_gen, refs = ready[0]
    assert ready_gen is gen1
    assert ray.get(refs) == [10, 11]

    # The returned refs are consumed from the stream.
    assert ray.get(next(gen1)) == 12


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_timeout(ray_start_regular):
    @ray.remote(num_cpus=0, max_concurrency=2)
    class Signal:
        def __init__(self):
            self.ready = False

        def wait(self):
            while not self.ready:
                time.sleep(0.01)

        def send(self):
            self.ready = True

    @ray.remote
    def slow_gen(signal):
        ray.get(signal.wait.remote())
        yield 1

    signal = Signal.remote()
    gen = slow_gen.remote(signal)
    assert _wait_generators_bulk([(gen, [False])], timeout=0.01) == []
    ray.get(signal.send.remote())
    ready = _wait_generators_bulk([(gen, [False])], timeout=5)
    assert len(ready) == 1
    ready_gen, refs = ready[0]
    assert ready_gen is gen
    assert ray.get(refs) == [1]


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_validation(ray_start_regular):
    @ray.remote
    def gen():
        yield 1

    gen = gen.remote()

    with pytest.raises(TypeError):
        _wait_generators_bulk({})
    with pytest.raises(TypeError):
        _wait_generators_bulk([(ray.put(1), [False])])
    with pytest.raises(TypeError):
        _wait_generators_bulk([(gen, False)])
    with pytest.raises(ValueError):
        _wait_generators_bulk([(gen, [])])
    with pytest.raises(ValueError):
        _wait_generators_bulk([(gen, [False])], num_return=2)


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__consume_next_ref_n_rejects_unready(ray_start_regular):
    """Consuming before the last requested ref is ready must raise rather than
    silently advancing past (and dropping) the not-yet-produced object."""

    @ray.remote(num_cpus=0, max_concurrency=2)
    class Signal:
        def __init__(self):
            self.ready = False

        def wait(self):
            while not self.ready:
                time.sleep(0.01)

        def send(self):
            self.ready = True

    @ray.remote
    def slow_gen(signal):
        ray.get(signal.wait.remote())
        yield 1

    signal = Signal.remote()
    gen = slow_gen.remote(signal)

    # Peek without waiting: the ref isn't produced yet, so consuming it is rejected.
    gen._get_next_ref_n(1)
    with pytest.raises(ValueError):
        gen._consume_next_ref_n(1)

    # After the value is produced, the same generator consumes normally.
    ray.get(signal.send.remote())
    ready = _wait_generators_bulk([(gen, [False])], timeout=10)
    assert len(ready) == 1
    assert ray.get(ready[0][1]) == [1]


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_after_eof_raise_EndOfStreamError(ray_start_regular):
    @ray.remote
    def empty_gen():
        if False:
            yield 1

    empty = empty_gen.remote()
    ready = _wait_generators_bulk([(empty, [True, True, True])], timeout=1)
    assert len(ready) == 1
    ready_gen, refs = ready[0]
    assert ready_gen is empty
    assert len(set(refs)) == 3
    for ref in refs:
        with pytest.raises(ObjectRefStreamEndOfStreamError):
            ray.get(ref)


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_after_partial_eof(ray_start_regular):
    @ray.remote
    def one_item_gen():
        yield 1

    one_item = one_item_gen.remote()
    ready = _wait_generators_bulk([(one_item, [False, False, False])], timeout=1)
    assert len(ready) == 1
    ready_gen, refs = ready[0]
    assert ready_gen is one_item
    assert len(set(refs)) == 3
    assert ray.get(refs[0]) == 1
    for ref in refs[1:]:
        with pytest.raises(ObjectRefStreamEndOfStreamError):
            ray.get(ref)


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_after_partial_error(ray_start_regular):
    @ray.remote
    def one_item_then_error_gen():
        yield 1
        raise ValueError("expected test error")

    one_item_then_error = one_item_then_error_gen.remote()
    ready = _wait_generators_bulk(
        [(one_item_then_error, [False, False, False])], timeout=1
    )
    assert len(ready) == 1
    ready_gen, refs = ready[0]
    assert ready_gen is one_item_then_error
    assert len(set(refs)) == 3
    assert ray.get(refs[0]) == 1
    with pytest.raises(RayTaskError) as exc_info:
        ray.get(refs[1])
    assert isinstance(exc_info.value.as_instanceof_cause(), ValueError)
    with pytest.raises(ObjectRefStreamEndOfStreamError):
        ray.get(refs[2])


def _assert_no_owned_refs_leak():
    """Wait until the owner holds no live references and the store is empty."""

    def check():
        gc.collect()
        core_worker = ray._private.worker.global_worker.core_worker
        ref_counts = core_worker.get_all_reference_counts()
        for rc in ref_counts.values():
            if rc["local"] != 0 or rc["submitted"] != 0:
                return False
        return core_worker.get_memory_store_size() == 0

    wait_for_condition(check, timeout=30)


@pytest.mark.skipif(client_test_enabled(), reason="util not available with ray client")
def test__wait_generators_bulk_no_ref_leak(ray_start_regular):
    """Draining a generator entirely via _wait_generators_bulk must not leak
    owner-side references for the consumed objects.

    The bulk peek (peek_object_ref_stream_n) hands back ObjectRefs that add their
    own local reference, while the owner-side stream reference taken at peek/report
    time is only released for *unconsumed* refs at stream teardown. This test
    confirms whether consumed refs leave that owner-side reference dangling.
    """

    @ray.remote
    def gen():
        for i in range(3):
            yield i

    g = gen.remote()

    collected = []
    saw_eof = False
    while not saw_eof:
        ready = _wait_generators_bulk([(g, [True, True, True])], timeout=10)
        assert len(ready) == 1
        # Avoid binding the generator object to a local (e.g. via tuple unpacking),
        # which would keep its stream alive and prevent teardown.
        refs = ready[0][1]
        for ref in refs:
            try:
                collected.append(ray.get(ref))
            except ObjectRefStreamEndOfStreamError:
                saw_eof = True
                break

    assert collected == [0, 1, 2]

    # Drop every handle to the generator and its (consumed) objects.
    del g, ready, refs, ref

    _assert_no_owned_refs_leak()


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