import asyncio
import datetime
import glob
import os
import random
import socket
import threading
import time
from contextlib import asynccontextmanager
from copy import copy, deepcopy
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
from unittest.mock import Mock

import grpc
import httpx
import requests
from starlette.requests import Request

import ray
from ray import serve
from ray._common.network_utils import build_address
from ray._common.test_utils import (
    PrometheusTimeseries,
    fetch_prometheus_metric_timeseries,
    wait_for_condition,
)
from ray._common.utils import TimerBase
from ray.actor import ActorHandle
from ray.serve._private.client import ServeControllerClient
from ray.serve._private.common import (
    CreatePlacementGroupRequest,
    DeploymentID,
    DeploymentStatus,
    ReplicaID,
    RequestProtocol,
    RunningReplicaInfo,
)
from ray.serve._private.constants import (
    RAY_SERVE_ENABLE_HA_PROXY,
    RAY_SERVE_HAPROXY_SOCKET_PATH,
    SERVE_DEFAULT_APP_NAME,
    SERVE_NAMESPACE,
)
from ray.serve._private.deployment_info import DeploymentInfo
from ray.serve._private.deployment_scheduler import ReplicaSchedulingRequest
from ray.serve._private.deployment_state import (
    ALL_REPLICA_STATES,
    DeploymentVersion,
    ReplicaStartupStatus,
    ReplicaState,
)
from ray.serve._private.haproxy import HAProxyApi
from ray.serve._private.proxy import DRAINING_MESSAGE
from ray.serve._private.replica_result import ReplicaResult
from ray.serve._private.request_router import (
    PendingRequest,
    RunningReplica,
)
from ray.serve._private.usage import ServeUsageTag
from ray.serve.context import _get_global_client
from ray.serve.gang import GangContext
from ray.serve.generated import serve_pb2, serve_pb2_grpc
from ray.serve.schema import ApplicationStatus, ReplicaRank, TargetGroup
from ray.util.state import list_actors

TELEMETRY_ROUTE_PREFIX = "/telemetry"
STORAGE_ACTOR_NAME = "storage"
PROMETHEUS_METRICS_TIMEOUT_S = 5


def skip_if_haproxy(reason: str):
    """Skip a test when the HAProxy ingress is enabled.

    The HAProxy ingress runs as a separate premerge step with
    RAY_SERVE_ENABLE_HA_PROXY=1. Some tests exercise behavior HAProxy does not
    support yet (e.g. gRPC ingress) or that probes the native Serve proxy
    directly. Mark those with this decorator instead of maintaining a separate
    test allowlist. The test still runs in the non-HAProxy steps.
    """
    import pytest

    return pytest.mark.skipif(
        RAY_SERVE_ENABLE_HA_PROXY, reason=f"HAProxy ingress: {reason}"
    )


# Global variable that is fetched during controller recovery that
# marks (simulates) which replicas have died since controller first
# recovered a list of live replica names.
# NOTE(zcin): This is necessary because the replica's `recover()` method
# is called in the controller's init function, instead of in the control
# loop, so we can't "mark" a replica dead through a method. This global
# state is cleared after each test that uses the fixtures in this file.
dead_replicas_context = set()
# Replicas registered in this set will report `was_initialized() == False` to
# the controller during recovery, simulating the case where a previous
# controller crashed before the actor finished its initial setup.
uninitialized_replicas_context = set()
replica_rank_context: Dict[str, ReplicaRank] = {}


class MockTimer(TimerBase):
    def __init__(self, start_time: Optional[float] = None):
        self._lock = threading.Lock()
        self.reset(start_time=start_time)

    def reset(self, start_time: Optional[float] = None):
        if start_time is None:
            start_time = time.time()
        self._curr = start_time

    def time(self) -> float:
        return self._curr

    def advance(self, by: float):
        with self._lock:
            self._curr += by

    def realistic_sleep(self, amt: float):
        with self._lock:
            self._curr += amt + 0.001


class MockAsyncTimer:
    def __init__(self, start_time: Optional[float] = 0):
        self.reset(start_time=start_time)
        self._num_sleepers = 0

    def reset(self, start_time: 0):
        self._curr = start_time

    def time(self) -> float:
        return self._curr

    async def sleep(self, amt: float):
        self._num_sleepers += 1
        end = self._curr + amt

        # Give up the event loop
        while self._curr < end:
            await asyncio.sleep(0)

        self._num_sleepers -= 1

    def advance(self, amt: float):
        self._curr += amt

    def num_sleepers(self):
        return self._num_sleepers


class MockKVStore:
    def __init__(self):
        self.store = dict()

    def put(self, key: str, val: Any) -> bool:
        if not isinstance(key, str):
            raise TypeError("key must be a string, got: {}.".format(type(key)))
        self.store[key] = val
        return True

    def get(self, key: str) -> Any:
        if not isinstance(key, str):
            raise TypeError("key must be a string, got: {}.".format(type(key)))
        return self.store.get(key, None)

    def delete(self, key: str) -> bool:
        if not isinstance(key, str):
            raise TypeError("key must be a string, got: {}.".format(type(key)))

        if key in self.store:
            del self.store[key]
            return True

        return False


class MockClusterNodeInfoCache:
    def __init__(self):
        self.alive_node_ids = set()
        self.total_resources_per_node = dict()
        self.available_resources_per_node = dict()
        self.draining_nodes = dict()
        self.node_labels = dict()

    def get_alive_node_ids(self):
        return self.alive_node_ids

    def get_draining_nodes(self):
        return self.draining_nodes

    def get_active_node_ids(self):
        return self.alive_node_ids - set(self.draining_nodes)

    def get_node_az(self, node_id):
        return None

    def get_available_resources_per_node(self):
        return self.available_resources_per_node

    def get_total_resources_per_node(self):
        return self.total_resources_per_node

    def add_node(self, node_id: str, resources: Dict = None, labels: Dict = None):
        self.alive_node_ids.add(node_id)
        self.total_resources_per_node[node_id] = deepcopy(resources) or {}
        self.available_resources_per_node[node_id] = deepcopy(resources) or {}
        self.node_labels[node_id] = labels or {}

    def set_available_resources_per_node(self, node_id: str, resources: Dict):
        self.available_resources_per_node[node_id] = deepcopy(resources)

    def get_node_labels(self, node_id: str):
        return self.node_labels.get(node_id, {})


