import asyncio
import pickle
import sys
from types import SimpleNamespace
from typing import Union

import pytest

import ray
from ray import ObjectRef, ObjectRefGenerator
from ray._common.test_utils import SignalActor, async_wait_for_condition
from ray._common.utils import get_or_create_event_loop
from ray.exceptions import ActorDiedError, ActorUnavailableError, TaskCancelledError
from ray.serve._private.common import (
    DeploymentID,
    ReplicaID,
    ReplicaQueueLengthInfo,
    RequestMetadata,
    RunningReplicaInfo,
)
from ray.serve._private.constants import SERVE_NAMESPACE
from ray.serve._private.request_router.common import PendingRequest
from ray.serve._private.request_router.replica_wrapper import RunningReplica
from ray.serve._private.test_utils import send_signal_on_cancellation
from ray.serve._private.utils import Semaphore


class _IntMetricsManager:
    """Minimal metrics manager that tracks only the in-flight count."""

    def __init__(self):
        self._n = 0

    def get_num_ongoing_requests(self):
        return self._n

    def inc_num_ongoing_requests(self, _):
        self._n += 1

    def dec_num_ongoing_requests(self, _):
        self._n -= 1


@ray.remote(num_cpus=0)
class SlotReservationActor:
    """Ray actor wrapping the real Replica.reserve_slot / release_slot.

    Used by integration tests that need production slot-reservation logic
    running under Ray's actor concurrency model — unit tests share one event
    loop and can't observe sync/async ordering on a real ReplicaActor.
    """

    def __init__(self, max_ongoing_requests: int):
        from ray.serve._private.replica import Replica

        replica = Replica.__new__(Replica)
        replica._deployment_config = SimpleNamespace(
            max_ongoing_requests=max_ongoing_requests
        )
        replica._reserved_slots = set()
        replica._semaphore = Semaphore(lambda: max_ongoing_requests)
        replica._metrics_manager = _IntMetricsManager()
        # __init__ is bypassed; set the quiesce flag read by
        # _can_accept_request (reservations are rejected once quiescing).
        replica._quiescing = False
        self._replica = replica

    async def reserve_slot(self, request_metadata, slot_token: str):
        return await self._replica.reserve_slot(request_metadata, slot_token)

    def release_slot(self, slot_token: str):
        return self._replica.release_slot(slot_token)

    def get_num_ongoing_requests(self) -> int:
        return self._replica.get_num_ongoing_requests()


@ray.remote(num_cpus=0)
class BlockingReserveActor:
    """Actor whose reserve_slot blocks on a SignalActor.

    Records every release_slot token it receives so a test can verify the
    cancellation cleanup path in RunningReplica.reserve_slot.
    """

    def __init__(self, signal_actor):
        self._signal = signal_actor
        self._released_tokens = []

    async def reserve_slot(self, request_metadata, slot_token: str):
        await self._signal.wait.remote()
        return True, 1

    def release_slot(self, slot_token: str):
        self._released_tokens.append(slot_token)
        return True, 0

    def get_released_tokens(self):
        return list(self._released_tokens)


@ray.remote(num_cpus=0)
class FakeReplicaActor:
    def __init__(self):
        self._replica_queue_length_info = None

    def set_replica_queue_length_info(self, info: ReplicaQueueLengthInfo):
        self._replica_queue_length_info = info

    async def handle_request(
        self,
        request_metadata: Union[bytes, RequestMetadata],
        message: str,
        *,
        is_streaming: bool,
    ):
        if isinstance(request_metadata, bytes):
            request_metadata = pickle.loads(request_metadata)

        assert not is_streaming and not request_metadata.is_streaming
        return message

    async def handle_request_streaming(
        self,
        request_metadata: Union[bytes, RequestMetadata],
        message: str,
        *,
        is_streaming: bool,
    ):
        if isinstance(request_metadata, bytes):
            request_metadata = pickle.loads(request_metadata)

        assert is_streaming and request_metadata.is_streaming
        for i in range(5):
            yield f"{message}-{i}"

    async def handle_request_with_rejection(
        self,
        pickled_request_metadata: bytes,
        *args,
        **kwargs,
    ):
        cancelled_signal_actor = kwargs.pop("cancelled_signal_actor", None)
        if cancelled_signal_actor is not None:
            executing_signal_actor = kwargs.pop("executing_signal_actor")
            async with send_signal_on_cancellation(cancelled_signal_actor):
                await executing_signal_actor.send.remote()

            return

        # Special case: if "raise_task_cancelled_error" is in kwargs, raise TaskCancelledError
        # This simulates the scenario where the underlying Ray task gets cancelled
        if kwargs.pop("raise_task_cancelled_error", False):
            raise TaskCancelledError()

        yield pickle.dumps(self._replica_queue_length_info)
        if not self._replica_queue_length_info.accepted:
            return

        request_metadata = pickle.loads(pickled_request_metadata)
        if request_metadata.is_streaming:
            async for result in self.handle_request_streaming(
                request_metadata, *args, **kwargs
            ):
                yield result
        else:
            yield await self.handle_request(request_metadata, *args, **kwargs)


