"""Client tests that run their own init (as with init_and_serve) live here"""

import json
import random
import subprocess
import sys
import time
from unittest.mock import patch

import pytest

import ray
import ray.core.generated.ray_client_pb2 as ray_client_pb2
import ray.util.client.server.server as ray_client_server
from ray._private.client_mode_hook import _explicitly_disable_client_mode
from ray._private.ray_logging.logging_config import LoggingConfig
from ray.cloudpickle.compat import pickle
from ray.cluster_utils import cluster_not_supported
from ray.job_config import JobConfig
from ray.util.client import _ClientContext, ray as ray_client
from ray.util.client.worker import prepare_init_request_args


@ray.remote
def hello_world():
    c1 = complex_task.remote(random.randint(1, 10))
    c2 = complex_task.remote(random.randint(1, 10))
    return sum(ray.get([c1, c2]))


@ray.remote
def complex_task(value):
    time.sleep(1)
    return value * 10


@ray.remote
class C:
    def __init__(self, x):
        self.val = x

    def double(self):
        self.val += self.val

    def get(self):
        return self.val


@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
@pytest.fixture
def init_and_serve_lazy(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1, num_gpus=0)
    cluster.wait_for_nodes(1)
    address = cluster.address

    def connect(job_config=None, **ray_init_kwargs):
        ray.init(address=address, job_config=job_config, **ray_init_kwargs)

    server_handle = ray_client_server.serve("localhost", 50051, connect)
    yield server_handle
    ray_client_server.shutdown_with_server(server_handle.grpc_server)


def test_validate_port():
    """Check that ports outside of 1024-65535 are rejected."""
    for port in [1000, 1023, 65536, 700000]:
        with pytest.raises(subprocess.CalledProcessError) as excinfo:
            subprocess.check_output(
                [
                    "ray",
                    "start",
                    "--head",
                    "--num-cpus",
                    "8",
                    "--ray-client-server-port",
                    f"{port}",
                ]
            )
            assert "ValueError" in str(excinfo.traceback)
            assert "65535" in str(excinfo.traceback)


def test_basic_preregister(init_and_serve):
    """Tests conversion of Ray actors and remote functions to client actors
    and client remote functions.

    Checks that the conversion works when disconnecting and reconnecting client
    sessions.
    """
    from ray.util.client import ray

    for _ in range(2):
        ray.connect("localhost:50051")
        val = ray.get(hello_world.remote())
        print(val)
        assert val >= 20
        assert val <= 200
        c = C.remote(3)
        x = c.double.remote()
        y = c.double.remote()
        ray.wait([x, y])
        val = ray.get(c.get.remote())
        assert val == 12
        ray.disconnect()


def test_idempotent_disconnect(init_and_serve):
    from ray.util.client import ray

    ray.disconnect()
    ray.disconnect()
    ray.connect("localhost:50051")
    ray.disconnect()
    ray.disconnect()


def test_public_shutdown_wait_for_processes_in_client_mode(init_and_serve):
    ray_client.connect("localhost:50051")
    try:
        assert ray_client.is_connected()
        # In client mode, wait_for_processes is a no-op
        ray.shutdown(wait_for_processes=True)
        assert not ray_client.is_connected()
    finally:
        _explicitly_disable_client_mode()


def test_num_clients(init_and_serve_lazy):
    # Tests num clients reporting; useful if you want to build an app that
    # load balances clients between Ray client servers.

    def get_job_id(api):
        return api.get_runtime_context().worker.current_job_id

    api1 = _ClientContext()
    info1 = api1.connect("localhost:50051")
    job_id_1 = get_job_id(api1)
    assert info1["num_clients"] == 1, info1
    api2 = _ClientContext()
    info2 = api2.connect("localhost:50051")
    job_id_2 = get_job_id(api2)
    assert info2["num_clients"] == 2, info2

    assert job_id_1 == job_id_2

    # Disconnect the first two clients.
    api1.disconnect()
    api2.disconnect()
    time.sleep(1)

    api3 = _ClientContext()
    info3 = api3.connect("localhost:50051")
    job_id_3 = get_job_id(api3)
    assert info3["num_clients"] == 1, info3
    assert job_id_1 != job_id_3

    # Check info contains ray and python version.
    assert isinstance(info3["ray_version"], str), info3
    assert isinstance(info3["ray_commit"], str), info3
    assert isinstance(info3["python_version"], str), info3
    api3.disconnect()