# Default value used by FakeRunningReplica when the test doesn't override.
FAKE_REPLICA_DEFAULT_MAX_ONGOING_REQUESTS = 10
# Every FakeRunningReplica is stamped with this deployment id. Router
# tests never care about app/name specifically; cross-test consistency
# matters more than differentiation.
FAKE_REPLICA_DEPLOYMENT_ID = DeploymentID(name="TEST_DEPLOYMENT")


class FakeRunningReplica(RunningReplica):
    """In-memory ``RunningReplica`` stand-in for request-router unit tests.

    Shared between ``test_pow_2_request_router.py`` and
    ``test_consistent_hash_router.py``. Supports enough knobs for pow-2's
    locality / cancellation / re-probe tests; consistent-hash only uses
    the queue-length response hooks.
    """

    def __init__(
        self,
        replica_unique_id: str,
        *,
        node_id: str = "",
        availability_zone: Optional[str] = None,
        reset_after_response: bool = False,
        model_ids: Optional[Set[str]] = None,
        sleep_time_s: float = 0.0,
        max_ongoing_requests: int = FAKE_REPLICA_DEFAULT_MAX_ONGOING_REQUESTS,
    ):
        self._replica_id = ReplicaID(
            unique_id=replica_unique_id,
            deployment_id=FAKE_REPLICA_DEPLOYMENT_ID,
        )
        self._node_id = node_id
        self._availability_zone = availability_zone
        self._queue_len = 0
        self._max_ongoing_requests = max_ongoing_requests
        self._has_queue_len_response = asyncio.Event()
        self._reset_after_response = reset_after_response
        self._model_ids = model_ids or set()
        self._sleep_time_s = sleep_time_s
        self._exception: Optional[Exception] = None

        self.get_queue_len_was_cancelled = False
        self.queue_len_deadline_history = list()
        self.num_get_queue_len_calls = 0

    @property
    def replica_id(self) -> ReplicaID:
        return self._replica_id

    @property
    def node_id(self) -> str:
        return self._node_id

    @property
    def availability_zone(self) -> Optional[str]:
        return self._availability_zone

    @property
    def multiplexed_model_ids(self) -> Set[str]:
        return self._model_ids

    def update_replica_info(self, replica_info: RunningReplicaInfo) -> None:
        self._model_ids = set(replica_info.multiplexed_model_ids)

    @property
    def max_ongoing_requests(self) -> int:
        return self._max_ongoing_requests

    def set_queue_len_response(
        self,
        queue_len: int,
        exception: Optional[Exception] = None,
    ):
        self._queue_len = queue_len
        self._exception = exception
        self._has_queue_len_response.set()

    def push_proxy_handle(self, handle: ActorHandle):
        pass

    async def get_queue_len(self, *, deadline_s: float) -> int:
        self.num_get_queue_len_calls += 1
        self.queue_len_deadline_history.append(deadline_s)
        try:
            while not self._has_queue_len_response.is_set():
                await self._has_queue_len_response.wait()

            if self._sleep_time_s > 0:
                await asyncio.sleep(self._sleep_time_s)

            if self._reset_after_response:
                self._has_queue_len_response.clear()

            if self._exception is not None:
                raise self._exception

            return self._queue_len
        except asyncio.CancelledError:
            self.get_queue_len_was_cancelled = True
            raise

    def try_send_request(
        self, pr: PendingRequest, with_rejection: bool
    ) -> ReplicaResult:
        raise NotImplementedError()

    def send_request_with_rejection(self, pr: PendingRequest) -> ReplicaResult:
        raise NotImplementedError()


class FakeRemoteFunction:
    def remote(self):
        pass


class MockActorHandle:
    def __init__(self, **kwargs):
        self._options = kwargs
        self._actor_id = "fake_id"
        self.initialize_and_get_metadata_called = False
        self.is_allocated_called = False

    @property
    def initialize_and_get_metadata(self):
        self.initialize_and_get_metadata_called = True
        # return a mock object so that we can call `remote()` on it.
        return FakeRemoteFunction()

    @property
    def is_allocated(self):
        self.is_allocated_called = True
        return FakeRemoteFunction()


class MockActorClass:
    def __init__(self):
        self._init_args = ()
        self._options = dict()

    def options(self, **kwargs):
        res = copy(self)

        for k, v in kwargs.items():
            res._options[k] = v

        return res

    def remote(self, *args) -> MockActorHandle:
        return MockActorHandle(init_args=args, **self._options)


class MockPlacementGroup:
    def __init__(self, request: CreatePlacementGroupRequest):
        self._bundles = request.bundles
        self._strategy = request.strategy
        self._soft_target_node_id = request.target_node_id
        self._name = request.name
        self._lifetime = "detached"


class MockDeploymentHandle:
    def __init__(self, deployment_name: str, app_name: str = SERVE_DEFAULT_APP_NAME):
        self._deployment_name = deployment_name
        self._app_name = app_name
        self._protocol = RequestProtocol.UNDEFINED
        self._running_replicas_populated = False
        self._initialized = False

    def is_initialized(self):
        return self._initialized

    def _init(self):
        if self._initialized:
            raise RuntimeError("already initialized")

        self._initialized = True

    def options(self, *args, **kwargs):
        return self

    def __eq__(self, dep: Tuple[str]):
        other_deployment_name, other_app_name = dep
        return (
            self._deployment_name == other_deployment_name
            and self._app_name == other_app_name
        )

    def _set_request_protocol(self, protocol: RequestProtocol):
        self._protocol = protocol

    def _get_or_create_router(self):
        pass

    def running_replicas_populated(self) -> bool:
        return self._running_replicas_populated

    def set_running_replicas_populated(self, val: bool):
        self._running_replicas_populated = val