@pytest.fixture
def setup_fake_replica(ray_instance) -> RunningReplica:
    replica_id = ReplicaID(
        "fake_replica", deployment_id=DeploymentID(name="fake_deployment")
    )
    actor_name = replica_id.to_full_id_str()
    # Create actor with a name so it can be retrieved by get_actor_handle()
    _ = FakeReplicaActor.options(
        name=actor_name, namespace=SERVE_NAMESPACE, lifetime="detached"
    ).remote()
    return RunningReplicaInfo(
        replica_id=replica_id,
        node_id=None,
        node_ip=None,
        availability_zone=None,
        actor_name=actor_name,
        max_ongoing_requests=10,
        is_cross_language=False,
    )


def test_update_replica_info_refreshes_backend_http_endpoint(setup_fake_replica):
    replica = RunningReplica(setup_fake_replica)
    assert replica.backend_http_endpoint is None

    updated_info = RunningReplicaInfo(
        replica_id=setup_fake_replica.replica_id,
        node_id=setup_fake_replica.node_id,
        node_ip="127.0.0.1",
        availability_zone=setup_fake_replica.availability_zone,
        actor_name=setup_fake_replica.actor_name,
        max_ongoing_requests=setup_fake_replica.max_ongoing_requests,
        is_cross_language=setup_fake_replica.is_cross_language,
        backend_http_port=8001,
    )

    replica.update_replica_info(updated_info)
    assert replica.backend_http_endpoint == ("127.0.0.1", 8001)


def test_backend_http_endpoint_requires_host_and_port(setup_fake_replica):
    replica = RunningReplica(setup_fake_replica)

    updated_info = RunningReplicaInfo(
        replica_id=setup_fake_replica.replica_id,
        node_id=setup_fake_replica.node_id,
        node_ip=None,
        availability_zone=setup_fake_replica.availability_zone,
        actor_name=setup_fake_replica.actor_name,
        max_ongoing_requests=setup_fake_replica.max_ongoing_requests,
        is_cross_language=setup_fake_replica.is_cross_language,
        backend_http_port=8001,
    )

    replica.update_replica_info(updated_info)
    assert replica.backend_http_endpoint is None


@pytest.mark.asyncio
@pytest.mark.parametrize("is_streaming", [False, True])
async def test_send_request_without_rejection(setup_fake_replica, is_streaming: bool):
    replica = RunningReplica(setup_fake_replica)

    pr = PendingRequest(
        args=["Hello"],
        kwargs={"is_streaming": is_streaming},
        metadata=RequestMetadata(
            request_id="abc",
            internal_request_id="def",
            is_streaming=is_streaming,
        ),
    )
    replica_result = replica.try_send_request(pr, with_rejection=False)
    if is_streaming:
        assert isinstance(replica_result.to_object_ref_gen(), ObjectRefGenerator)
        for i in range(5):
            assert await replica_result.__anext__() == f"Hello-{i}"
    else:
        assert isinstance(replica_result.to_object_ref(), ObjectRef)
        assert isinstance(await replica_result.to_object_ref_async(), ObjectRef)
        assert await replica_result.get_async() == "Hello"


@pytest.mark.asyncio
@pytest.mark.parametrize("accepted", [False, True])
@pytest.mark.parametrize("is_streaming", [False, True])
async def test_send_request_with_rejection(
    setup_fake_replica, accepted: bool, is_streaming: bool
):
    actor_handle = setup_fake_replica.get_actor_handle()
    replica = RunningReplica(setup_fake_replica)
    ray.get(
        actor_handle.set_replica_queue_length_info.remote(
            ReplicaQueueLengthInfo(accepted=accepted, num_ongoing_requests=10),
        )
    )

    pr = PendingRequest(
        args=["Hello"],
        kwargs={"is_streaming": is_streaming},
        metadata=RequestMetadata(
            request_id="abc",
            internal_request_id="def",
            is_streaming=is_streaming,
        ),
    )
    replica_result = replica.try_send_request(pr, with_rejection=True)
    info = await replica_result.get_rejection_response()
    assert info.accepted == accepted
    assert info.num_ongoing_requests == 10
    if not accepted:
        pass
    elif is_streaming:
        assert isinstance(replica_result.to_object_ref_gen(), ObjectRefGenerator)
        for i in range(5):
            assert await replica_result.__anext__() == f"Hello-{i}"
    else:
        assert isinstance(replica_result.to_object_ref(), ObjectRef)
        assert isinstance(await replica_result.to_object_ref_async(), ObjectRef)
        assert await replica_result.get_async() == "Hello"


