import sys

import pytest

import ray


def test_multi_raise_exception_on_non_client(shutdown_only):
    with pytest.raises(
        RuntimeError,
        match=("Do not pass the `allow_multiple` to `ray.init` to fix the issue."),
    ):
        ray.init(allow_multiple=True)


@pytest.mark.skipif(
    sys.platform == "win32", reason="PSUtil does not work the same on windows."
)
@pytest.mark.parametrize(
    "call_ray_start",
    ["ray start --head --ray-client-server-port 25001 --port 0"],
    indirect=True,
)
def test_multi_cli_basic(call_ray_start):
    ray.init("ray://localhost:25001")
    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)
    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)
    with cli1:
        a = ray.put(10)

    with cli2:
        b = ray.put(20)

    # TODO better error message.
    # Right now, it's EOFError actually
    with pytest.raises(Exception):
        ray.get(a)

    with pytest.raises(Exception), cli2:
        ray.get(a)

    with pytest.raises(Exception), cli1:
        ray.get(b)

    c = ray.put(30)

    with cli1:
        assert 10 == ray.get(a)

    with cli2:
        assert 20 == ray.get(b)

    with pytest.raises(Exception), cli1:
        ray.get(c)

    with pytest.raises(Exception), cli2:
        ray.get(c)


@pytest.mark.skipif(
    sys.platform == "win32", reason="PSUtil does not work the same on windows."
)
@pytest.mark.parametrize(
    "call_ray_start",
    ["ray start --head --ray-client-server-port 25001 --port 0"],
    indirect=True,
)
def test_multi_cli_init(call_ray_start):
    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)  # noqa
    with pytest.raises(
        ValueError,
        match="The client has already connected to the cluster "
        "with allow_multiple=True. Please set allow_multiple=True"
        " to proceed",
    ):
        ray.init("ray://localhost:25001")
    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)  # noqa

    cli1.disconnect()
    cli2.disconnect()

    ray.init("ray://localhost:25001")
    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)  # noqa


@pytest.mark.skipif(
    sys.platform == "win32", reason="PSUtil does not work the same on windows."
)
@pytest.mark.parametrize(
    "call_ray_start",
    ["ray start --head --ray-client-server-port 25001 --port 0"],
    indirect=True,
)
def test_multi_cli_func(call_ray_start):
    @ray.remote
    def hello():
        return "world"

    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)
    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)

    # TODO better error message.
    # Right now, it's EOFError actually
    with pytest.raises(Exception):
        ray.get(hello.remote())

    with cli1:
        o1 = hello.remote()
        assert "world" == ray.get(o1)

    with cli2:
        o2 = hello.remote()
        assert "world" == ray.get(o2)

    with pytest.raises(Exception), cli1:
        ray.get(o2)

    with pytest.raises(Exception), cli2:
        ray.get(o1)


@pytest.mark.skipif(
    sys.platform == "win32", reason="PSUtil does not work the same on windows."
)
@pytest.mark.parametrize(
    "call_ray_start",
    ["ray start --head --ray-client-server-port 25001 --port 0"],
    indirect=True,
)
def test_multi_cli_actor(call_ray_start):
    @ray.remote
    class Actor:
        def __init__(self, v):
            self.v = v

        def double(self):
            return self.v * 2

    cli1 = ray.init("ray://localhost:25001", allow_multiple=True)
    cli2 = ray.init("ray://localhost:25001", allow_multiple=True)

    # TODO better error message.
    # Right now, it's EOFError actually
    with pytest.raises(Exception):
        a = Actor.remote(10)
        ray.get(a.double.remote())

    with cli1:
        a1 = Actor.remote(10)
        o1 = a1.double.remote()
        assert 20 == ray.get(o1)

    with cli2:
        a2 = Actor.remote(20)
        o2 = a2.double.remote()
        assert 40 == ray.get(o2)

    with pytest.raises(Exception), cli1:
        ray.get(a2.double.remote())

    with pytest.raises(Exception), cli1:
        ray.get(o2)

    with pytest.raises(Exception), cli2:
        ray.get(a1.double.remote())

    with pytest.raises(Exception), cli2:
        ray.get(o1)


@pytest.mark.skipif(
    sys.platform == "win32", reason="PSUtil does not work the same on windows."
)
@pytest.mark.parametrize(
    "call_ray_start",
    ["ray start --head --ray-client-server-port 25001 --port 0"],
    indirect=True,
)
def test_multi_cli_threading(call_ray_start):
    import threading

    b = threading.Barrier(2)
    ret = [None, None]

    def get(idx):
        cli = ray.init("ray://localhost:25001", allow_multiple=True)
        with cli:
            a = ray.put(idx)
            b.wait()
            v = ray.get(a)
            assert idx == v
            b.wait()
            ret[idx] = v

    t1 = threading.Thread(target=get, args=(0,))
    t2 = threading.Thread(target=get, args=(1,))
    t1.start()
    t2.start()
    t1.join()
    t2.join()
    assert ret == [0, 1]


if __name__ == "__main__":
    # The following should be removed after
    # https://github.com/ray-project/ray/issues/20355
    # is fixed.
    from ray._private import auto_init_hook

    auto_init_hook.enable_auto_connect = False
    sys.exit(pytest.main(["-sv", __file__]))