import json
import os
import shutil
import time
import uuid

import pytest
import ray

import sentry_sdk
from sentry_sdk.envelope import Envelope
from sentry_sdk.integrations.ray import RayIntegration
from sentry_sdk.integrations.stdlib import StdlibIntegration
from sentry_sdk.traces import SegmentNameSource
from tests.conftest import TestTransport


@pytest.fixture(autouse=True)
def shutdown_ray(tmpdir):
    yield
    ray.shutdown()


class RayTestTransport(TestTransport):
    def __init__(self):
        self.envelopes = []
        super().__init__()

    def capture_envelope(self, envelope: Envelope) -> None:
        self.envelopes.append(envelope)


class RayLoggingTransport(TestTransport):
    def capture_envelope(self, envelope: Envelope) -> None:
        print(envelope.serialize().decode("utf-8", "replace"))


def setup_sentry_with_logging_transport():
    setup_sentry(transport=RayLoggingTransport())


def setup_sentry_with_logging_transport_and_span_streaming():
    setup_sentry(span_streaming=True, transport=RayLoggingTransport())


def setup_sentry(span_streaming=False, transport=None):
    if span_streaming:
        sentry_sdk._span_batcher.SpanBatcher.MAX_BEFORE_FLUSH = 1

    sentry_sdk.init(
        integrations=[RayIntegration()],
        transport=RayTestTransport() if transport is None else transport,
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )


def read_error_from_log(job_id, ray_temp_dir):
    # Find the actual session directory that Ray created
    session_dirs = [d for d in os.listdir(ray_temp_dir) if d.startswith("session_")]
    if not session_dirs:
        raise FileNotFoundError(f"No session directory found in {ray_temp_dir}")

    session_dir = os.path.join(ray_temp_dir, session_dirs[0])
    log_dir = os.path.join(session_dir, "logs")

    if not os.path.exists(log_dir):
        raise FileNotFoundError(f"No logs directory found at {log_dir}")

    log_file = [
        f
        for f in os.listdir(log_dir)
        if "worker" in f and job_id in f and f.endswith(".out")
    ][0]

    next_line_is_event_payload = False
    with open(os.path.join(log_dir, log_file), "r") as file:
        for line in file:
            try:
                payload = json.loads(line)
            except ValueError:
                continue

            if next_line_is_event_payload:
                return payload

            if isinstance(payload, dict) and payload.get("type") == "event":
                next_line_is_event_payload = True

    return None


def _parse_spans_from_log(job_id, ray_temp_dir):
    # Find the actual session directory that Ray created
    session_dirs = [d for d in os.listdir(ray_temp_dir) if d.startswith("session_")]
    if not session_dirs:
        raise FileNotFoundError(f"No session directory found in {ray_temp_dir}")

    session_dir = os.path.join(ray_temp_dir, session_dirs[0])
    log_dir = os.path.join(session_dir, "logs")

    if not os.path.exists(log_dir):
        raise FileNotFoundError(f"No logs directory found at {log_dir}")

    log_files = [
        f
        for f in os.listdir(log_dir)
        if "worker" in f and job_id in f and f.endswith(".out")
    ]
    if not log_files:
        return []

    spans = []
    next_line_is_span_payload = False
    with open(os.path.join(log_dir, log_files[0]), "r") as file:
        for line in file:
            try:
                payload = json.loads(line)
            except ValueError:
                continue
            if next_line_is_span_payload:
                spans.extend(payload["items"])
                next_line_is_span_payload = False
                continue
            if isinstance(payload, dict) and payload.get("type") == "span":
                next_line_is_span_payload = True

    return spans


def read_spans_from_log(job_id, ray_temp_dir, min_spans=1, timeout=10):
    deadline = time.monotonic() + timeout
    spans = []
    while True:
        try:
            spans = _parse_spans_from_log(job_id, ray_temp_dir)
        except FileNotFoundError:
            spans = []
        if len(spans) >= min_spans or time.monotonic() >= deadline:
            return spans
        time.sleep(0.1)


