import asyncio
import glob
import json
import os
import shutil
import sys
from unittest.mock import patch

import pytest

import ray
from ray._private.test_utils import check_call_ray
from ray.util.tracing.setup_local_tmp_tracing import setup_tracing, spans_dir
from ray.util.tracing.tracing_helper import _enable_tracing

setup_tracing_path = "ray.util.tracing.setup_local_tmp_tracing:setup_tracing"


@pytest.fixture()
def cleanup_dirs():
    """Cleanup temporary spans_dir folder at beginning and end of test."""
    if os.path.exists(spans_dir):
        shutil.rmtree(spans_dir)
    os.makedirs(spans_dir)
    yield
    # Enable tracing only sets up tracing once per driver process.
    # We set ray.__traced__ to False here so that each
    # test will re-set up tracing.
    ray.__traced__ = False
    if os.path.exists(spans_dir):
        shutil.rmtree(spans_dir)


@pytest.fixture()
def ray_start_cli_tracing(scope="function"):
    """Start ray with tracing-startup-hook, and clean up at end of test."""
    check_call_ray(
        ["stop", "--force"],
    )
    check_call_ray(
        ["start", "--head", "--tracing-startup-hook", setup_tracing_path],
    )
    yield
    ray.shutdown()
    check_call_ray(["stop", "--force"])


@pytest.fixture()
def ray_start_cli_predefined_actor_tracing(scope="function"):
    """Start ray with tracing-startup-hook, and clean up at end of test."""
    check_call_ray(
        ["stop", "--force"],
    )
    check_call_ray(
        ["start", "--head", "--tracing-startup-hook", setup_tracing_path],
    )
    yield
    ray.shutdown()
    check_call_ray(["stop", "--force"])


@pytest.fixture()
def ray_start_cli_tracing_with_client(scope="function"):
    """Start ray with tracing startup hook and Ray Client server."""
    check_call_ray(["stop", "--force"])
    check_call_ray(
        [
            "start",
            "--head",
            "--ray-client-server-port=50057",
            "--tracing-startup-hook",
            setup_tracing_path,
        ],
    )
    yield "ray://localhost:50057"
    ray.shutdown()
    check_call_ray(["stop", "--force"])


@pytest.fixture()
def ray_start_init_tracing(scope="function"):
    """Call ray.init with tracing-startup-hook, and clean up at end of test."""
    ray.init(_tracing_startup_hook=setup_tracing_path)
    yield
    ray.shutdown()


def get_span_list():
    """Read span files and return list of span names."""
    span_list = []
    for entry in glob.glob(f"{spans_dir}/**/*.txt", recursive=True):
        with open(entry) as f:
            for line in f.readlines():
                span_list.append(json.loads(line))
    return span_list


def get_span_dict(span_list):
    """Given a list of span names, return dictionary of span names."""
    strip_prefix = "python.ray.tests."
    span_names = {}
    for span in span_list:
        span_name = span["name"]
        if span_name.startswith(strip_prefix):
            span_name = span_name[len(strip_prefix) :]
        if span_name in span_names:
            span_names[span_name] += 1
        else:
            span_names[span_name] = 1
    return span_names


def task_helper():
    """Run a Ray task and check the spans produced."""

    @ray.remote
    def f(value):
        return value + 1

    obj_ref = f.remote(2)
    ray.get(obj_ref)

    span_list = get_span_list()
    assert len(span_list) == 2, span_list

    # The spans could show up in a different order, so just check that
    # all spans are as expected
    span_names = get_span_dict(span_list)
    assert span_names == {
        "test_tracing.f ray.remote": 1,
        "test_tracing.f ray.remote_worker": 1,
    }


def sync_actor_helper():
    """Run a Ray sync actor and check the spans produced."""

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

        def increment(self):
            self.value += 1
            return self.value

    # Create an actor from this class.
    counter = Counter.remote()
    obj_ref = counter.increment.remote()
    assert ray.get(obj_ref) == 1

    span_list = get_span_list()
    assert len(span_list) == 4

    # The spans could show up in a different order, so just check that
    # all spans are as expected
    span_names = get_span_dict(span_list)
    assert span_names == {
        "sync_actor_helper.<locals>.Counter.__init__ ray.remote": 1,
        "sync_actor_helper.<locals>.Counter.increment ray.remote": 1,
        "Counter.__init__ ray.remote_worker": 1,
        "Counter.increment ray.remote_worker": 1,
    }, span_names


