import asyncio
import threading

import pytest

import ray
from ray import serve
from ray._common.test_utils import async_wait_for_condition, wait_for_condition
from ray.exceptions import RayError
from ray.serve._private.common import DeploymentStatus
from ray.serve._private.constants import (
    REPLICA_HEALTH_CHECK_UNHEALTHY_THRESHOLD,
    SERVE_DEFAULT_APP_NAME,
)
from ray.serve.config import GangSchedulingConfig


class Counter:
    def __init__(self):
        self._count = 0

    def get(self):
        return self._count

    def inc(self):
        self._count += 1
        return self._count

    def reset(self):
        self._count = 0


@serve.deployment(health_check_period_s=1, health_check_timeout_s=1)
class Patient:
    def __init__(self):
        self.healthy = True
        self.should_hang = False

    def check_health(self):
        if self.should_hang:
            import time

            time.sleep(10000)
        elif not self.healthy:
            raise Exception("intended to fail")

    def __call__(self, *args):
        return ray.get_runtime_context().current_actor

    def set_should_fail(self):
        self.healthy = False
        return ray.get_runtime_context().current_actor

    def set_should_hang(self):
        self.should_hang = True
        return ray.get_runtime_context().current_actor


async def check_new_actor_started(handle, original_actors):
    if not isinstance(original_actors, set):
        original_actors = {original_actors._actor_id}
    try:
        return (await handle.remote())._actor_id not in original_actors
    except RayError:
        return False


@pytest.mark.parametrize("use_class", [True, False])
def test_no_user_defined_method(serve_instance, use_class):
    """Check the default behavior when an actor crashes."""

    if use_class:

        @serve.deployment
        class A:
            def __call__(self, *args):
                return ray.get_runtime_context().current_actor

    else:

        @serve.deployment
        def A(*args):
            return ray.get_runtime_context().current_actor

    h = serve.run(A.bind())
    actor = h.remote().result()
    ray.kill(actor)

    # This would time out if we wait for multiple health check failures.
    wait_for_condition(check_new_actor_started, handle=h, original_actors=actor)


@pytest.mark.asyncio
async def test_user_defined_method_fails(serve_instance):
    h = serve.run(Patient.bind())
    actor = await h.remote()
    await h.set_should_fail.remote()

    await async_wait_for_condition(
        check_new_actor_started, handle=h, original_actors=actor
    )
    await asyncio.gather(*[h.remote() for _ in range(100)])


@pytest.mark.asyncio
async def test_user_defined_method_hangs(serve_instance):
    h = serve.run(Patient.options(graceful_shutdown_timeout_s=0).bind())
    actor = await h.remote()
    await h.set_should_hang.remote()

    await async_wait_for_condition(
        check_new_actor_started, handle=h, original_actors=actor
    )
    await asyncio.gather(*[h.remote() for _ in range(100)])


@pytest.mark.asyncio
async def test_multiple_replicas(serve_instance):
    h = serve.run(Patient.options(num_replicas=2).bind())
    actors = {
        a._actor_id for a in await asyncio.gather(*[h.remote() for _ in range(100)])
    }
    assert len(actors) == 2

    await h.set_should_fail.remote()

    await async_wait_for_condition(
        check_new_actor_started, handle=h, original_actors=actors
    )

    new_actors = {
        a._actor_id for a in await asyncio.gather(*[h.remote() for _ in range(100)])
    }
    assert len(new_actors) == 2
    assert len(new_actors.intersection(actors)) == 1


def test_inherit_healthcheck(serve_instance):
    class Parent:
        def __init__(self):
            self.should_fail = False

        def check_health(self):
            if self.should_fail:
                raise Exception("intended to fail")

        def set_should_fail(self):
            self.should_fail = True

    @serve.deployment(health_check_period_s=1)
    class Child(Parent):
        def __call__(self, *args):
            return ray.get_runtime_context().current_actor

    h = serve.run(Child.bind())
    actors = {h.remote().result()._actor_id for _ in range(100)}
    assert len(actors) == 1

    h.set_should_fail.remote().result()
    wait_for_condition(check_new_actor_started, handle=h, original_actors=actors)


def test_nonconsecutive_failures(serve_instance):
    counter = ray.remote(Counter).remote()

    # Test that a health check failing every other call isn't marked unhealthy.
    @serve.deployment(health_check_period_s=0.1)
    class FlakyHealthCheck:
        def check_health(self):
            curr_count = ray.get(counter.inc.remote())
            if curr_count % 2 == 0:
                raise Exception("Ah! I had evens!")

        def __call__(self, *args):
            return ray.get_runtime_context().current_actor

    h = serve.run(FlakyHealthCheck.bind())
    a1 = h.remote().result()

    # Wait for 10 health check periods, should never get marked unhealthy.
    wait_for_condition(lambda: ray.get(counter.get.remote()) > 10)
    assert h.remote().result()._actor_id == a1._actor_id