def example_task(span_streaming: bool):
    if span_streaming:
        with sentry_sdk.traces.start_span(
            name="example task step",
            attributes={
                "sentry.op": "task",
            },
        ):
            ...
    else:
        with sentry_sdk.start_span(op="task", name="example task step"):
            ...

        return sentry_sdk.get_client().transport.envelopes


# RayIntegration must leave variadic keyword arguments at the end
def example_task_with_kwargs(span_streaming: bool, **kwargs):
    if span_streaming:
        with sentry_sdk.traces.start_span(
            name="example task step", attributes={"sentry.op": "task"}
        ):
            ...
    else:
        with sentry_sdk.start_span(op="task", name="example task step"):
            ...

        return sentry_sdk.get_client().transport.envelopes


@pytest.mark.parametrize(
    "task_options", [{}, {"num_cpus": 0, "memory": 1024 * 1024 * 10}]
)
@pytest.mark.parametrize(
    "task",
    [example_task, example_task_with_kwargs],
)
@pytest.mark.parametrize("span_streaming", [True, False])
def test_tracing_in_ray_tasks(task_options, task, span_streaming):
    sentry_sdk.init(
        integrations=[RayIntegration()],
        disabled_integrations=[StdlibIntegration],
        transport=RayTestTransport(),
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    # Setup ray task, calling decorator directly instead of @,
    # to accommodate for test parametrization
    if task_options:
        example_task = ray.remote(**task_options)(task)
    else:
        example_task = ray.remote(task)

    # Function name shouldn't be overwritten by Sentry wrapper
    assert (
        example_task._function_name
        == f"tests.integrations.ray.test_ray.{task.__name__}"
    )

    if span_streaming:
        ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}")
        os.makedirs(ray_temp_dir, exist_ok=True)

        try:
            ray.init(
                runtime_env={
                    "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming,
                    "working_dir": "./",
                },
                _temp_dir=ray_temp_dir,
            )

            with sentry_sdk.traces.start_span(
                name="ray test parent", attributes={"sentry.op": "task"}
            ):
                future = example_task.remote(span_streaming)
                ray.get(future)

            job_id = future.job_id().hex()
            worker_spans = read_spans_from_log(job_id, ray_temp_dir, min_spans=2)
        finally:
            if os.path.exists(ray_temp_dir):
                shutil.rmtree(ray_temp_dir, ignore_errors=True)

        sentry_sdk.flush()
        client_envelope = sentry_sdk.get_client().transport.envelopes[0]
        client_spans = [
            span
            for item in client_envelope.items
            for span in item.payload.json["items"]
        ]
        assert client_spans[1]["name"] == "ray test parent"

        assert (
            worker_spans[1]["name"]
            == f"tests.integrations.ray.test_ray.{task.__name__}"
        )
        assert (
            worker_spans[1]["attributes"]["sentry.segment.name.source"]["value"]
            == SegmentNameSource.TASK
        )

        span = client_spans[0]
        assert span["attributes"]["sentry.op"]["value"] == "queue.submit.ray"
        assert span["attributes"]["sentry.origin"]["value"] == "auto.queue.ray"
        assert span["name"] == f"tests.integrations.ray.test_ray.{task.__name__}"
        assert span["parent_span_id"] == client_spans[1]["span_id"]
        assert span["trace_id"] == client_spans[1]["trace_id"]

        span = worker_spans[0]
        assert span["attributes"]["sentry.op"]["value"] == "task"
        assert span["attributes"]["sentry.origin"]["value"] == "manual"
        assert span["name"] == "example task step"
        assert span["parent_span_id"] == worker_spans[1]["span_id"]
        assert span["trace_id"] == worker_spans[1]["trace_id"]

        assert client_spans[1]["trace_id"] == worker_spans[1]["trace_id"]
    else:
        ray.init(
            runtime_env={
                "worker_process_setup_hook": setup_sentry,
                "working_dir": "./",
            }
        )

        with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
            worker_envelopes = ray.get(example_task.remote(span_streaming))

        client_envelope = sentry_sdk.get_client().transport.envelopes[0]
        client_transaction = client_envelope.get_transaction_event()
        assert client_transaction["transaction"] == "ray test transaction"
        assert client_transaction["transaction_info"] == {"source": "custom"}

        worker_envelope = worker_envelopes[0]
        worker_transaction = worker_envelope.get_transaction_event()
        assert (
            worker_transaction["transaction"]
            == f"tests.integrations.ray.test_ray.{task.__name__}"
        )
        assert worker_transaction["transaction_info"] == {"source": "task"}

        (span,) = client_transaction["spans"]
        assert span["op"] == "queue.submit.ray"
        assert span["origin"] == "auto.queue.ray"
        assert span["description"] == f"tests.integrations.ray.test_ray.{task.__name__}"
        assert (
            span["parent_span_id"] == client_transaction["contexts"]["trace"]["span_id"]
        )
        assert span["trace_id"] == client_transaction["contexts"]["trace"]["trace_id"]

        (span,) = worker_transaction["spans"]
        assert span["op"] == "task"
        assert span["origin"] == "manual"
        assert span["description"] == "example task step"
        assert (
            span["parent_span_id"] == worker_transaction["contexts"]["trace"]["span_id"]
        )
        assert span["trace_id"] == worker_transaction["contexts"]["trace"]["trace_id"]

        assert (
            client_transaction["contexts"]["trace"]["trace_id"]
            == worker_transaction["contexts"]["trace"]["trace_id"]
        )


