from concurrent.futures import Future

import pytest
from ray.util.client import _ClientContext
from ray.util.client.common import ClientActorRef, ClientObjectRef
from ray.util.client.ray_client_helpers import ray_start_client_server
from ray.util.client.ray_client_helpers import (
    ray_start_client_server_pair,
    ray_start_cluster_client_server_pair,
)
from ray._private.test_utils import wait_for_condition, object_memory_usage
import ray as real_ray
from ray.core.generated.gcs_pb2 import ActorTableData
from ray._raylet import ActorID, ObjectRef


def test_client_object_ref_basics(ray_start_regular):
    with ray_start_client_server_pair() as pair:
        ray, server = pair
        ref = ray.put("Hello World")
        # Make sure ClientObjectRef is a subclass of ObjectRef
        assert isinstance(ref, ClientObjectRef)
        assert isinstance(ref, ObjectRef)

        # Invalid ref format.
        with pytest.raises(Exception):
            ClientObjectRef(b"\0")

        obj_id = b"\0" * 28
        fut = Future()
        fut.set_result(obj_id)
        server_ref = ObjectRef(obj_id)
        for client_ref in [ClientObjectRef(obj_id), ClientObjectRef(fut)]:
            client_members = set(client_ref.__dir__())
            server_members = set(server_ref.__dir__())
            client_members = {m for m in client_ref.__dir__() if not m.startswith("_")}
            server_members = {m for m in server_ref.__dir__() if not m.startswith("_")}
            assert client_members.difference(server_members) == {"id"}
            assert server_members.difference(client_members) == set()

            # Test __eq__()
            assert client_ref == ClientObjectRef(obj_id)
            assert client_ref != ref
            assert client_ref != server_ref

            # Test other methods
            assert client_ref.__repr__() == f"ClientObjectRef({obj_id.hex()})"
            assert client_ref.binary() == obj_id
            assert client_ref.hex() == obj_id.hex()
            assert not client_ref.is_nil()
            assert client_ref.task_id() == server_ref.task_id()
            assert client_ref.job_id() == server_ref.job_id()


def test_client_actor_ref_basics(ray_start_regular):
    with ray_start_client_server_pair() as pair:
        ray, server = pair

        @ray.remote
        class Counter:
            def __init__(self):
                self.acc = 0

            def inc(self):