"""Unit tests for RDTManager."""
import re
import sys
from dataclasses import dataclass
from typing import Any, List

import pytest

from ray.exceptions import GetTimeoutError
from ray.experimental import (
    CommunicatorMetadata,
    TensorTransportManager,
    TensorTransportMetadata,
    register_tensor_transport,
)
from ray.experimental.rdt.rdt_manager import RDTManager, RDTMeta
from ray.experimental.rdt.tensor_transport_manager import FetchRequest

_BACKEND_NAME = "TEST_PIPELINE"
_TWO_SIDED_BACKEND_NAME = "TEST_TWO_SIDED"


@dataclass
class _TestCommMeta(CommunicatorMetadata):
    pass


@dataclass
class _TrackedFetchRequest(FetchRequest):
    """FetchRequest subclass that records when it is deleted."""

    def __del__(self):
        _PipelineCheckingTransport.deleted_requests.add(self.obj_id)


class _PipelineCheckingTransport(TensorTransportManager):
    """Fake one-sided transport that records the order of fetch/wait calls.

    Each fetch_multiple_tensors call appends ("fetch", obj_id) to call_log,
    and each wait_fetch_complete call appends ("wait", obj_id).  The test
    asserts that all fetch entries appear before any wait entry.

    call_log is a class-level list so the singleton instance created by
    get_tensor_transport_manager records to the same list across all tests.
    """

    call_log: List = []
    fail_on_wait: set = set()
    wait_delay: float = 0
    deleted_requests: set = set()

    def tensor_transport_backend(self) -> str:
        return _BACKEND_NAME

    @staticmethod
    def is_one_sided() -> bool:
        return True

    @staticmethod
    def can_abort_transport() -> bool:
        return False

    def actor_has_tensor_transport(self, actor) -> bool:
        return True

    def extract_tensor_transport_metadata(self, obj_id, rdt_object):
        return TensorTransportMetadata(tensor_meta=[], tensor_device="cpu")

    def get_communicator_metadata(self, src_actor, dst_actor, backend=None):
        return _TestCommMeta()

    def fetch_multiple_tensors(
        self,
        obj_id: str,
        tensor_transport_metadata,
        communicator_metadata,
        target_buffers=None,
    ) -> FetchRequest:
        self.__class__.call_log.append(("fetch", obj_id))
        return _TrackedFetchRequest(obj_id=obj_id, tensors=[f"val:{obj_id}"])

    def wait_fetch_complete(
        self, fetch_request: FetchRequest, timeout: float = -1
    ) -> List[Any]:
        if self.__class__.wait_delay > 0:
            import time

            time.sleep(self.__class__.wait_delay)
        self.__class__.call_log.append(("wait", fetch_request.obj_id))
        if fetch_request.obj_id in self.__class__.fail_on_wait:
            raise RuntimeError(f"wait failed for {fetch_request.obj_id}")
        return fetch_request.tensors

    def recv_multiple_tensors(self, obj_id, meta, comm_meta, target_buffers=None):
        return []

    def send_multiple_tensors(self, tensors, meta, comm_meta):
        pass

    def garbage_collect(self, obj_id, meta, tensors):
        pass

    def abort_transport(self, obj_id, comm_meta):
        pass


class _TwoSidedTransport(TensorTransportManager):
    """Fake two-sided transport (e.g. NCCL/GLOO style)."""

    def tensor_transport_backend(self) -> str:
        return _TWO_SIDED_BACKEND_NAME

    @staticmethod
    def is_one_sided() -> bool:
        return False

    @staticmethod
    def can_abort_transport() -> bool:
        return False

    def actor_has_tensor_transport(self, actor) -> bool:
        return True

    def extract_tensor_transport_metadata(self, obj_id, rdt_object):
        return TensorTransportMetadata(tensor_meta=[], tensor_device="cpu")

    def get_communicator_metadata(self, src_actor, dst_actor, backend=None):
        return _TestCommMeta()

    def fetch_multiple_tensors(self, obj_id, meta, comm_meta, target_buffers=None):
        raise NotImplementedError

    def recv_multiple_tensors(self, obj_id, meta, comm_meta, target_buffers=None):
        raise NotImplementedError

    def send_multiple_tensors(self, tensors, meta, comm_meta):
        raise NotImplementedError

    def garbage_collect(self, obj_id, meta, tensors):
        pass

    def abort_transport(self, obj_id, comm_meta):
        pass


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module", autouse=True)
def register_test_transports():
    """Register both test transports once for the lifetime of the module."""
    try:
        register_tensor_transport(
            _BACKEND_NAME, ["cpu"], _PipelineCheckingTransport, list
        )
    except ValueError:
        pass  # already registered (e.g. test module loaded more than once)
    try:
        register_tensor_transport(
            _TWO_SIDED_BACKEND_NAME, ["cpu"], _TwoSidedTransport, list
        )
    except ValueError:
        pass