@pytest.mark.parametrize("span_streaming", [True, False])
def test_errors_in_ray_tasks(span_streaming):
    sentry_sdk.init(
        integrations=[RayIntegration()],
        disabled_integrations=[StdlibIntegration],
        transport=RayTestTransport(),
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}")
    os.makedirs(ray_temp_dir, exist_ok=True)

    try:
        ray.init(
            runtime_env={
                "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming
                if span_streaming
                else setup_sentry_with_logging_transport,
                "working_dir": "./",
            },
            _temp_dir=ray_temp_dir,
        )

        # Setup ray task
        @ray.remote
        def example_task():
            1 / 0

        if span_streaming:
            with sentry_sdk.traces.start_span(
                name="ray test parent", attributes={"sentry.op": "task"}
            ):
                with pytest.raises(ZeroDivisionError):
                    future = example_task.remote()
                    ray.get(future)
        else:
            with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
                with pytest.raises(ZeroDivisionError):
                    future = example_task.remote()
                    ray.get(future)

        job_id = future.job_id().hex()
        error = read_error_from_log(job_id, ray_temp_dir)

        assert error["level"] == "error"
        assert (
            error["transaction"]
            == "tests.integrations.ray.test_ray.test_errors_in_ray_tasks.<locals>.example_task"
        )
        assert error["exception"]["values"][0]["mechanism"]["type"] == "ray"
        assert not error["exception"]["values"][0]["mechanism"]["handled"]

    finally:
        if os.path.exists(ray_temp_dir):
            shutil.rmtree(ray_temp_dir, ignore_errors=True)