class MockDeploymentActorWrapper:
    """Mock for DeploymentActorWrapper with per-instance setters."""

    def __init__(
        self,
        deployment_id: DeploymentID,
        config,
        code_version: str,
        recovered_handle=None,
    ):
        self._deployment_id = deployment_id
        self._config = config
        self._code_version = code_version
        self._handle = recovered_handle
        self._ready = False  # Recovered starts pending until set_ready()
        self._start_error_msg: Optional[str] = None
        self._check_ready_error_msg: Optional[str] = None
        self.killed = False
        self.pending_killed = False
        self._start_fail_count = 0
        self._start_fail_msg = None
        self._health_ok = True
        self.reset_health_state_after_running_count = 0

    @property
    def actor_logical_name(self) -> str:
        return self._config.name

    @property
    def code_version(self) -> str:
        return self._code_version

    def set_ready(self):
        self._ready = True

    def set_start_error(self, error_msg: str):
        """Make start() fail with this error (persists until cleared)."""
        self._start_error_msg = error_msg

    def set_check_ready_error(self, error_msg: str):
        self._check_ready_error_msg = error_msg

    def set_failed_to_start(self, error_msg: str = "constructor failed"):
        """Simulate start failure (like replica's set_failed_to_start). check_ready will return this error."""
        self._check_ready_error_msg = error_msg

    def start(self, deployment_runtime_env=None):
        if self._start_error_msg is not None:
            return False, self._start_error_msg
        # Match real behavior: created actor starts as pending until ready.
        return True, None

    def check_ready(self):
        if self._check_ready_error_msg is not None:
            return False, self._check_ready_error_msg
        if self._ready:
            return True, None
        return False, None

    def set_health_ok(self, ok: bool) -> None:
        """If False, ``check_health`` returns False (failed health reconciliation)."""
        self._health_ok = ok

    def reset_health_state_after_running(self) -> None:
        """Match DeploymentActorWrapper; increments counter for unit tests."""
        self.reset_health_state_after_running_count += 1

    def check_health(self) -> bool:
        """Match DeploymentActorWrapper; controlled via ``set_health_ok``."""
        return self._health_ok

    def kill(self) -> None:
        self.killed = True


class MockReplicaActorWrapper:
    def __init__(
        self,
        replica_id: ReplicaID,
        version: DeploymentVersion,
    ):
        self._replica_id = replica_id
        self._actor_name = replica_id.to_full_id_str()
        # Will be set when `start()` is called.
        self.started = False
        # Will be set when `recover()` is called.
        self.recovering = False
        # Will be set when `start()` is called.
        self.version = version
        # Initial state for a replica is PENDING_ALLOCATION.
        self.status = ReplicaStartupStatus.PENDING_ALLOCATION
        # Will be set when `graceful_stop()` is called.
        self.stopped = False
        # Expected to be set in the test.
        self.done_stopping = False
        # Will be set when `force_stop()` is called.
        self.force_stopped_counter = 0
        # Will be set when `check_health()` is called.
        self.health_check_called = False
        # Returned by the health check.
        self.healthy = True
        self._is_cross_language = False
        self._actor_handle = MockActorHandle()
        self._node_id = None
        self._node_ip = None
        self._node_instance_id = None
        self._node_id_is_set = False
        self._log_file_path = None
        self._actor_id = None
        self._internal_grpc_port = None
        self._http_port = None
        self._pg_bundles = None
        self._initialization_latency_s = -1
        self._docs_path = None
        self._rank = replica_rank_context.get(replica_id.unique_id, None)
        self._assign_rank_callback = None
        self._ingress = False
        self._gang_context = None
        self._gang_pg_index = None
        self._unrecoverable = False

    @property
    def is_cross_language(self) -> bool:
        return self._is_cross_language

    @property
    def replica_id(self) -> ReplicaID:
        return self._replica_id

    @property
    def deployment_name(self) -> str:
        return self._replica_id.deployment_id.name

    @property
    def actor_handle(self) -> MockActorHandle:
        return self._actor_handle

    @property
    def max_ongoing_requests(self) -> int:
        return self.version.deployment_config.max_ongoing_requests

    @property
    def graceful_shutdown_timeout_s(self) -> float:
        return self.version.deployment_config.graceful_shutdown_timeout_s

    @property
    def health_check_period_s(self) -> float:
        return self.version.deployment_config.health_check_period_s

    @property
    def health_check_timeout_s(self) -> float:
        return self.version.deployment_config.health_check_timeout_s

    @property
    def pid(self) -> Optional[int]:
        return None

    @property
    def actor_id(self) -> Optional[str]:
        return self._actor_id

    @property
    def worker_id(self) -> Optional[str]:
        return None

    @property
    def node_id(self) -> Optional[str]:
        if self._node_id_is_set:
            return self._node_id
        if self.status == ReplicaStartupStatus.SUCCEEDED or self.started:
            return "node-id"
        return None

    @property
    def availability_zone(self) -> Optional[str]:
        return None

    @property
    def node_ip(self) -> Optional[str]:
        return None

    @property
    def node_instance_id(self) -> Optional[str]:
        return None

    @property
    def log_file_path(self) -> Optional[str]:
        return self._log_file_path

    @property
    def grpc_port(self) -> Optional[int]:
        return None

    @property
    def placement_group_bundles(self) -> Optional[List[Dict[str, float]]]:
        return None

    @property
    def initialization_latency_s(self) -> float:
        return self._initialization_latency_s

    @property
    def reconfigure_start_time(self) -> Optional[float]:
        return None

    @property
    def last_health_check_latency_ms(self) -> Optional[float]:
        return None

    @property
    def last_health_check_failed(self) -> bool:
        return False

    def set_docs_path(self, docs_path: str):
        self._docs_path = docs_path

    @property
    def docs_path(self) -> Optional[str]:
        return self._docs_path

    def set_status(self, status: ReplicaStartupStatus):
        self.status = status

    def set_ready(self, version: DeploymentVersion = None):
        self.status = ReplicaStartupStatus.SUCCEEDED
        # Mirror the real actor: a started replica has allocated a log file.
        self._log_file_path = "serve/replica.log"
        if version:
            self.version_to_be_fetched_from_actor = version
        else:
            self.version_to_be_fetched_from_actor = self.version

    def set_failed_to_start(self):
        self.status = ReplicaStartupStatus.FAILED

    def set_done_stopping(self):
        self.done_stopping = True

    def set_unhealthy(self):
        self.healthy = False

    def set_starting_version(self, version: DeploymentVersion):
        """Mocked deployment_worker return version from reconfigure()"""
        self.starting_version = version

    def set_node_id(self, node_id: str):
        self._node_id = node_id
        self._node_id_is_set = True

    def set_actor_id(self, actor_id: str):
        self._actor_id = actor_id

    def start(
        self,
        deployment_info: DeploymentInfo,
        assign_rank_callback: Callable[[ReplicaID], ReplicaRank],
        gang_placement_group=None,
        gang_pg_index=None,
        gang_context=None,
    ):
        self.started = True
        self._gang_context = gang_context
        self._gang_pg_index = gang_pg_index
        self._assign_rank_callback = assign_rank_callback
        self._rank = assign_rank_callback(self._replica_id.unique_id, node_id=-1)
        replica_rank_context[self._replica_id.unique_id] = self._rank

        def _on_scheduled_stub(*args, **kwargs):
            pass

        return ReplicaSchedulingRequest(
            replica_id=self._replica_id,
            actor_def=Mock(),
            actor_resources={},
            actor_options={"name": "placeholder"},
            actor_init_args=(),
            placement_group_bundles=(
                deployment_info.replica_config.placement_group_bundles
            ),
            on_scheduled=_on_scheduled_stub,
            gang_placement_group=gang_placement_group,
            gang_pg_index=gang_pg_index,
        )

    @property
    def rank(self) -> Optional[ReplicaRank]:
        return self._rank

    def reconfigure(
        self,
        version: DeploymentVersion,
        rank: ReplicaRank = None,
    ):
        self.started = True
        updating = self.version.requires_actor_reconfigure(version)
        self.version = version
        self._rank = rank
        replica_rank_context[self._replica_id.unique_id] = rank
        return updating

    def recover(self, ingress: bool = False):
        if self.replica_id in dead_replicas_context:
            return False

        self._ingress = ingress
        self.recovering = True
        self.started = False
        self._rank = replica_rank_context.get(self._replica_id.unique_id, None)
        # Mirror production: the `was_initialized` probe is fired here and
        # observed asynchronously in `check_ready()`. Tests register
        # uninitialized replicas via `uninitialized_replicas_context`.
        self._unrecoverable = self.replica_id in uninitialized_replicas_context
        return True

    def check_ready(self) -> ReplicaStartupStatus:
        # If the controller's async `was_initialized` probe came back False,
        # report a failed-but-unrecoverable startup so the reconciler drops
        # and replaces the replica without recording a deploy failure.
        if self.recovering and self._unrecoverable:
            return (
                ReplicaStartupStatus.FAILED,
                f"{self._replica_id} was found alive but never finished "
                "its initial setup.",
            )

        ready = self.status
        self.status = ReplicaStartupStatus.PENDING_INITIALIZATION
        if ready == ReplicaStartupStatus.SUCCEEDED and self.recovering:
            self.recovering = False
            self.started = True
            self.version = self.version_to_be_fetched_from_actor
        return ready, None

    def resource_requirements(self) -> Tuple[str, str]:
        assert self.started
        return str({"REQUIRED_RESOURCE": 1.0}), str({"AVAILABLE_RESOURCE": 1.0})

    @property
    def actor_resources(self) -> Dict[str, float]:
        return {"CPU": 0.1}

    @property
    def available_resources(self) -> Dict[str, float]:
        # Only used to print a warning.
        return {}

    def graceful_stop(self) -> None:
        # `started` is only set after a successful `check_ready` transition
        # to RUNNING. A replica force-stopped while still RECOVERING (e.g.,
        # because the `was_initialized` probe failed) is a legitimate stop.
        assert self.started or self.recovering
        self.stopped = True
        return self.graceful_shutdown_timeout_s

    def check_stopped(self) -> bool:
        return self.done_stopping

    def force_stop(self, log_shutdown_message: bool = False):
        self.force_stopped_counter += 1

    def check_health(self):
        self.health_check_called = True
        return self.healthy

    def get_routing_stats(self) -> Dict[str, Any]:
        return {}

    def get_outbound_deployments(self) -> Optional[List[DeploymentID]]:
        return getattr(self, "_outbound_deployments", None)

    @property
    def route_patterns(self) -> Optional[List[str]]:
        return None

    @property
    def gang_context(self) -> Optional[GangContext]:
        return self._gang_context

    @property
    def unrecoverable(self) -> bool:
        return self._unrecoverable