@pytest.fixture(autouse=True)
def clear_call_log():
    """Reset the pipeline transport's call log and error config before each test."""
    _PipelineCheckingTransport.call_log.clear()
    _PipelineCheckingTransport.fail_on_wait.clear()
    _PipelineCheckingTransport.wait_delay = 0
    _PipelineCheckingTransport.deleted_requests.clear()


def _build_manager(object_ids: List[str], backend: str = _BACKEND_NAME) -> RDTManager:
    """Return an RDTManager pre-populated with fake RDT metadata.

    Uses a real RDTStore so no Ray cluster is required.
    All objects are non-primary copies (pop_object=True in fetch_and_get_rdt_objects).
    """
    manager = RDTManager()

    meta = TensorTransportMetadata(tensor_meta=[], tensor_device="cpu")
    for obj_id in object_ids:
        manager.set_rdt_metadata(
            obj_id,
            RDTMeta(
                src_actor=None,
                tensor_transport_backend=backend,
                tensor_transport_meta=meta,
                sent_dest_actors=set(),
                sent_to_src_actor_and_others_warned=False,
                target_buffers=None,
            ),
        )

    return manager


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


def test_fetch_and_get():
    object_ids = ["obj1", "obj2", "obj3"]
    manager = _build_manager(object_ids)
    result = manager.fetch_and_get_rdt_objects(object_ids)

    call_log = _PipelineCheckingTransport.call_log
    fetch_indices = [i for i, (kind, _) in enumerate(call_log) if kind == "fetch"]
    wait_indices = [i for i, (kind, _) in enumerate(call_log) if kind == "wait"]

    # All fetch_multiple_tensors calls must precede all wait_fetch_complete
    # calls.
    assert len(fetch_indices) == len(object_ids), f"call_log={call_log}"
    assert len(wait_indices) == len(object_ids), f"call_log={call_log}"
    assert max(fetch_indices) < min(
        wait_indices
    ), f"Expected all fetches before all waits, got call_log={call_log}"

    # One entry per requested object ID.
    assert set(result.keys()) == set(object_ids)

    call_log = _PipelineCheckingTransport.call_log
    # Each object ID triggers exactly one fetch and one wait.
    fetched = [oid for kind, oid in call_log if kind == "fetch"]
    waited = [oid for kind, oid in call_log if kind == "wait"]

    assert sorted(fetched) == sorted(object_ids)
    assert sorted(waited) == sorted(object_ids)


def test_primary_copy_objects_skip_fetch():
    """Objects already in the store must not trigger a fetch."""
    secondary_ids = ["secondary1", "secondary2"]
    primary_id = "primary1"
    manager = _build_manager(secondary_ids + [primary_id])

    # Add the primary-copy and one secondary-copy object to the store directly.
    # Phase 1 of fetch_and_get_rdt_objects skips objects in store.
    manager.rdt_store.add_object(primary_id, ["primary_value"], is_primary=True)
    manager.rdt_store.add_object(secondary_ids[0], ["secondary"], is_primary=False)
    result = manager.fetch_and_get_rdt_objects(secondary_ids + [primary_id])

    call_log = _PipelineCheckingTransport.call_log
    fetched = [oid for kind, oid in call_log if kind == "fetch"]
    assert set(fetched) == set(
        secondary_ids[1:]
    ), f"objects in store should not be fetched; got fetched={fetched}"
    # One fetch + one wait for each secondary object; zero for the primary one.
    assert len(call_log) == 2, f"call_log={call_log}"
    # All objects should be returned in the results.
    assert set(result.keys()) == set(secondary_ids + [primary_id])
    assert result[primary_id] == ["primary_value"]
    assert result[secondary_ids[0]] == ["secondary"]