# Arbitrary keyword argument to test all decorator paths
@pytest.mark.parametrize("remote_kwargs", [{}, {"namespace": "actors"}])
@pytest.mark.parametrize("span_streaming", [True, False])
def test_tracing_in_ray_actors(remote_kwargs, span_streaming):
    sentry_sdk.init(
        integrations=[RayIntegration()],
        disabled_integrations=[StdlibIntegration],
        transport=RayTestTransport(),
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    # Setup ray actor
    if remote_kwargs:

        @ray.remote(**remote_kwargs)
        class Counter:
            def __init__(self):
                self.n = 0

            def increment(self):
                if span_streaming:
                    with sentry_sdk.traces.start_span(
                        name="example actor execution", attributes={"sentry.op": "task"}
                    ):
                        self.n += 1
                else:
                    with sentry_sdk.start_span(
                        op="task", name="example actor execution"
                    ):
                        self.n += 1

                    return sentry_sdk.get_client().transport.envelopes
    else:

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

            def increment(self):
                if span_streaming:
                    with sentry_sdk.traces.start_span(
                        name="example actor execution", attributes={"sentry.op": "task"}
                    ):
                        self.n += 1
                else:
                    with sentry_sdk.start_span(
                        op="task", name="example actor execution"
                    ):
                        self.n += 1

                    return sentry_sdk.get_client().transport.envelopes

    if span_streaming:
        ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}")
        os.makedirs(ray_temp_dir, exist_ok=True)

        try:
            ray.init(
                runtime_env={
                    "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming,
                    "working_dir": "./",
                },
                _temp_dir=ray_temp_dir,
            )

            with sentry_sdk.traces.start_span(
                name="ray test parent", attributes={"sentry.op": "task"}
            ):
                counter = Counter.remote()
                future = counter.increment.remote()
                ray.get(future)

            job_id = future.job_id().hex()
            worker_spans = read_spans_from_log(job_id, ray_temp_dir)
        finally:
            if os.path.exists(ray_temp_dir):
                shutil.rmtree(ray_temp_dir, ignore_errors=True)

        sentry_sdk.flush()
        client_envelope = sentry_sdk.get_client().transport.envelopes[0]
        client_spans = [
            span
            for item in client_envelope.items
            for span in item.payload.json["items"]
        ]

        # Spans for submitting the actor task are not created (actors are not supported yet)
        # Only the manual "example actor execution" span is recorded.
        assert len(client_spans) == 1

        # Transaction are not yet created when executing ray actors (actors are not supported yet)
        # Only the manual "example actor execution" span is recorded.
        assert len(worker_spans) == 1
    else:
        ray.init(
            runtime_env={
                "worker_process_setup_hook": setup_sentry,
                "working_dir": "./",
            }
        )

        with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
            counter = Counter.remote()
            worker_envelopes = ray.get(counter.increment.remote())

        client_envelope = sentry_sdk.get_client().transport.envelopes[0]
        client_transaction = client_envelope.get_transaction_event()

        # Spans for submitting the actor task are not created (actors are not supported yet)
        assert client_transaction["spans"] == []

        # Transaction are not yet created when executing ray actors (actors are not supported yet)
        assert worker_envelopes == []


@pytest.mark.parametrize("span_streaming", [True, False])
def test_errors_in_ray_actors(span_streaming):
    sentry_sdk.init(
        integrations=[RayIntegration()],
        disabled_integrations=[StdlibIntegration],
        transport=RayLoggingTransport(),
        traces_sample_rate=1.0,
        trace_lifecycle="stream" if span_streaming else "static",
    )

    ray_temp_dir = os.path.join("/tmp", f"ray_test_{uuid.uuid4().hex[:8]}")
    os.makedirs(ray_temp_dir, exist_ok=True)

    try:
        ray.init(
            runtime_env={
                "worker_process_setup_hook": setup_sentry_with_logging_transport_and_span_streaming
                if span_streaming
                else setup_sentry_with_logging_transport,
                "working_dir": "./",
            },
            _temp_dir=ray_temp_dir,
        )

        # Setup ray actor
        @ray.remote
        class Counter:
            def __init__(self):
                self.n = 0

            def increment(self):
                if span_streaming:
                    with sentry_sdk.traces.start_span(
                        name="example actor execution", attributes={"sentry.op": "task"}
                    ):
                        1 / 0
                else:
                    with sentry_sdk.start_span(
                        op="task", name="example actor execution"
                    ):
                        1 / 0

                return sentry_sdk.get_client().transport.envelopes

        if span_streaming:
            with sentry_sdk.traces.start_span(
                name="ray test parent", attributes={"sentry.op": "task"}
            ):
                with pytest.raises(ZeroDivisionError):
                    counter = Counter.remote()
                    future = counter.increment.remote()
                    ray.get(future)
        else:
            with sentry_sdk.start_transaction(op="task", name="ray test transaction"):
                with pytest.raises(ZeroDivisionError):
                    counter = Counter.remote()
                    future = counter.increment.remote()
                    ray.get(future)

        job_id = future.job_id().hex()
        error = read_error_from_log(job_id, ray_temp_dir)

        # We do not capture errors in ray actors yet
        assert error is None

    finally:
        if os.path.exists(ray_temp_dir):
            shutil.rmtree(ray_temp_dir, ignore_errors=True)
