import asyncio
import logging
import os
import re
import sys
import time

import httpx
import pytest

import ray
from ray import serve
from ray._common.test_utils import SignalActor, wait_for_condition
from ray.exceptions import RayTaskError
from ray.serve._private.common import DeploymentID, ReplicaState
from ray.serve._private.constants import (
    RAY_SERVE_ENABLE_HA_PROXY,
    SERVE_CONTROLLER_NAME,
    SERVE_DEFAULT_APP_NAME,
    SERVE_NAMESPACE,
    SERVE_PROXY_NAME,
)
from ray.serve._private.test_utils import (
    check_replica_counts,
    get_application_url,
    request_with_retries,
)
from ray.serve.schema import LoggingConfig, ServeDeploySchema
from ray.util.state import list_actors


@pytest.mark.parametrize(
    "deployment_options",
    [
        {"num_replicas": 2},
        {"autoscaling_config": {"min_replicas": 2, "max_replicas": 2}},
    ],
)
def test_recover_start_from_replica_actor_names(serve_instance, deployment_options):
    """Test controller is able to recover starting -> running replicas from
    actor names.
    """

    # Test failed to deploy with total of 2 replicas,
    # but first constructor call fails.
    @serve.deployment(
        name="recover_start_from_replica_actor_names", **deployment_options
    )
    class TransientConstructorFailureDeployment:
        def __init__(self):
            pass

        def __call__(self, *args):
            return "hii"

    serve.run(TransientConstructorFailureDeployment.bind(), name="app")
    for _ in range(10):
        response = request_with_retries(timeout=30, app_name="app")
        assert response.text == "hii"
    # Assert 2 replicas are running in deployment deployment after partially
    # successful deploy() call with transient error
    deployment_dict = ray.get(serve_instance._controller._all_running_replicas.remote())
    id = DeploymentID(name="recover_start_from_replica_actor_names", app_name="app")
    assert len(deployment_dict[id]) == 2

    replica_version_hash = None
    for replica in deployment_dict[id]:
        ref = replica.get_actor_handle().initialize_and_get_metadata.remote()
        _, version, *_ = ray.get(ref)
        if replica_version_hash is None:
            replica_version_hash = hash(version)
        assert replica_version_hash == hash(version), (
            "Replica version hash should be the same for "
            "same code version and user config."
        )

    # Sample: [
    # 'TransientConstructorFailureDeployment#xlituP',
    # 'SERVE_CONTROLLER_ACTOR',
    # 'TransientConstructorFailureDeployment#NosHNA',
    # 'SERVE_CONTROLLER_ACTOR:SERVE_PROXY_ACTOR-node:192.168.86.165-0']
    actor_infos = list_actors(filters=[("state", "=", "ALIVE")])
    replica_names = [
        actor_info["name"]
        for actor_info in actor_infos
        if (
            SERVE_CONTROLLER_NAME not in actor_info["name"]
            and SERVE_PROXY_NAME not in actor_info["name"]
        )
    ]
    assert (
        len(replica_names) == 2
    ), "Should have two running replicas fetched from ray API."

    # Kill controller and wait for endpoint to be available again
    ray.kill(serve.context._global_client._controller, no_restart=False)
    wait_for_condition(
        lambda: get_application_url("HTTP", "app", use_localhost=True) is not None
    )
    for _ in range(10):
        response = request_with_retries(timeout=30, app_name="app")
        assert response.text == "hii"

    # Ensure recovered replica names are the same
    recovered_actor_infos = list_actors(filters=[("state", "=", "ALIVE")])
    recovered_replica_names = [
        actor_info["name"]
        for actor_info in recovered_actor_infos
        if (
            SERVE_CONTROLLER_NAME not in actor_info["name"]
            and SERVE_PROXY_NAME not in actor_info["name"]
        )
    ]
    assert (
        recovered_replica_names == replica_names
    ), "Running replica actor names after recovery must match"

    # Ensure recovered replica version has are the same
    for replica_name in recovered_replica_names:
        actor_handle = ray.get_actor(replica_name, namespace=SERVE_NAMESPACE)
        ref = actor_handle.initialize_and_get_metadata.remote()
        _, version, *_ = ray.get(ref)
        assert replica_version_hash == hash(
            version
        ), "Replica version hash should be the same after recover from actor names"