def test_empty_object_list_returns_empty_dict():
    """Calling fetch_and_get_rdt_objects with an empty list returns an empty dict."""
    manager = _build_manager([])
    result = manager.fetch_and_get_rdt_objects([])

    assert result == {}
    assert _PipelineCheckingTransport.call_log == []


def test_two_sided_transport_raises_on_fetch_and_get_rdt_objects():
    """ray.get (use_object_store=False) must raise ValueError for two-sided transports."""
    obj_id = "two_sided_obj"
    manager = _build_manager([obj_id], backend=_TWO_SIDED_BACKEND_NAME)

    with pytest.raises(
        ValueError,
        match=re.escape(
            f"ray.get is not allowed on RDT objects using the two-sided transport {_TWO_SIDED_BACKEND_NAME}. "
            "Either use a one-sided RDT transport or pass _use_object_store=True to ray.get to fetch the object through the object store instead."
        ),
    ):
        manager.fetch_and_get_rdt_objects([obj_id], use_object_store=False)


def test_fetch_requests_deleted_on_exception():
    """If _wait_fetch raises, all FetchRequests are deleted (cleaning up resources via __del__)."""
    import gc

    object_ids = ["obj1", "obj2", "obj3"]
    manager = _build_manager(object_ids)
    _PipelineCheckingTransport.fail_on_wait.add("obj1")

    with pytest.raises(RuntimeError, match="wait failed for obj1"):
        manager.fetch_and_get_rdt_objects(object_ids)

    gc.collect()
    assert _PipelineCheckingTransport.deleted_requests == set(object_ids), (
        f"All FetchRequests must be GCed even if one fails; "
        f"deleted={_PipelineCheckingTransport.deleted_requests}"
    )


def test_object_fetch_timed_out_error():
    """fetch_and_get_rdt_objects raises ObjectFetchTimedOutError when RDT timeout is hit."""
    from ray.exceptions import ObjectFetchTimedOutError

    object_ids = ["obj1", "obj2"]
    manager = _build_manager(object_ids)
    # Make wait_fetch_complete slow enough to exceed a very short timeout.
    _PipelineCheckingTransport.wait_delay = 0.2

    with pytest.raises(ObjectFetchTimedOutError):
        # timeout=None means no user timeout, so only RDT timeout applies.
        # We monkeypatch the constant to a very small value.
        import ray._private.ray_constants as rc

        original = rc.RDT_FETCH_FAIL_TIMEOUT_SECONDS
        rc.RDT_FETCH_FAIL_TIMEOUT_SECONDS = 0.1
        try:
            manager.fetch_and_get_rdt_objects(object_ids)
        finally:
            rc.RDT_FETCH_FAIL_TIMEOUT_SECONDS = original


def test_get_timed_out_error():
    """fetch_and_get_rdt_objects raises GetTimeoutError when user timeout is hit."""
    object_ids = ["obj1", "obj2"]
    manager = _build_manager(object_ids)
    # Make wait_fetch_complete slow enough to exceed a very short timeout.
    _PipelineCheckingTransport.wait_delay = 0.2

    # Check that user timeout triggers before fetch fail timeout.
    with pytest.raises(GetTimeoutError):
        import ray._private.ray_constants as rc

        original = rc.RDT_FETCH_FAIL_TIMEOUT_SECONDS
        rc.RDT_FETCH_FAIL_TIMEOUT_SECONDS = 1
        try:
            manager.fetch_and_get_rdt_objects(object_ids, timeout=0.1)
        finally:
            rc.RDT_FETCH_FAIL_TIMEOUT_SECONDS = original


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