def test_consecutive_failures(serve_instance):
    # Test that the health check must fail N times before being restarted.

    counter = ray.remote(Counter).remote()

    @serve.deployment(health_check_period_s=1)
    class ChronicallyUnhealthy:
        def __init__(self):
            self._actor_id = ray.get_runtime_context().current_actor._actor_id
            self._should_fail = False

        def check_health(self):
            if self._should_fail:
                ray.get(counter.inc.remote())
                raise Exception("intended to fail")

        def set_should_fail(self):
            self._should_fail = True
            return self._actor_id

        def __call__(self, *args):
            return self._actor_id

    h = serve.run(ChronicallyUnhealthy.bind())

    def check_fails_3_times():
        original_actor_id = h.set_should_fail.remote().result()

        # Wait until a new actor is started.
        wait_for_condition(lambda: h.remote().result() != original_actor_id)

        # Check that the health check failed N times before replica was killed.
        assert ray.get(counter.get.remote()) == REPLICA_HEALTH_CHECK_UNHEALTHY_THRESHOLD

    # Run the check twice to see that the counter gets reset after a
    # replica is killed.
    check_fails_3_times()
    ray.get(counter.reset.remote())
    check_fails_3_times()


def test_health_check_failure_cause_deploy_failure(serve_instance):
    """If a deployment always fails health check, the deployment should be unhealthy."""

    @serve.deployment
    class AlwaysUnhealthy:
        def check_health(self):
            raise Exception("intended to fail")

        def __call__(self, *args):
            return ray.get_runtime_context().current_actor

    with pytest.raises(RuntimeError):
        serve.run(AlwaysUnhealthy.bind())

    app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
    assert (
        app_status.deployments["AlwaysUnhealthy"].status
        == DeploymentStatus.DEPLOY_FAILED
    )


def test_health_check_failure_makes_deployment_unhealthy_transition(serve_instance):
    """
    If a deployment transitions to unhealthy, then continues to fail health check after
    being restarted, the deployment should be unhealthy.
    """

    class Toggle:
        def __init__(self):
            self._should_fail = False

        def set_should_fail(self):
            self._should_fail = True

        def should_fail(self):
            return self._should_fail

    @serve.deployment(health_check_period_s=1, health_check_timeout_s=1)
    class WillBeUnhealthy:
        def __init__(self, toggle):
            self._toggle = toggle

        def check_health(self):
            if ray.get(self._toggle.should_fail.remote()):
                raise Exception("intended to fail")

        def __call__(self, *args):
            return ray.get_runtime_context().current_actor

    def check_status(expected_status: DeploymentStatus):
        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
        assert app_status.deployments["WillBeUnhealthy"].status == expected_status
        return True

    toggle = ray.remote(Toggle).remote()
    serve.run(WillBeUnhealthy.bind(toggle))

    # Check that deployment is healthy initially
    assert check_status(DeploymentStatus.HEALTHY)

    ray.get(toggle.set_should_fail.remote())

    # Check that deployment is now unhealthy
    wait_for_condition(check_status, expected_status=DeploymentStatus.UNHEALTHY)

    # Check that deployment stays unhealthy
    for _ in range(5):
        assert check_status(DeploymentStatus.UNHEALTHY)


def test_replica_stalled_in_user_code_marked_unhealthy(serve_instance):
    """
    When a replica stalls in the request-serving path and the user-loop watchdog is
    enabled (RAY_SERVE_USER_HEALTH_CHECK_PROBE_MAX_FAIL > 0), repeated probe timeouts
    cause call_user_health_check() to raise, the controller marks the replica unhealthy,
    and a fresh replica is started (issue #61263).

    The watchdog is on by default (MAX_FAIL=3). We use short intervals here so probe
    failures accumulate quickly within the test window.
    """
    # threading.Event.wait() blocks the asyncio event loop (unlike asyncio.Event which
    # yields control). This simulates a replica stuck in a long synchronous call.
    @serve.deployment(
        health_check_period_s=1,
        health_check_timeout_s=3,
        graceful_shutdown_timeout_s=0,
        ray_actor_options={
            "runtime_env": {
                "env_vars": {
                    # Enable the user-loop watchdog with short intervals so that
                    # probe failures accumulate quickly within the test window.
                    "RAY_SERVE_USER_HEALTH_CHECK_PROBE_MAX_FAIL": "2",
                    "RAY_SERVE_USER_HEALTH_CHECK_PROBE_INTERVAL_S": "0.5",
                    "RAY_SERVE_USER_HEALTH_CHECK_PROBE_TIMEOUT_S": "1",
                }
            }
        },
    )
    class StalledReplica:
        def __init__(self):
            self._unblock = threading.Event()

        async def __call__(self):
            self._unblock.wait()
            return "ok"

        def set_unblock(self):
            self._unblock.set()

    handle = serve.run(StalledReplica.bind())

    # Send a request so the replica blocks on _unblock.wait(), wedging the user loop.
    ref = handle.remote()

    def deployment_unhealthy():
        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
        return (
            app_status.deployments["StalledReplica"].status
            == DeploymentStatus.UNHEALTHY
        )

    wait_for_condition(deployment_unhealthy, timeout=60)

    # Unblock so replicas can finish (stalled replica may already be replaced).
    handle.set_unblock.remote()
    try:
        ray.get(ref, timeout=2)
    except Exception:
        pass

    # Wait for deployment to recover (new replica is healthy).
    def deployment_healthy():
        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
        return (
            app_status.deployments["StalledReplica"].status == DeploymentStatus.HEALTHY
        )

    wait_for_condition(deployment_healthy, timeout=30)