def test_recover_rolling_update_from_replica_actor_names(serve_instance):
    """Test controller can recover replicas during rolling update.

    Replicas starting -> updating -> running
    replicas from actor names, with right replica versions during rolling
    update.
    """

    signal = SignalActor.remote()

    @serve.deployment(name="test", num_replicas=2)
    class V1:
        async def __call__(self):
            await signal.wait.remote()
            return "1", os.getpid()

    @serve.deployment(name="test", num_replicas=2)
    class V2:
        async def __call__(self):
            return "2", os.getpid()

    h = serve.run(V1.bind(), name="app")

    # Send requests to get pids of initial 2 replicas
    signal.send.remote()
    refs = [h.remote() for _ in range(10)]
    versions, pids = zip(*[ref.result() for ref in refs])
    assert versions.count("1") == 10
    initial_pids = set(pids)
    assert len(initial_pids) == 2

    # blocked_ref will block a single replica until the signal is sent.
    signal.send.remote(clear=True)
    blocked_ref = h.remote()

    # Kill the controller
    ray.kill(serve.context._global_client._controller, no_restart=False)

    # Redeploy new version.
    serve._run(V2.bind(), _blocking=False, name="app")

    # One replica of the old version should be stuck in stopping because
    # of the blocked request. Two replicas of the new version should be
    # brought up without waiting for the old replica to stop.
    wait_for_condition(
        check_replica_counts,
        controller=serve_instance._controller,
        deployment_id=DeploymentID(name="test", app_name="app"),
        total=3,
        by_state=[
            (ReplicaState.STOPPING, 1, lambda r: r._actor.pid in initial_pids),
            (ReplicaState.RUNNING, 2, lambda r: r._actor.pid not in initial_pids),
        ],
    )

    # All new requests should be sent to the new running replicas
    refs = [h.remote() for _ in range(10)]
    versions, pids = zip(*[ref.result(timeout_s=5) for ref in refs])
    assert versions.count("2") == 10
    pids2 = set(pids)
    assert len(pids2 & initial_pids) == 0

    # Kill the controller
    ray.kill(serve.context._global_client._controller, no_restart=False)

    # Release the signal so that the old replica can shutdown
    ray.get(signal.send.remote())
    val, pid = blocked_ref.result()
    assert val == "1"
    assert pid in initial_pids

    # Now the goal and requests to the new version should complete.
    # We should have two running replicas of the new version.
    serve_instance._wait_for_application_running("app")
    check_replica_counts(
        controller=serve_instance._controller,
        deployment_id=DeploymentID(name="test", app_name="app"),
        total=2,
        by_state=(
            [(ReplicaState.RUNNING, 2, lambda r: r._actor.pid not in initial_pids)]
        ),
    )


def test_controller_recover_initializing_actor(serve_instance):
    """Controller crash while a replica is still in `__init__`.

    The previous controller never finished sending the replica its initial
    `initialize_and_get_metadata(rank=...)` call, so the live actor has
    neither a rank nor a fully-initialized user callable. The new
    controller must detect this via the `was_initialized` probe and
    replace the half-initialized actor with a fresh one rather than try to
    recover it (which would silently complete its initialization with
    `rank=None` and break rank tracking).
    """

    signal = SignalActor.remote()
    init_started = SignalActor.remote()
    client = serve_instance

    @ray.remote
    def pending_init_indicator():
        ray.get(init_started.wait.remote())
        return True

    @serve.deployment
    class V1:
        async def __init__(self):
            ray.get(init_started.send.remote())
            await signal.wait.remote()

        def __call__(self, request):
            return f"1|{os.getpid()}"

    serve._run(V1.bind(), _blocking=False, name="app")
    ray.get(pending_init_indicator.remote())

    def get_actor_info(name: str):
        all_current_actors = list_actors(filters=[("state", "=", "ALIVE")])
        for actor in all_current_actors:
            if SERVE_PROXY_NAME in actor["name"]:
                continue
            if name in actor["name"]:
                return actor["name"], actor["pid"]

    original_actor_tag, _ = get_actor_info(f"app#{V1.name}")
    _, controller1_pid = get_actor_info(SERVE_CONTROLLER_NAME)
    ray.kill(serve.context._global_client._controller, no_restart=False)

    # Wait for the controller to be replaced. `list_actors` can briefly
    # report the killed controller as ALIVE (and any new controller as
    # PENDING_CREATION) right after `ray.kill`, so wait specifically for
    # the pid to change.
    def controller_replaced():
        info = get_actor_info(SERVE_CONTROLLER_NAME)
        return info is not None and info[1] != controller1_pid

    wait_for_condition(controller_replaced)

    # The new controller's `was_initialized` probe will report False for
    # the half-initialized actor, so it is killed and replaced. Wait for
    # the replacement replica to start and report it has reached its
    # constructor. Unblock its `__init__` once it has.
    ray.get(pending_init_indicator.remote())
    ray.get(signal.send.remote())
    client._wait_for_application_running("app")

    # The original half-initialized actor should have been replaced with a
    # fresh one (different replica id baked into the actor name).
    new_actor_tag, _ = get_actor_info(f"app#{V1.name}")
    assert new_actor_tag != original_actor_tag

    # And the original actor should actually be dead -- not just hidden by
    # the ALIVE filter on a different replica id. `list_actors` may take a
    # moment to reflect the kill in its state.
    def original_actor_dead():
        matching = list_actors(filters=[("name", "=", original_actor_tag)])
        # Either the entry has been pruned entirely, or it is reported DEAD.
        return not matching or all(a["state"] == "DEAD" for a in matching)

    wait_for_condition(original_actor_dead)


