import asyncio
import concurrent.futures
import sys
import time

import numpy as np
import pytest

import ray
from ray._common.test_utils import wait_for_condition
from ray._common.utils import (
    get_or_create_event_loop,
)


@pytest.fixture
def init():
    ray.init(num_cpus=4)
    get_or_create_event_loop().set_debug(False)
    yield
    ray.shutdown()


def gen_tasks(time_scale=0.1):
    @ray.remote
    def f(n):
        time.sleep(n * time_scale)
        return n, np.zeros(1024 * 1024, dtype=np.uint8)

    return [f.remote(i) for i in range(5)]


def test_simple(init):
    @ray.remote
    def f():
        time.sleep(1)
        return np.zeros(1024 * 1024, dtype=np.uint8)

    future = asyncio.ensure_future(f.remote())
    result = get_or_create_event_loop().run_until_complete(future)
    assert isinstance(result, np.ndarray)


def test_gather(init):
    loop = get_or_create_event_loop()
    tasks = gen_tasks()
    results = loop.run_until_complete(asyncio.gather(*tasks))
    assert all(a[0] == b[0] for a, b in zip(results, ray.get(tasks)))


def test_wait(init):
    loop = get_or_create_event_loop()
    tasks = gen_tasks()
    futures = [asyncio.ensure_future(obj_ref) for obj_ref in tasks]
    results, _ = loop.run_until_complete(asyncio.wait(futures))
    assert set(results) == set(futures)


def test_wait_timeout(init):
    loop = get_or_create_event_loop()
    tasks = gen_tasks(10)
    futures = [asyncio.ensure_future(obj_ref) for obj_ref in tasks]
    fut = asyncio.wait(futures, timeout=1)
    results, _ = loop.run_until_complete(fut)
    assert list(results)[0] == futures[0]


def test_gather_mixup(init):
    loop = get_or_create_event_loop()

    @ray.remote
    def f(n):
        time.sleep(n * 0.1)
        return n, np.zeros(1024 * 1024, dtype=np.uint8)

    async def g(n):
        await asyncio.sleep(n * 0.1)
        return n, np.zeros(1024 * 1024, dtype=np.uint8)

    tasks = [
        f.remote(1),
        g(2),
        f.remote(3),
        g(4),
    ]
    results = loop.run_until_complete(asyncio.gather(*tasks))
    assert [result[0] for result in results] == [1, 2, 3, 4]


def test_wait_mixup(init):
    loop = get_or_create_event_loop()

    @ray.remote
    def f(n):
        time.sleep(n)
        return n, np.zeros(1024 * 1024, dtype=np.uint8)

    def g(n):
        async def _g(_n):
            await asyncio.sleep(_n)
            return _n

        return asyncio.ensure_future(_g(n))

    tasks = [
        asyncio.ensure_future(f.remote(0.1)),
        g(7),
        asyncio.ensure_future(f.remote(5)),
        g(2),
    ]
    ready, _ = loop.run_until_complete(asyncio.wait(tasks, timeout=4))
    assert set(ready) == {tasks[0], tasks[-1]}


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "ray_start_regular",
    [
        {
            "object_store_memory": 100 * 1024 * 1024,
        }
    ],
    indirect=True,
)
async def test_garbage_collection(ray_start_regular):
    # This is a regression test for
    # https://github.com/ray-project/ray/issues/9134

    @ray.remote
    def f():
        return np.zeros(20 * 1024 * 1024, dtype=np.uint8)

    for _ in range(10):
        await f.remote()
    for _ in range(10):
        put_id = ray.put(np.zeros(20 * 1024 * 1024, dtype=np.uint8))
        await put_id


def test_concurrent_future(init):
    ref = ray.put(1)

    fut = ref.future()
    assert isinstance(fut, concurrent.futures.Future)
    wait_for_condition(lambda: fut.done())

    global_result = None

    def cb(fut):
        nonlocal global_result
        global_result = fut.result()

    fut.add_done_callback(cb)
    assert global_result == 1
    assert fut.result() == 1


def test_concurrent_future_many(init):
    @ray.remote
    def task(i):
        return i

    refs = [task.remote(i) for i in range(100)]
    futs = [ref.future() for ref in refs]
    result = set()

    for fut in concurrent.futures.as_completed(futs):
        assert fut.done()
        result.add(fut.result())
    assert result == set(range(100))


if __name__ == "__main__":

    sys.exit(pytest.main(["-sv", __file__]))