@serve.deployment
class GetPID:
    def __call__(self):
        return os.getpid()


get_pid_entrypoint = GetPID.bind()


def check_ray_stopped():
    try:
        requests.get("http://localhost:8265/api/ray/version")
        return False
    except Exception:
        return True


def check_ray_started():
    return requests.get("http://localhost:8265/api/ray/version").status_code == 200


def check_deployment_status(
    name: str, expected_status: DeploymentStatus, app_name=SERVE_DEFAULT_APP_NAME
) -> bool:
    app_status = serve.status().applications[app_name]
    assert app_status.deployments[name].status == expected_status
    return True


def get_num_alive_replicas(
    deployment_name: str, app_name: str = SERVE_DEFAULT_APP_NAME
) -> int:
    """Get the replicas currently running for the given deployment."""

    dep_id = DeploymentID(name=deployment_name, app_name=app_name)
    actors = list_actors(
        filters=[
            ("class_name", "=", dep_id.to_replica_actor_class_name()),
            ("state", "=", "ALIVE"),
        ]
    )
    return len(actors)


def expected_proxy_actors(num_proxy_nodes: int = 1) -> Dict[str, int]:
    """Proxy actors expected ALIVE by class name for the given number of proxy nodes.

    Natively each proxy node runs one ProxyActor. Under HAProxy each proxy node runs an
    HAProxyManager and the head node also runs a single fallback ProxyActor. Set-based
    callers can take the keys, count-based callers use the counts directly.
    """
    if RAY_SERVE_ENABLE_HA_PROXY:
        return {"HAProxyManager": num_proxy_nodes, "ProxyActor": 1}
    return {"ProxyActor": num_proxy_nodes}