@pytest.mark.asyncio
async def test_send_request_with_rejection_cancellation(setup_fake_replica):
    """
    Verify that the downstream actor method call is cancelled if the call to send the
    request to the replica is cancelled.
    """
    replica = RunningReplica(setup_fake_replica)

    executing_signal_actor = SignalActor.remote()
    cancelled_signal_actor = SignalActor.remote()

    pr = PendingRequest(
        args=["Hello"],
        kwargs={
            "cancelled_signal_actor": cancelled_signal_actor,
            "executing_signal_actor": executing_signal_actor,
        },
        metadata=RequestMetadata(
            request_id="abc",
            internal_request_id="def",
        ),
    )

    # Send request should hang because the downstream actor method call blocks
    # before sending the system message.
    replica_result = replica.try_send_request(pr, with_rejection=True)
    request_task = get_or_create_event_loop().create_task(
        replica_result.get_rejection_response()
    )

    # Check that the downstream actor method call has started.
    await executing_signal_actor.wait.remote()

    _, pending = await asyncio.wait([request_task], timeout=0.001)
    assert len(pending) == 1

    # Cancel the task. This should cause the downstream actor method call to
    # be cancelled (verified via signal actor).
    request_task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await request_task

    await cancelled_signal_actor.wait.remote()


@pytest.mark.asyncio
async def test_send_request_with_rejection_task_cancelled_error(setup_fake_replica):
    """
    Test that TaskCancelledError from the underlying Ray task gets converted to
    asyncio.CancelledError when sending request with rejection.
    """
    actor_handle = setup_fake_replica.get_actor_handle()
    replica = RunningReplica(setup_fake_replica)

    # Set up the replica to accept the request
    ray.get(
        actor_handle.set_replica_queue_length_info.remote(
            ReplicaQueueLengthInfo(accepted=True, num_ongoing_requests=5),
        )
    )

    pr = PendingRequest(
        args=["Hello"],
        kwargs={
            "raise_task_cancelled_error": True
        },  # This will trigger TaskCancelledError
        metadata=RequestMetadata(
            request_id="abc",
            internal_request_id="def",
        ),
    )

    # The TaskCancelledError should be caught and converted to asyncio.CancelledError
    replica_result = replica.try_send_request(pr, with_rejection=True)
    with pytest.raises(asyncio.CancelledError):
        await replica_result.get_rejection_response()


def _spawn_running_replica(actor_cls, replica_id_str: str, *actor_args, **actor_kwargs):
    """Spawn a named actor and wrap it in a RunningReplica.

    Returns ``(running_replica, actor_handle)``. The actor must be created
    with the canonical replica-id name so RunningReplica can resolve it
    through its normal GCS lookup.
    """
    replica_id = ReplicaID(
        replica_id_str, deployment_id=DeploymentID(name="slot_reservation_test")
    )
    actor_name = replica_id.to_full_id_str()
    actor_handle = actor_cls.options(
        name=actor_name, namespace=SERVE_NAMESPACE, lifetime="detached"
    ).remote(*actor_args, **actor_kwargs)
    info = RunningReplicaInfo(
        replica_id=replica_id,
        node_id=None,
        node_ip=None,
        availability_zone=None,
        actor_name=actor_name,
        max_ongoing_requests=10,
        is_cross_language=False,
    )
    return RunningReplica(info), actor_handle


def _dummy_request_metadata() -> RequestMetadata:
    return RequestMetadata(request_id="abc", internal_request_id="def")


@pytest.mark.asyncio
async def test_reserve_slot_cancellation_releases_slot_on_actor(ray_instance):
    """If the awaiting reserve_slot task is cancelled, the wrapper must fire a
    follow-up release_slot.remote(token) so the actor doesn't leak the slot.
    """
    signal = SignalActor.remote()
    replica, actor = _spawn_running_replica(
        BlockingReserveActor, "blocking-replica", signal
    )

    task = get_or_create_event_loop().create_task(
        replica.reserve_slot(_dummy_request_metadata())
    )

    # Let the actor enter reserve_slot and start awaiting the signal.
    _, pending = await asyncio.wait([task], timeout=0.5)
    assert len(pending) == 1

    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task

    # Unblock the actor so it can process the follow-up release_slot.remote().
    await signal.send.remote()

    # The wrapper's cancellation cleanup fires release_slot.remote(token)
    # without awaiting it; wait until the actor records the call.
    async def _release_received():
        return bool(await actor.get_released_tokens.remote())

    await async_wait_for_condition(_release_received, timeout=5)

    released_tokens = await actor.get_released_tokens.remote()
    assert len(released_tokens) == 1


@pytest.mark.asyncio
async def test_reserve_slot_propagates_actor_died_error(ray_instance):
    """If the replica actor is dead, RunningReplica.reserve_slot must raise
    ActorDiedError so AsyncioRouter.choose_replica can retry against another
    replica. ActorUnavailableError is also acceptable on the brief window
    before the actor failure has propagated.
    """
    replica, actor = _spawn_running_replica(
        SlotReservationActor, "doomed-replica", max_ongoing_requests=1
    )

    # Confirm liveness via a successful reservation first.
    _, info = await replica.reserve_slot(_dummy_request_metadata())
    assert info.accepted

    ray.kill(actor)

    with pytest.raises((ActorDiedError, ActorUnavailableError)):
        await replica.reserve_slot(_dummy_request_metadata())


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