# coding: utf-8
import asyncio
import sys
import threading
import time
from typing import Any, Tuple

import pytest

import ray
from ray._common.test_utils import (
    SignalActor,
    run_string_as_driver,
)
from ray._common.utils import get_or_create_event_loop


# This tests the methods are executed in the correct eventloop.
def test_basic(ray_start_regular_shared):
    @ray.remote(concurrency_groups={"io": 2, "compute": 4})
    class AsyncActor:
        def __init__(self):
            self.eventloop_f1 = None
            self.eventloop_f2 = None
            self.eventloop_f3 = None
            self.eventloop_f4 = None
            self.default_eventloop = get_or_create_event_loop()

        @ray.method(concurrency_group="io")
        async def f1(self):
            self.eventloop_f1 = get_or_create_event_loop()
            return threading.current_thread().ident

        @ray.method(concurrency_group="io")
        def f2(self):
            self.eventloop_f2 = get_or_create_event_loop()
            return threading.current_thread().ident

        @ray.method(concurrency_group="compute")
        def f3(self):
            self.eventloop_f3 = get_or_create_event_loop()
            return threading.current_thread().ident

        @ray.method(concurrency_group="compute")
        def f4(self):
            self.eventloop_f4 = get_or_create_event_loop()
            return threading.current_thread().ident

        def f5(self):
            # If this method is executed in default eventloop.
            assert get_or_create_event_loop() == self.default_eventloop
            return threading.current_thread().ident

        @ray.method(concurrency_group="io")
        def do_assert(self):
            if self.eventloop_f1 != self.eventloop_f2:
                return False
            if self.eventloop_f3 != self.eventloop_f4:
                return False
            if self.eventloop_f1 == self.eventloop_f3:
                return False
            if self.eventloop_f1 == self.eventloop_f4:
                return False
            return True

    ###############################################
    a = AsyncActor.remote()
    f1_thread_id = ray.get(a.f1.remote())  # executed in the "io" group.
    f2_thread_id = ray.get(a.f2.remote())  # executed in the "io" group.
    f3_thread_id = ray.get(a.f3.remote())  # executed in the "compute" group.
    f4_thread_id = ray.get(a.f4.remote())  # executed in the "compute" group.

    assert f1_thread_id == f2_thread_id
    assert f3_thread_id == f4_thread_id
    assert f1_thread_id != f3_thread_id

    assert ray.get(a.do_assert.remote())

    assert ray.get(a.f5.remote())  # executed in the default group.

    # It also has the ability to specify it at runtime.
    # This task will be invoked in the `compute` thread pool.
    result = ray.get(a.f2.options(concurrency_group="compute").remote())
    assert result == f3_thread_id


# The case tests that the asyncio count down works well in one concurrency
# group.
def test_async_methods_in_concurrency_group(ray_start_regular_shared):
    @ray.remote(concurrency_groups={"async": 3})
    class AsyncBatcher:
        def __init__(self):
            self.batch = []
            self.event = None

        @ray.method(concurrency_group="async")
        def init_event(self):
            self.event = asyncio.Event()
            return True

        @ray.method(concurrency_group="async")
        async def add(self, x):
            self.batch.append(x)
            if len(self.batch) >= 3:
                self.event.set()
            else:
                await self.event.wait()
            return sorted(self.batch)

    a = AsyncBatcher.remote()
    ray.get(a.init_event.remote())

    x1 = a.add.remote(1)
    x2 = a.add.remote(2)
    x3 = a.add.remote(3)
    r1 = ray.get(x1)
    r2 = ray.get(x2)
    r3 = ray.get(x3)
    assert r1 == [1, 2, 3]
    assert r1 == r2 == r3


# This case tests that if blocking task in default group blocks
# tasks in other groups.
# See https://github.com/ray-project/ray/issues/20475
def test_default_concurrency_group_does_not_block_others(ray_start_regular_shared):
    @ray.remote(concurrency_groups={"my_group": 1})
    class AsyncActor:
        def __init__(self):
            pass

        async def f1(self):
            time.sleep(10000)
            return "never return"

        @ray.method(concurrency_group="my_group")
        def f2(self):
            return "ok"

    async_actor = AsyncActor.remote()
    async_actor.f1.remote()
    assert "ok" == ray.get(async_actor.f2.remote())


# This case tests that a blocking group doesn't blocks
# tasks in other groups.
# See https://github.com/ray-project/ray/issues/19593
def test_blocking_group_does_not_block_others(ray_start_regular_shared):
    @ray.remote(concurrency_groups={"group1": 1, "group2": 1})
    class AsyncActor:
        def __init__(self):
            pass

        @ray.method(concurrency_group="group1")
        async def f1(self):
            time.sleep(10000)
            return "never return"

        @ray.method(concurrency_group="group2")
        def f2(self):
            return "ok"

    async_actor = AsyncActor.remote()
    # Execute f1 twice for blocking the group1.
    obj_0 = async_actor.f1.remote()
    obj_1 = async_actor.f1.remote()
    # Wait a while to make sure f2 is scheduled after f1.
    ray.wait([obj_0, obj_1], timeout=5)
    # f2 should work well even if group1 is blocking.
    assert "ok" == ray.get(async_actor.f2.remote())


def test_system_concurrency_group(ray_start_regular_shared):
    @ray.remote
    class NormalActor:
        def block_forever(self):
            time.sleep(9999)
            return "never"

        def ping(self):
            return "pong"

    n = NormalActor.remote()
    n.block_forever.options(concurrency_group="_ray_system").remote()
    print(ray.get(n.ping.remote()))