def test_python_version(init_and_serve):
    server_handle = init_and_serve
    ray = _ClientContext()
    info1 = ray.connect("localhost:50051")
    assert info1["python_version"] == ".".join(
        [str(x) for x in list(sys.version_info)[:3]]
    )
    ray.disconnect()
    time.sleep(1)

    def mock_connection_response():
        return ray_client_pb2.ConnectionInfoResponse(
            num_clients=1,
            python_version="2.7.12",
            ray_version="",
            ray_commit="",
        )

    # inject mock connection function
    server_handle.data_servicer._build_connection_response = mock_connection_response

    ray = _ClientContext()
    with pytest.raises(RuntimeError):
        _ = ray.connect("localhost:50051")

    ray = _ClientContext()
    info3 = ray.connect("localhost:50051", ignore_version=True)
    assert info3["num_clients"] == 1, info3
    ray.disconnect()


@patch("ray.util.client.server.dataservicer.CLIENT_SERVER_MAX_THREADS", 6)
@patch("ray.util.client.server.logservicer.CLIENT_SERVER_MAX_THREADS", 6)
def test_max_clients(init_and_serve):
    # Tests max clients. Should support up to CLIENT_SERVER_MAX_THREADS / 2.
    def get_job_id(api):
        return api.get_runtime_context().worker.current_job_id

    for i in range(3):
        api1 = _ClientContext()
        info1 = api1.connect("localhost:50051")

        assert info1["num_clients"] == i + 1, info1

    with pytest.raises(ConnectionError):
        api = _ClientContext()
        _ = api.connect("localhost:50051")


@pytest.mark.parametrize(
    "call_ray_start",
    ["ray start --head --ray-client-server-port=50056"],
    indirect=True,
)
def test_logging_config_in_client_mode(call_ray_start):
    """logging_config reaches task workers (JobConfig) when using Ray Client."""

    @ray.remote
    def job_logging_encoding():
        w = ray.get_runtime_context().worker
        jlc = w.job_logging_config
        return None if jlc is None else jlc.encoding

    ray.init(
        "ray://localhost:50056",
        logging_config=LoggingConfig(encoding="TEXT", log_level="INFO"),
        log_to_driver=False,
        configure_logging=False,
    )

    assert ray.get(job_logging_encoding.remote()) == "TEXT"


def test_prepare_init_request_args():
    """``prepare_init_request_args`` copies kwargs, JSON-serializes logging config,
    propagates it onto ``JobConfig``, and pickles the job for ``InitRequest``."""
    # no job config, no ray init kwargs
    serialized, kwargs = prepare_init_request_args(None, None)
    assert serialized is None and kwargs == {}

    # invalid logging config type
    with pytest.raises(
        TypeError, match="logging_config must be a dict or LoggingConfig"
    ):
        prepare_init_request_args(JobConfig(), {"logging_config": 42})

    # ray init kwargs contains serialized logging config
    lc = LoggingConfig(encoding="JSON", log_level="DEBUG")
    orig = {"logging_config": lc, "num_cpus": 1}
    job = JobConfig()
    _, out = prepare_init_request_args(job, orig)
    assert isinstance(orig["logging_config"], LoggingConfig)
    assert isinstance(out["logging_config"], dict)
    assert out["logging_config"]["encoding"] == "JSON"
    assert orig["num_cpus"] == 1
    json.dumps(out)
    assert job.py_logging_config is not None
    assert job.py_logging_config.encoding == "JSON"

    # write logging config to JobConfig
    job_dict = JobConfig()
    prepare_init_request_args(
        job_dict,
        {"logging_config": {"encoding": "TEXT", "log_level": "ERROR"}},
    )
    assert job_dict.py_logging_config.encoding == "TEXT"
    assert job_dict.py_logging_config.log_level == "ERROR"

    # don't overwrite existing logging config on JobConfig
    existing = LoggingConfig(encoding="TEXT")
    job_keep = JobConfig()
    job_keep.set_py_logging_config(existing)
    prepare_init_request_args(
        job_keep,
        {"logging_config": LoggingConfig(encoding="JSON")},
    )
    assert job_keep.py_logging_config is existing

    # pickle the job for InitRequest
    job_pickle = JobConfig()
    blob, _ = prepare_init_request_args(job_pickle, {})
    assert isinstance(blob, bytes) and len(blob) > 0
    assert isinstance(pickle.loads(blob), JobConfig)


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