def test_replica_deletion_after_controller_recover(serve_instance):
    """Test that replicas are deleted when controller is recovered"""

    controller = serve.context._global_client._controller

    @serve.deployment(graceful_shutdown_timeout_s=3)
    class V1:
        async def __call__(self):
            while True:
                await asyncio.sleep(0.1)

    handle = serve.run(V1.bind(), name="app")
    _ = handle.remote()
    serve.delete("app", _blocking=False)

    def check_replica(replica_state=None):
        id = DeploymentID(name="V1", app_name="app")
        try:
            replicas = ray.get(controller._dump_replica_states_for_testing.remote(id))
        except RayTaskError as ex:
            # Deployment is not existed any more.
            if isinstance(ex, KeyError):
                return []
            # Unexpected exception raised.
            raise ex
        if replica_state is None:
            replica_state = list(ReplicaState)
        else:
            replica_state = [replica_state]
        return replicas.get(replica_state)

    # Make sure the replica is in STOPPING state.
    wait_for_condition(lambda: len(check_replica(ReplicaState.STOPPING)) > 0)
    ray.kill(serve.context._global_client._controller, no_restart=False)
    # Make sure the replica is in STOPPING state.
    wait_for_condition(lambda: len(check_replica(ReplicaState.STOPPING)) > 0)

    # The graceful shutdown timeout of 3 seconds should be used
    wait_for_condition(lambda: len(check_replica()) == 0, timeout=20)
    # Application should be removed soon after
    wait_for_condition(lambda: "app" not in serve.status().applications, timeout=20)


def test_recover_deleting_application(serve_instance):
    """Test that replicas that are stuck on __del__ when the controller crashes,
    is properly recovered when the controller is recovered.

    This is similar to the test test_replica_deletion_after_controller_recover,
    except what's blocking the deployment is __del__ instead of ongoing requests
    """

    signal = SignalActor.remote()

    @serve.deployment
    class A:
        async def __del__(self):
            await signal.wait.remote()

    id = DeploymentID(name="A")
    serve.run(A.bind())

    @ray.remote
    def delete_task():
        serve.delete(SERVE_DEFAULT_APP_NAME)

    # Delete application and make sure it is stuck on deleting
    delete_ref = delete_task.remote()
    print("Started task to delete application `default`")

    def application_deleting():
        # Confirm application is in deleting state
        app_status = serve.status().applications[SERVE_DEFAULT_APP_NAME]
        assert app_status.status == "DELETING"

        # Confirm deployment is in updating state
        status = serve_instance.get_all_deployment_statuses()[0]
        assert status.name == "A" and status.status == "UPDATING"

        # Confirm replica is stopping
        replicas = ray.get(
            serve_instance._controller._dump_replica_states_for_testing.remote(id)
        )
        assert replicas.count(states=[ReplicaState.STOPPING]) == 1

        # Confirm delete task is still blocked
        finished, pending = ray.wait([delete_ref], timeout=0)
        assert pending and not finished
        return True

    def check_deleted():
        deployment_statuses = serve_instance.get_all_deployment_statuses()
        if len(deployment_statuses) != 0:
            return False

        finished, pending = ray.wait([delete_ref], timeout=0)
        return finished and not pending

    wait_for_condition(application_deleting)
    for _ in range(10):
        time.sleep(0.1)
        assert application_deleting()

    print("Confirmed that application `default` is stuck on deleting.")

    # Kill controller while the application is stuck on deleting
    ray.kill(serve.context._global_client._controller, no_restart=False)
    print("Finished killing the controller (with restart).")

    def check_controller_alive():
        all_current_actors = list_actors(filters=[("state", "=", "ALIVE")])
        for actor in all_current_actors:
            if actor["class_name"] == "ServeController":
                return True
        return False

    wait_for_condition(check_controller_alive)
    print("Controller is back alive.")

    wait_for_condition(application_deleting)
    # Before we send the signal, the application should still be deleting
    for _ in range(10):
        time.sleep(0.1)
        assert application_deleting()

    print("Confirmed that application is still stuck on deleting.")

    # Since we've confirmed the replica is in a stopping state, we can grab
    # the reference to the in-progress graceful shutdown task
    replicas = ray.get(
        serve_instance._controller._dump_replica_states_for_testing.remote(id)
    )
    graceful_shutdown_ref = replicas.get()[0]._actor._graceful_shutdown_ref

    signal.send.remote()
    print("Sent signal to unblock deletion of application")
    wait_for_condition(check_deleted)
    print("Confirmed that application finished deleting and delete task has returned.")

    # Make sure graceful shutdown ran successfully
    ray.get(graceful_shutdown_ref)
    print("Confirmed that graceful shutdown ran successfully.")