def wait_for_haproxy_routing_to_replica(timeout: int = 30):
    """Block until HAProxy marks a replica's primary backend server UP.

    serve.run returns once replicas are RUNNING, but until a replica's
    direct-ingress server passes HAProxy's health check, HAProxy serves requests
    from the head-node backup fallback proxy. That path adds proxy and router
    spans, so trace-topology assertions must wait for direct routing first. No-op
    without HAProxy, which has no fallback to race.
    """
    if not RAY_SERVE_ENABLE_HA_PROXY:
        return

    socket_glob = os.path.join(
        os.path.dirname(RAY_SERVE_HAPROXY_SOCKET_PATH), "*", "admin.sock"
    )

    def replica_backend_up():
        for sock_path in glob.glob(socket_glob):
            try:
                with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
                    client.settimeout(5.0)
                    client.connect(sock_path)
                    client.sendall(b"show stat\n")
                    data = b""
                    while chunk := client.recv(65536):
                        data += chunk
            except OSError:
                continue  # stale socket from a prior run
            stats = HAProxyApi._parse_haproxy_csv_stats(data.decode(errors="replace"))
            if any(
                name.startswith("SERVE_REPLICA") and server.is_up
                for servers in stats.values()
                for name, server in servers.items()
            ):
                return True
        return False

    wait_for_condition(replica_backend_up, timeout=timeout)


def alive_actor_counts() -> Dict[str, int]:
    """Count of ALIVE actors by class name in the current Ray session."""
    counts: Dict[str, int] = {}
    for actor in list_actors(filters=[("STATE", "=", "ALIVE")]):
        counts[actor["class_name"]] = counts.get(actor["class_name"], 0) + 1
    return counts


def check_num_replicas_gte(
    name: str, target: int, app_name: str = SERVE_DEFAULT_APP_NAME
) -> int:
    """Check if num replicas is >= target."""

    assert get_num_alive_replicas(name, app_name) >= target
    return True


def check_num_replicas_eq(
    name: str,
    target: int,
    app_name: str = SERVE_DEFAULT_APP_NAME,
    use_controller: bool = False,
) -> bool:
    """Check if num replicas is == target."""

    if use_controller:
        dep = serve.status().applications[app_name].deployments[name]
        num_running_replicas = dep.replica_states.get(ReplicaState.RUNNING, 0)
        assert num_running_replicas == target
    else:
        assert get_num_alive_replicas(name, app_name) == target

    return True


def check_num_replicas_lte(
    name: str, target: int, app_name: str = SERVE_DEFAULT_APP_NAME
) -> int:
    """Check if num replicas is <= target."""

    assert get_num_alive_replicas(name, app_name) <= target
    return True


def check_apps_running(apps: List):
    status = serve.status()

    for app_name in apps:
        assert status.applications[app_name].status == ApplicationStatus.RUNNING

    return True


def check_replica_counts(
    controller: ActorHandle,
    deployment_id: DeploymentID,
    total: Optional[int] = None,
    by_state: Optional[List[Tuple[ReplicaState, int, Callable]]] = None,
):
    """Uses _dump_replica_states_for_testing to check replica counts.

    Args:
        controller: A handle to the Serve controller.
        deployment_id: The deployment to check replica counts for.
        total: The total number of expected replicas for the deployment.
        by_state: A list of tuples of the form
            (replica state, number of replicas, filter function).
            Used for more fine grained checks.

    Returns:
        True when all assertions pass (raises ``AssertionError`` otherwise).
    """
    replicas = ray.get(
        controller._dump_replica_states_for_testing.remote(deployment_id)
    )

    if total is not None:
        replica_counts = {
            state: len(replicas.get([state]))
            for state in ALL_REPLICA_STATES
            if replicas.get([state])
        }
        assert replicas.count() == total, replica_counts

    if by_state is not None:
        for state, count, check in by_state:
            assert isinstance(state, ReplicaState)
            assert isinstance(count, int) and count >= 0
            if check:
                filtered = {r for r in replicas.get(states=[state]) if check(r)}
                curr_count = len(filtered)
            else:
                curr_count = replicas.count(states=[state])
            msg = f"Expected {count} for state {state} but got {curr_count}."
            assert curr_count == count, msg

    return True


@ray.remote(name=STORAGE_ACTOR_NAME, namespace=SERVE_NAMESPACE, num_cpus=0)
class TelemetryStorage:
    def __init__(self):
        self.reports_received = 0
        self.current_report = dict()

    def store_report(self, report: Dict) -> None:
        self.reports_received += 1
        self.current_report = report

    def get_report(self) -> Dict:
        return self.current_report

    def get_reports_received(self) -> int:
        return self.reports_received


@serve.deployment(ray_actor_options={"num_cpus": 0})
class TelemetryReceiver:
    def __init__(self):
        self.storage = ray.get_actor(name=STORAGE_ACTOR_NAME, namespace=SERVE_NAMESPACE)

    async def __call__(self, request: Request) -> bool:
        report = await request.json()
        ray.get(self.storage.store_report.remote(report))
        return True


receiver_app = TelemetryReceiver.bind()


def start_telemetry_app():
    """Start a telemetry Serve app.

    Ray should be initialized before calling this method.

    NOTE: If you're running the TelemetryReceiver Serve app to check telemetry,
    remember that the receiver itself is counted in the telemetry. E.g. if you
    deploy a Serve app other than the receiver, the number of apps in the
    cluster is 2- not 1– since the receiver is also running.

    Returns a handle to a TelemetryStorage actor. You can use this actor
    to access the latest telemetry reports.
    """

    storage = TelemetryStorage.remote()
    serve.run(receiver_app, name="telemetry", route_prefix=TELEMETRY_ROUTE_PREFIX)
    return storage


def check_telemetry(
    tag: ServeUsageTag, expected: Any, storage_actor_name: str = STORAGE_ACTOR_NAME
):
    storage_handle = ray.get_actor(storage_actor_name, namespace=SERVE_NAMESPACE)
    report = ray.get(storage_handle.get_report.remote())
    print(report["extra_usage_tags"])
    assert tag.get_value_from_report(report) == expected
    return True


def ping_grpc_list_applications(channel, app_names, test_draining=False):
    import pytest

    stub = serve_pb2_grpc.RayServeAPIServiceStub(channel)
    request = serve_pb2.ListApplicationsRequest()
    if test_draining:
        with pytest.raises(grpc.RpcError) as exception_info:
            _, _ = stub.ListApplications.with_call(request=request)
        rpc_error = exception_info.value
        assert rpc_error.code() == grpc.StatusCode.UNAVAILABLE
        assert rpc_error.details() == DRAINING_MESSAGE
    else:
        response, call = stub.ListApplications.with_call(request=request)
        assert call.code() == grpc.StatusCode.OK
        assert response.application_names == app_names
    return True