@ray.remote(concurrency_groups={"io": 1, "compute": 1})
class Actor:
    def __init__(self):
        self._thread_local_data = threading.local()

    def set_thread_local(self, value: Any) -> int:
        self._thread_local_data.value = value
        return threading.current_thread().ident

    def get_thread_local(self) -> Tuple[Any, int]:
        return self._thread_local_data.value, threading.current_thread().ident


class TestThreadingLocalData:
    """
    This test verifies that synchronous tasks can access thread-local data that
    was set by previous synchronous tasks when the concurrency group has only
    one thread. For concurrency groups with multiple threads, it doesn't promise
    access to the same thread-local data because Ray currently doesn't expose APIs
    for users to specify which thread the task will be scheduled on in the same
    concurrency group.
    """

    def test_tasks_on_default_executor(self, ray_start_regular_shared):
        a = Actor.remote()
        tid_1 = ray.get(a.set_thread_local.remote("f1"))
        value, tid_2 = ray.get(a.get_thread_local.remote())
        assert tid_1 == tid_2
        assert value == "f1"

    def test_tasks_on_specific_executor(self, ray_start_regular_shared):
        a = Actor.remote()
        tid_1 = ray.get(a.set_thread_local.options(concurrency_group="io").remote("f1"))
        value, tid_2 = ray.get(
            a.get_thread_local.options(concurrency_group="io").remote()
        )
        assert tid_1 == tid_2
        assert value == "f1"

    def test_tasks_on_different_executors(self, ray_start_regular_shared):
        a = Actor.remote()
        tid_1 = ray.get(a.set_thread_local.options(concurrency_group="io").remote("f1"))
        tid_3 = ray.get(
            a.set_thread_local.options(concurrency_group="compute").remote("f2")
        )
        value, tid_2 = ray.get(
            a.get_thread_local.options(concurrency_group="io").remote()
        )
        assert tid_1 == tid_2
        assert value == "f1"

        value, tid_4 = ray.get(
            a.get_thread_local.options(concurrency_group="compute").remote()
        )
        assert tid_3 == tid_4
        assert value == "f2"


def test_multiple_threads_in_same_group(ray_start_regular_shared):
    """
    This test verifies that all threads in the same concurrency group are still
    alive from the Python interpreter's perspective even if Ray tasks have finished, so that
    thread-local data will not be garbage collected.
    """

    @ray.remote
    class Actor:
        def __init__(self, signal: SignalActor, max_concurrency: int):
            self._thread_local_data = threading.local()
            self.signal = signal
            self.thread_id_to_data = {}
            self.max_concurrency = max_concurrency

        def set_thread_local(self, value: int) -> int:
            # If the thread-local data were garbage collected after the previous
            # task on the same thread finished, `self.data` would be incremented
            # more than once for the same thread.
            assert not hasattr(self._thread_local_data, "value")
            self._thread_local_data.value = value
            self.thread_id_to_data[threading.current_thread().ident] = value
            ray.get(self.signal.wait.remote())

        def check_thread_local_data(self) -> bool:
            assert len(self.thread_id_to_data) == self.max_concurrency
            assert hasattr(self._thread_local_data, "value")
            assert (
                self._thread_local_data.value
                == self.thread_id_to_data[threading.current_thread().ident]
            )
            ray.get(self.signal.wait.remote())

    max_concurrency = 5
    signal = SignalActor.remote()
    a = Actor.options(max_concurrency=max_concurrency).remote(signal, max_concurrency)

    refs = []
    for i in range(max_concurrency):
        refs.append(a.set_thread_local.remote(i))

    ray.get(signal.send.remote())
    ray.get(refs)

    refs = []
    for _ in range(max_concurrency):
        refs.append(a.check_thread_local_data.remote())

    ray.get(signal.send.remote())
    ray.get(refs)


def test_invalid_concurrency_group():
    """Verify that when a concurrency group has max concurrency set to 0,
    an error is raised when the actor is created. This test uses
    `run_string_as_driver` and checks whether the error message appears in the
    driver's stdout. Since the error in the core worker process does not raise
    an exception in the driver process, we need to check the driver process's
    stdout.
    """

    script = """
import ray

ray.init()

@ray.remote(concurrency_groups={"io": 0, "compute": 0})
class A:
    def __init__(self):
        pass

actor = A.remote()
    """

    output = run_string_as_driver(script)
    assert "max_concurrency must be greater than 0" in output


def test_per_group_independent_ordering(ray_start_regular_shared):
    """
    Verify that a blocking task on one concurrency group doesn't affect another
    even on an in-order actor. Sequencing should be independent per concurrency group.
    When group_a has a long-running task, multiple sequential tasks submitted to group_b
    should still be able to complete in order.
    """

    @ray.remote
    def dependency_task(signal):
        ray.get(signal.wait.remote())

    @ray.remote(concurrency_groups={"group_a": 1, "group_b": 1})
    class MyActor:
        @ray.method(concurrency_group="group_a")
        async def slow_a(self, _blocking_ref):
            return "a_done"

        @ray.method(concurrency_group="group_b")
        def fast_b(self, index):
            return f"b_{index}"

    signal = SignalActor.remote()
    actor = MyActor.remote()

    # ref_a will be blocked in dependency resolution
    ref = dependency_task.remote(signal)
    ref_a = actor.slow_a.remote(ref)

    # Tasks on group_b should still be able to finish
    refs_b = [actor.fast_b.remote(i) for i in range(5)]
    results_b = ray.get(refs_b)
    assert results_b == [f"b_{i}" for i in range(5)]

    # Unblock group_a and verify it completes.
    ray.get(signal.send.remote())
    assert ray.get(ref_a) == "a_done"


if __name__ == "__main__":

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