def test_gang_health_check_restarts_gang(serve_instance):
    """RESTART_GANG tears down the entire gang on failure while the deployment
    keeps serving traffic with no downtime."""

    class Toggle:
        def __init__(self):
            self._should_fail = False

        def set_should_fail(self):
            self._should_fail = True

        def unset_should_fail(self):
            self._should_fail = False

        def should_fail(self):
            return self._should_fail

    toggle = ray.remote(Toggle).remote()

    @serve.deployment(health_check_period_s=1, health_check_timeout_s=1)
    class GangPatient:
        def __init__(self):
            self._fail = False

        def check_health(self):
            if self._fail and ray.get(toggle.should_fail.remote()):
                raise Exception("intended to fail")

        def __call__(self):
            ctx = ray.serve.context._get_internal_replica_context()
            gc = ctx.gang_context
            return {
                "replica_id": ctx.replica_id.unique_id,
                "gang_id": gc.gang_id if gc else None,
            }

        def set_should_fail(self):
            self._fail = True
            ctx = ray.serve.context._get_internal_replica_context()
            gc = ctx.gang_context
            return {
                "replica_id": ctx.replica_id.unique_id,
                "gang_id": gc.gang_id if gc else None,
            }

    num_replicas = 4
    gang_size = 2
    num_gangs = num_replicas // gang_size

    h = serve.run(
        GangPatient.options(
            num_replicas=num_replicas,
            gang_scheduling_config=GangSchedulingConfig(gang_size=gang_size),
        ).bind()
    )

    # Collect initial replica state.
    initial_replicas = {}
    for _ in range(100):
        result = h.remote().result()
        initial_replicas[result["replica_id"]] = result
        if len(initial_replicas) == num_replicas:
            break
    assert len(initial_replicas) == num_replicas

    # Identify the two distinct gang IDs.
    gang_ids = {r["gang_id"] for r in initial_replicas.values()}
    assert len(gang_ids) == num_gangs

    # Make one replica fail health checks.
    fail_info = h.set_should_fail.remote().result()
    target_gang_id = fail_info["gang_id"]
    surviving_gang_id = (gang_ids - {target_gang_id}).pop()
    ray.get(toggle.set_should_fail.remote())

    # Wait for deployment to become UNHEALTHY.
    def check_unhealthy():
        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
        assert (
            app_status.deployments["GangPatient"].status == DeploymentStatus.UNHEALTHY
        )
        return True

    wait_for_condition(check_unhealthy, timeout=10)

    # Zero-downtime check.
    # While the failed gang is being torn down and before the replacement
    # gang comes up, the surviving gang must keep serving traffic.
    for _ in range(30):
        result = h.remote().result()
        assert result["gang_id"] == surviving_gang_id

    # Turn off failures so replacement replicas start healthy.
    ray.get(toggle.unset_should_fail.remote())

    # Wait for deployment to recover.
    def check_healthy():
        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
        assert app_status.status == "RUNNING"
        assert app_status.deployments["GangPatient"].status == DeploymentStatus.HEALTHY
        return True

    wait_for_condition(check_healthy, timeout=10)

    # Collect final replica state.
    final_replicas = {}
    for _ in range(100):
        result = h.remote().result()
        final_replicas[result["replica_id"]] = result
        if len(final_replicas) == num_replicas:
            break
    assert len(final_replicas) == num_replicas

    # Both replicas from the failed gang should have been replaced.
    old_gang_ids = {
        r["replica_id"]
        for r in initial_replicas.values()
        if r["gang_id"] == target_gang_id
    }
    assert len(old_gang_ids) == gang_size
    assert old_gang_ids.isdisjoint(final_replicas.keys())


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", "-s", __file__]))