def ping_grpc_healthz(channel, test_draining=False):
    import pytest

    stub = serve_pb2_grpc.RayServeAPIServiceStub(channel)
    request = serve_pb2.HealthzRequest()
    if test_draining:
        with pytest.raises(grpc.RpcError) as exception_info:
            _, _ = stub.Healthz.with_call(request=request)
        rpc_error = exception_info.value
        assert rpc_error.code() == grpc.StatusCode.UNAVAILABLE
        assert rpc_error.details() == DRAINING_MESSAGE
    else:
        response, call = stub.Healthz.with_call(request=request)
        assert call.code() == grpc.StatusCode.OK
        if not RAY_SERVE_ENABLE_HA_PROXY:
            assert response.message == "success"
    return True


def ping_grpc_call_method(channel, app_name, test_not_found=False):
    import pytest

    stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
    request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
    metadata = (("application", app_name),)
    if test_not_found:
        with pytest.raises(grpc.RpcError) as exception_info:
            _, _ = stub.__call__.with_call(request=request, metadata=metadata)
        rpc_error = exception_info.value
        assert rpc_error.code() == grpc.StatusCode.NOT_FOUND, rpc_error.code()
        assert f"Application '{app_name}' not found." in rpc_error.details()
    else:
        response, call = stub.__call__.with_call(request=request, metadata=metadata)
        assert call.code() == grpc.StatusCode.OK, call.code()
        assert response.greeting == "Hello foo from bar", response.greeting


def ping_grpc_another_method(channel, app_name):
    stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
    request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
    metadata = (("application", app_name),)
    response = stub.Method1(request=request, metadata=metadata)
    assert response.greeting == "Hello foo from method1"


def ping_grpc_model_multiplexing(channel, app_name):
    stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
    request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
    multiplexed_model_id = "999"
    metadata = (
        ("application", app_name),
        ("multiplexed_model_id", multiplexed_model_id),
    )
    response = stub.Method2(request=request, metadata=metadata)
    assert (
        response.greeting
        == f"Method2 called model, loading model: {multiplexed_model_id}"
    )


def ping_grpc_streaming(channel, app_name):
    stub = serve_pb2_grpc.UserDefinedServiceStub(channel)
    request = serve_pb2.UserDefinedMessage(name="foo", num=30, foo="bar")
    metadata = (("application", app_name),)
    responses = stub.Streaming(request=request, metadata=metadata)
    for idx, response in enumerate(responses):
        assert response.greeting == f"{idx}: Hello foo from bar"


def ping_fruit_stand(channel, app_name):
    stub = serve_pb2_grpc.FruitServiceStub(channel)
    request = serve_pb2.FruitAmounts(orange=4, apple=8)
    metadata = (("application", app_name),)
    response = stub.FruitStand(request=request, metadata=metadata)
    assert response.costs == 32


@asynccontextmanager
async def send_signal_on_cancellation(signal_actor: ActorHandle):
    cancelled = False
    try:
        yield
        await asyncio.sleep(100)
    except asyncio.CancelledError:
        cancelled = True
        # Clear the context var to avoid Ray recursively cancelling this method call.
        ray._raylet.async_task_id.set(None)
        await signal_actor.send.remote()

    if not cancelled:
        raise RuntimeError(
            "CancelledError wasn't raised during `send_signal_on_cancellation` block"
        )


class FakeGrpcContext:
    def __init__(self):
        self._auth_context = {"key": "value"}
        self._invocation_metadata = [("key", "value")]
        self._peer = "peer"
        self._peer_identities = b"peer_identities"
        self._peer_identity_key = "peer_identity_key"
        self._code = None
        self._details = None
        self._trailing_metadata = []
        self._invocation_metadata = []

    def auth_context(self):
        return self._auth_context

    def code(self):
        return self._code

    def details(self):
        return self._details

    def peer(self):
        return self._peer

    def peer_identities(self):
        return self._peer_identities

    def peer_identity_key(self):
        return self._peer_identity_key

    def trailing_metadata(self):
        return self._trailing_metadata

    def set_code(self, code):
        self._code = code

    def set_details(self, details):
        self._details = details

    def set_trailing_metadata(self, trailing_metadata):
        self._trailing_metadata = trailing_metadata

    def invocation_metadata(self):
        return self._invocation_metadata


class FakeGauge:
    def __init__(self, name: str = None, tag_keys: Tuple[str] = None):
        self.name = name
        self.values = dict()

        self.tags = tag_keys or ()
        self.default_tags = dict()

    def set_default_tags(self, tags: Dict[str, str]):
        for key, tag in tags.items():
            assert key in self.tags
            self.default_tags[key] = tag

    def set(self, value: Union[int, float], tags: Dict[str, str] = None):
        merged_tags = self.default_tags.copy()
        merged_tags.update(tags or {})
        assert set(merged_tags.keys()) == set(self.tags)

        d = self.values
        for tag in self.tags[:-1]:
            tag_value = merged_tags[tag]
            if tag_value not in d:
                d[tag_value] = dict()
            d = d[tag_value]

        d[merged_tags[self.tags[-1]]] = value

    def get_value(self, tags: Dict[str, str]):
        value = self.values
        for tag in self.tags:
            tag_value = tags[tag]
            value = value.get(tag_value)
            if value is None:
                return

        return value


class FakeCounter:
    def __init__(self, name: str = None, tag_keys: Tuple[str] = None):
        self.name = name
        self.counts = dict()

        self.tags = tag_keys or ()
        self.default_tags = dict()

    def set_default_tags(self, tags: Dict[str, str]):
        for key, tag in tags.items():
            assert key in self.tags
            self.default_tags[key] = tag

    def inc(self, value: Union[int, float] = 1.0, tags: Dict[str, str] = None):
        merged_tags = self.default_tags.copy()
        merged_tags.update(tags or {})
        assert set(merged_tags.keys()) == set(self.tags)

        d = self.counts
        for tag in self.tags[:-1]:
            tag_value = merged_tags[tag]
            if tag_value not in d:
                d[tag_value] = dict()
            d = d[tag_value]

        key = merged_tags[self.tags[-1]]
        d[key] = d.get(key, 0) + value

    def get_count(self, tags: Dict[str, str]) -> int:
        value = self.counts
        for tag in self.tags:
            tag_value = tags[tag]
            value = value.get(tag_value)
            if value is None:
                return

        return value

    def get_tags(self):
        return self.tags