def test_controller_crashes_with_logging_config(serve_instance):
    """Controller persists logging config into kv store, and when controller recover
    from crash, it will read logging config from kv store and apply to the
    controller and proxy.
    """

    @serve.deployment
    class Model:
        def __init__(self):
            self.logger = logging.getLogger("ray.serve")

        def __call__(self):
            self.logger.debug("this_is_debug_info")
            return

    serve.run(Model.bind())

    client = serve_instance

    # Update the logging config
    client.update_global_logging_config(
        LoggingConfig(encoding="JSON", log_level="DEBUG")
    )

    def check_log_file(log_file: str, expected_regex: list):
        with open(log_file, "r") as f:
            s = f.read()
            for regex in expected_regex:
                assert re.findall(regex, s) != []
        return True

    # Check the controller update
    def check_log_state():
        logging_config, _ = ray.get(client._controller._get_logging_config.remote())
        assert logging_config.encoding == "JSON"
        assert logging_config.log_level == "DEBUG"
        return True

    wait_for_condition(check_log_state, timeout=60)
    _, log_file_path = ray.get(client._controller._get_logging_config.remote())
    # DEBUG level check & JSON check
    check_log_file(
        log_file_path,
        [".*Configure the serve controller logger.*", '.*"component_name":.*'],
    )

    ray.kill(client._controller, no_restart=False)

    def check_controller_alive():
        all_current_actors = list_actors(filters=[("state", "=", "ALIVE")])
        for actor in all_current_actors:
            if actor["class_name"] == "ServeController":
                return True
        return False

    wait_for_condition(check_controller_alive)

    # Check the controller log config
    wait_for_condition(check_log_state)
    _, new_log_file_path = ray.get(client._controller._get_logging_config.remote())
    assert new_log_file_path != log_file_path
    # Check again, make sure the logging config is recovered.
    check_log_file(new_log_file_path, ['.*"component_name":.*'])

    # Check proxy logging
    def check_proxy_handle_in_controller():
        proxy_handles = ray.get(client._controller.get_proxies.remote())
        expected_num_proxies = 1
        if RAY_SERVE_ENABLE_HA_PROXY:
            # fallback proxy
            expected_num_proxies += 1
        assert len(proxy_handles) == expected_num_proxies
        return True

    wait_for_condition(check_proxy_handle_in_controller)
    proxy_handles = ray.get(client._controller.get_proxies.remote())
    proxy_handle = list(proxy_handles.values())[0]
    file_path = ray.get(proxy_handle._get_logging_config.remote())
    # We should see the health check debug log in the proxy logs.
    wait_for_condition(
        check_log_file,
        log_file=file_path,
        expected_regex=['"message": "Received health check."'],
        timeout=15,  # The health check period is 10 seconds.
    )


def test_controller_recover_and_deploy(serve_instance):
    """Ensure that in-progress deploy can finish even after controller dies."""
    client = serve_instance
    signal = SignalActor.options(name="signal123").remote()

    config_json = {
        "applications": [
            {
                "name": SERVE_DEFAULT_APP_NAME,
                "import_path": "ray.serve.tests.test_config_files.hangs.app",
            }
        ]
    }
    config = ServeDeploySchema.model_validate(config_json)
    client.deploy_apps(config)

    wait_for_condition(
        lambda: serve.status().applications["default"].status == "DEPLOYING"
    )
    ray.kill(client._controller, no_restart=False)

    signal.send.remote()

    # When controller restarts, it should redeploy config automatically
    wait_for_condition(
        lambda: httpx.get(f"{get_application_url()}/").text == "hello world"
    )


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