def async_actor_helper():
    """Run a Ray async actor and check the spans produced."""

    @ray.remote
    class AsyncActor:
        # multiple invocation of this method can be running in
        # the event loop at the same time
        async def run_concurrent(self):
            await asyncio.sleep(2)  # concurrent workload here

    actor = AsyncActor.remote()
    ray.get([actor.run_concurrent.remote() for _ in range(4)])

    span_list = get_span_list()
    assert len(span_list) == 10

    # The spans could show up in a different order, so just check that
    # all spans are as expected
    span_names = get_span_dict(span_list)
    assert span_names == {
        "async_actor_helper.<locals>.AsyncActor.__init__ ray.remote": 1,
        "async_actor_helper.<locals>.AsyncActor.run_concurrent ray.remote": 4,
        "AsyncActor.__init__ ray.remote_worker": 1,
        "AsyncActor.run_concurrent ray.remote_worker": 4,
    }, span_names


def test_tracing_task_init_workflow(cleanup_dirs, ray_start_init_tracing):
    task_helper()


def test_tracing_task_start_workflow(cleanup_dirs, ray_start_cli_tracing):
    task_helper()


def test_tracing_sync_actor_init_workflow(cleanup_dirs, ray_start_init_tracing):
    sync_actor_helper()


def test_tracing_sync_actor_start_workflow(cleanup_dirs, ray_start_cli_tracing):
    sync_actor_helper()


def test_tracing_async_actor_init_workflow(cleanup_dirs, ray_start_init_tracing):
    async_actor_helper()


def test_tracing_async_actor_start_workflow(cleanup_dirs, ray_start_cli_tracing):
    async_actor_helper()


def test_tracing_predefined_actor(cleanup_dirs, ray_start_cli_predefined_actor_tracing):
    sync_actor_helper()


def test_tracing_actor_client_mode(cleanup_dirs, ray_start_cli_tracing_with_client):
    ray.init(ray_start_cli_tracing_with_client)
    assert ray.util.client.ray.is_connected()
    setup_tracing()
    _enable_tracing()

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

        def increment(self):
            self.value += 1
            return self.value

    counter = Counter.remote()
    assert ray.get(counter.increment.remote()) == 1

    span_list = get_span_list()
    span_names = get_span_dict(span_list)
    assert span_names == {
        "test_tracing_actor_client_mode.<locals>.Counter.__init__ ray.remote": 1,
        "Counter.__init__ ray.remote": 1,
        "Counter.increment ray.remote": 1,
        "Counter.__init__ ray.remote_worker": 1,
        "Counter.increment ray.remote_worker": 1,
    }, span_names

    actor_creation_spans = [
        span
        for span in span_list
        if span["name"].endswith("Counter.__init__ ray.remote")
    ]
    actor_method_span = next(
        span for span in span_list if span["name"] == "Counter.increment ray.remote"
    )
    actor_id = counter._actor_id.hex()
    for span in actor_creation_spans + [actor_method_span]:
        assert span["attributes"]["ray.actor_id"] == actor_id


def test_wrapping(ray_start_init_tracing):
    @ray.remote
    def f(**_kwargs):
        pass


def test_deserialization_works_without_opentelemetry(ray_start_regular):
    """
    Test that if a function is serialized with opentelemetry, it can be
    deserialized in a context without opentelemetry
    """

    @ray.remote
    def f():
        return 30

    fn_bytes = ray.cloudpickle.dumps(f)
    with patch.dict("sys.modules", opentelemetry=None):
        # Ensure that opentelemetry cannot be imported
        with pytest.raises(ModuleNotFoundError):
            import opentelemetry.trace  # noqa: F401
        ray.cloudpickle.loads(fn_bytes)


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