@ray.remote
def get_node_id():
    return ray.get_runtime_context().get_node_id()


def check_num_alive_nodes(target: int):
    alive_nodes = [node for node in ray.nodes() if node["Alive"]]
    assert len(alive_nodes) == target
    return True


def get_deployment_details(
    deployment_name: str,
    app_name: str = SERVE_DEFAULT_APP_NAME,
    _client: ServeControllerClient = None,
):
    client = _client or _get_global_client()
    details = client.get_serve_details()
    return details["applications"][app_name]["deployments"][deployment_name]


@ray.remote(num_cpus=0)
class Barrier:
    """Block until n participants have called wait()."""

    def __init__(self, n: int):
        self.n = n
        self.count = 0
        self.event = asyncio.Event()

    async def wait(self):
        self.count += 1
        if self.count >= self.n:
            self.event.set()
        else:
            await self.event.wait()

    def get_count(self) -> int:
        return self.count


@ray.remote(num_cpus=0)
class Accumulator:
    """Collect items from multiple actors/replicas."""

    def __init__(self):
        self.items = []

    def add(self, item):
        self.items.append(item)

    def get(self) -> list:
        return list(self.items)

    def count(self) -> int:
        return len(self.items)


@ray.remote(num_cpus=0)
class FailedReplicaStore:
    """Controls replica constructor failure behavior for constructor failure tests.

    Behavior is determined by ``fail_first``:
      - ``fail_first=True``  (transient): first caller fails, rest succeed.
      - ``fail_first=False`` (partial):   first caller succeeds, rest fail.

    All decisions are made in a single atomic actor call to avoid races
    between concurrent replicas.
    """

    def __init__(self, fail_first: bool = False):
        self._fail_first = fail_first
        self._first_caller = False
        self._fail_count = 0

    def get_fail_count(self) -> int:
        """Returns the number replica startup failures"""
        return self._fail_count

    def should_fail(self) -> bool:
        """Returns whether this replica should raise in its constructor."""
        if not self._first_caller:
            self._first_caller = True
            result = self._fail_first
        else:
            result = not self._fail_first
        if result:
            self._fail_count += 1
        return result


@ray.remote(num_cpus=0)
class FailedGangReplicaStore:
    """
    Controls replica constructor failure behavior for gang scheduling tests.
        - The first gang seen is allowed to run.
        - The next gang has first replica flagged to fail.
        - Retry gangs after the first failed gang have first replica flagged.
    """

    def __init__(self):
        self._good_gang_chosen = False
        self._successful_gang_id = None
        self._failed_gang_ids = set()

    def mark_first_failing_gang(self, gang_id: str) -> bool:
        """Picks the good gang on the first call and flags the first replica of the next gang to fail."""
        if not self._good_gang_chosen:
            self._good_gang_chosen = True
            self._successful_gang_id = gang_id
            return False
        if gang_id == self._successful_gang_id:
            return False
        if len(self._failed_gang_ids) == 0:
            self._failed_gang_ids.add(gang_id)
            return True
        return False

    def mark_retry_failing_gang(self, gang_id: str) -> bool:
        """Flags the first replica of each new retry gang.
        This is called after ``mark_first_failing_gang``."""
        if gang_id == self._successful_gang_id:
            return False
        if gang_id not in self._failed_gang_ids:
            self._failed_gang_ids.add(gang_id)
            return True
        return False

    def get_failed_gang_count(self) -> int:
        """Returns the number of distinct gangs that failed."""
        return len(self._failed_gang_ids)


@ray.remote(num_cpus=0)
class SharedFlag:
    """Non-blocking boolean flag shared across actors."""

    def __init__(self):
        self.value = False

    def set(self):
        self.value = True

    def clear(self):
        self.value = False

    def is_set(self) -> bool:
        return self.value


@ray.remote(num_cpus=0)
class SharedCounter:
    """Simple shared counter for cross-actor coordination."""

    def __init__(self):
        self.count = 0

    def inc(self) -> int:
        self.count += 1
        return self.count

    def get(self) -> int:
        return self.count


@ray.remote(num_cpus=0)
class Counter:
    def __init__(self, target: int):
        self.count = 0
        self.target = target
        self.ready_event = asyncio.Event()

    def inc(self):
        self.count += 1
        if self.count == self.target:
            self.ready_event.set()

    async def wait(self):
        await self.ready_event.wait()


def tlog(s: str, level: str = "INFO"):
    """Convenient logging method for testing."""

    now = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
    print(f"[{level}] {now} {s}")


def check_target_groups_ready(
    client: ServeControllerClient,
    app_name: str,
    protocol: Union[str, RequestProtocol] = RequestProtocol.HTTP,
):
    """Wait for target groups to be ready for the given app and protocol.

    Target groups are ready when there are at least one target for the given protocol. And it's
    possible that target groups are not ready immediately. An example is when the controller
    is recovering from a crash.
    """
    target_groups = ray.get(client._controller.get_target_groups.remote(app_name))
    target_groups = [
        target_group
        for target_group in target_groups
        if target_group.protocol == protocol
    ]
    all_targets = [
        target for target_group in target_groups for target in target_group.targets
    ]
    return len(all_targets) > 0


def get_application_urls(
    protocol: Union[str, RequestProtocol] = RequestProtocol.HTTP,
    app_name: str = SERVE_DEFAULT_APP_NAME,
    use_localhost: bool = True,
    is_websocket: bool = False,
    exclude_route_prefix: bool = False,
    from_proxy_manager: bool = False,
) -> List[str]:
    """Get the URL of the application.

    Args:
        protocol: The protocol to use for the application.
        app_name: The name of the application.
        use_localhost: Whether to use localhost instead of the IP address.
            Set to True if Serve deployments are not exposed publicly or
            for low latency benchmarking.
        is_websocket: Whether the url should be served as a websocket.
        exclude_route_prefix: The route prefix to exclude from the application.
        from_proxy_manager: Whether the caller is a proxy manager.
    Returns:
        The URLs of the application.
    """
    client = _get_global_client()
    serve_details = client.get_serve_details()
    assert (
        app_name in serve_details["applications"]
    ), f"App {app_name} not found in serve details. Use this method only when the app is known to be running."
    route_prefix = serve_details["applications"][app_name]["route_prefix"]
    # route_prefix is set to None when route_prefix value is specifically set to None
    # in the config used to deploy the app.
    if exclude_route_prefix or route_prefix is None:
        route_prefix = ""
    if isinstance(protocol, str):
        protocol = RequestProtocol(protocol)
    target_groups: List[TargetGroup] = ray.get(
        client._controller.get_target_groups.remote(app_name, from_proxy_manager)
    )
    target_groups = [
        target_group
        for target_group in target_groups
        if target_group.protocol == protocol
    ]
    if len(target_groups) == 0:
        raise ValueError(
            f"No target group found for app {app_name} with protocol {protocol} and route prefix {route_prefix}"
        )
    urls = []
    for target_group in target_groups:
        for target in target_group.targets:
            ip = "localhost" if use_localhost else target.ip
            if protocol == RequestProtocol.HTTP:
                scheme = "ws" if is_websocket else "http"
                url = f"{scheme}://{build_address(ip, target.port)}{route_prefix}"
            elif protocol == RequestProtocol.GRPC:
                if is_websocket:
                    raise ValueError(
                        "is_websocket=True is not supported with gRPC protocol."
                    )
                url = build_address(ip, target.port)
            else:
                raise ValueError(f"Unsupported protocol: {protocol}")
            url = url.rstrip("/")
            urls.append(url)
    return urls


def get_application_url(
    protocol: Union[str, RequestProtocol] = RequestProtocol.HTTP,
    app_name: str = SERVE_DEFAULT_APP_NAME,
    use_localhost: bool = True,
    is_websocket: bool = False,
    exclude_route_prefix: bool = False,
    from_proxy_manager: bool = False,
) -> str:
    """Get the URL of the application.

    Args:
        protocol: The protocol to use for the application.
        app_name: The name of the application.
        use_localhost: Whether to use localhost instead of the IP address.
            Set to True if Serve deployments are not exposed publicly or
            for low latency benchmarking.
        is_websocket: Whether the url should be served as a websocket.
        exclude_route_prefix: The route prefix to exclude from the application.
        from_proxy_manager: Whether the caller is a proxy manager.
    Returns:
        The URL of the application. If there are multiple URLs, a random one is returned.
    """
    return random.choice(
        get_application_urls(
            protocol,
            app_name,
            use_localhost,
            is_websocket,
            exclude_route_prefix,
            from_proxy_manager,
        )
    )


def check_running(app_name: str = SERVE_DEFAULT_APP_NAME):
    assert serve.status().applications[app_name].status == ApplicationStatus.RUNNING
    return True


def request_with_retries(timeout=30, app_name=SERVE_DEFAULT_APP_NAME):
    result_holder = {"resp": None}

    def _attempt() -> bool:
        try:
            url = get_application_url("HTTP", app_name=app_name)
            result_holder["resp"] = httpx.get(url, timeout=timeout)
            return True
        except (httpx.RequestError, IndexError):
            return False

    try:
        wait_for_condition(_attempt, timeout=timeout)
        return result_holder["resp"]
    except RuntimeError as e:
        # Preserve previous API by raising TimeoutError on expiry
        raise TimeoutError from e


# Metrics test utilities
TEST_METRICS_EXPORT_PORT = 9999


def get_metric_float(
    metric: str,
    expected_tags: Optional[Dict[str, str]],
    timeseries: Optional[PrometheusTimeseries] = None,
    timeout: float = 20,
) -> float:
    """Gets the float value of metric.

    If tags is specified, searched for metric with matching tags.

    Returns -1 if the metric isn't available.
    """
    if timeseries is None:
        timeseries = PrometheusTimeseries()
    samples = fetch_prometheus_metric_timeseries(
        [f"localhost:{TEST_METRICS_EXPORT_PORT}"],
        timeseries,
        timeout=timeout,
    ).get(metric, [])
    for sample in samples:
        if expected_tags.items() <= sample.labels.items():
            return sample.value
    return -1


def check_metric_float_eq(
    metric: str,
    expected: float,
    expected_tags: Optional[Dict[str, str]],
    timeseries: Optional[PrometheusTimeseries] = None,
    timeout: float = 20,
) -> bool:
    """Check if a metric's float value equals the expected value."""
    metric_value = get_metric_float(metric, expected_tags, timeseries, timeout=timeout)
    assert float(metric_value) == expected
    return True


def get_metric_dictionaries(
    name: str,
    timeout: float = 20,
    timeseries: Optional[PrometheusTimeseries] = None,
    wait: bool = True,
) -> List[Dict]:
    """Get metric samples as list of label dicts.

    Args:
        name: Metric name to fetch.
        timeout: Timeout for each fetch attempt.
        timeseries: Optional shared timeseries to populate.
        wait: If True (default), wait for metric to appear before returning.
            If False, fetch once and return immediately (possibly empty).

    Returns:
        List of metric samples as label dicts.
    """
    if timeseries is None:
        timeseries = PrometheusTimeseries()

    def metric_available() -> bool:
        prom_timeseries = fetch_prometheus_metric_timeseries(
            [f"localhost:{TEST_METRICS_EXPORT_PORT}"],
            timeseries,
            # pass timeout to fetch_prometheus_metric_timeseries
            # so the test doesn't hang on requests.get
            timeout=timeout,
        )
        assert (
            name in prom_timeseries
        ), f"Metric {name} not found. Available metrics: {list(prom_timeseries.keys())}"
        return True

    if wait:
        wait_for_condition(
            metric_available, retry_interval_ms=1000, timeout=timeout * 4
        )
    else:
        # Fetch once without asserting - allows outer wait_for_condition to retry
        fetch_prometheus_metric_timeseries(
            [f"localhost:{TEST_METRICS_EXPORT_PORT}"],
            timeseries,
            timeout=PROMETHEUS_METRICS_TIMEOUT_S,
        )

    metric_dicts = []
    for sample in timeseries.metric_samples.values():
        if sample.name == name:
            metric_dicts.append(sample.labels)

    return metric_dicts