import os
import random
import socket
import subprocess
import tempfile
from contextlib import contextmanager
from copy import deepcopy
from typing import Any, Dict, Generator

import httpx
import pytest
import pytest_asyncio

import ray
from ray import serve
from ray._common.test_utils import SignalActor, wait_for_condition
from ray._common.usage import usage_lib
from ray._common.utils import reset_ray_address
from ray.cluster_utils import AutoscalingCluster, Cluster
from ray.serve._private.test_utils import (
    TELEMETRY_ROUTE_PREFIX,
    TEST_METRICS_EXPORT_PORT,
    check_ray_started,
    check_ray_stopped,
    start_telemetry_app,
)
from ray.serve.config import HTTPOptions, ProxyLocation, gRPCOptions
from ray.serve.context import _get_global_client
from ray.tests.conftest import (  # noqa
    external_redis,
    propagate_logs,
    pytest_runtest_makereport,
)

# https://tools.ietf.org/html/rfc6335#section-6
MIN_DYNAMIC_PORT = 49152
MAX_DYNAMIC_PORT = 65535

TEST_GRPC_SERVICER_FUNCTIONS = [
    "ray.serve.generated.serve_pb2_grpc.add_UserDefinedServiceServicer_to_server",
    "ray.serve.generated.serve_pb2_grpc.add_FruitServiceServicer_to_server",
]

if os.environ.get("RAY_SERVE_INTENTIONALLY_CRASH", False) == 1:
    serve.controller._CRASH_AFTER_CHECKPOINT_PROBABILITY = 0.5


@pytest.fixture(autouse=True)
def _clear_stale_ray_address():
    # Serve CI runs several test targets per container sharing /tmp/ray; a target
    # killed mid-run can leave a ray_current_cluster pointing at a dead cluster.
    # Drop it before each test so an address-less ray.init() starts fresh.
    reset_ray_address()
    yield


@pytest.fixture
def ray_shutdown():
    serve.shutdown()
    if ray.is_initialized():
        ray.shutdown()
    yield
    serve.shutdown()
    if ray.is_initialized():
        ray.shutdown()


@pytest.fixture
def ray_cluster():
    cluster = Cluster()
    yield cluster
    serve.shutdown()
    ray.shutdown()
    cluster.shutdown()


@pytest.fixture
def ray_autoscaling_cluster(request):
    # NOTE(zcin): We have to make a deepcopy here because AutoscalingCluster
    # modifies the dictionary that's passed in.
    params = deepcopy(request.param)
    cluster = AutoscalingCluster(**params)
    cluster.start()
    yield
    serve.shutdown()
    ray.shutdown()
    cluster.shutdown()


@pytest.fixture
def ray_start(scope="module"):
    port = random.randint(MIN_DYNAMIC_PORT, MAX_DYNAMIC_PORT)
    subprocess.check_output(
        [
            "ray",
            "start",
            "--head",
            "--num-cpus",
            "16",
            "--ray-client-server-port",
            f"{port}",
        ]
    )
    try:
        yield f"localhost:{port}"
    finally:
        subprocess.check_output(["ray", "stop", "--force"])


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


@contextmanager
def start_and_shutdown_ray_cli():
    subprocess.check_output(["ray", "stop", "--force"])
    wait_for_condition(_check_ray_stop, timeout=15)
    subprocess.check_output(["ray", "start", "--head"])

    yield

    subprocess.check_output(["ray", "stop", "--force"])
    wait_for_condition(_check_ray_stop, timeout=15)


@pytest.fixture(scope="module")
def start_and_shutdown_ray_cli_module():
    with start_and_shutdown_ray_cli():
        yield


@pytest.fixture
def tmp_dir():
    with tempfile.TemporaryDirectory() as tmp_dir:
        old_dir = os.getcwd()
        os.chdir(tmp_dir)
        yield tmp_dir
        os.chdir(old_dir)


@pytest.fixture(scope="session")
def _shared_serve_instance():
    # Note(simon):
    # This line should be not turned on on master because it leads to very
    # spammy and not useful log in case of a failure in CI.
    # To run locally, please use this instead.
    # SERVE_DEBUG_LOG=1 pytest -v -s test_api.py
    # os.environ["SERVE_DEBUG_LOG"] = "1" <- Do not uncomment this.

    # Overriding task_retry_delay_ms to relaunch actors more quickly
    ray.init(
        address="local",
        num_cpus=36,
        namespace="default_test_namespace",
        _metrics_export_port=9999,
        _system_config={"metrics_report_interval_ms": 1000, "task_retry_delay_ms": 50},
    )
    serve.start(
        proxy_location=ProxyLocation.HeadOnly,
        http_options={"host": "0.0.0.0"},
        grpc_options={
            "port": 9000,
            "grpc_servicer_functions": TEST_GRPC_SERVICER_FUNCTIONS,
        },
    )
    yield _get_global_client()
    # Shutdown Serve and Ray when the session ends so that proxy actors
    # (e.g. HAProxyManager) run their shutdown logic and stop subprocesses.
    serve.shutdown()


@pytest_asyncio.fixture
async def serve_instance_async(_shared_serve_instance):
    yield _shared_serve_instance
    # Clear all state for 2.x applications and deployments.
    _shared_serve_instance.delete_all_apps()
    # Clear the ServeHandle cache between tests to avoid them piling up.
    await _shared_serve_instance.shutdown_cached_handles_async()


@pytest.fixture
def serve_instance(_shared_serve_instance):
    yield _shared_serve_instance
    # Clear all state for 2.x applications and deployments.
    _shared_serve_instance.delete_all_apps()
    # Clear the ServeHandle cache between tests to avoid them piling up.
    _shared_serve_instance.shutdown_cached_handles()


@pytest.fixture
def serve_instance_with_signal(serve_instance):
    client = serve_instance

    signal = SignalActor.options(name="signal123").remote()
    yield client, signal

    # Delete signal actor so there is no conflict between tests
    ray.kill(signal)


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


@pytest.fixture(scope="function")
def ray_start_stop():
    subprocess.check_output(["ray", "stop", "--force"])
    ray.shutdown()
    wait_for_condition(
        check_ray_stop,
        timeout=15,
    )
    subprocess.check_output(["ray", "start", "--head"])
    wait_for_condition(
        lambda: httpx.get("http://localhost:8265/api/ray/version").status_code == 200,
        timeout=15,
    )
    ray.init("auto")
    yield
    serve.shutdown()
    ray.shutdown()
    subprocess.check_output(["ray", "stop", "--force"])
    wait_for_condition(
        check_ray_stop,
        timeout=15,
    )


@pytest.fixture(scope="function")
def ray_start_stop_in_specific_directory(request):
    original_working_dir = os.getcwd()

    # Change working directory so Ray will start in the requested directory.
    new_working_dir = request.param
    os.chdir(new_working_dir)
    print(f"\nChanged working directory to {new_working_dir}\n")

    subprocess.check_output(["ray", "start", "--head"])
    wait_for_condition(
        lambda: httpx.get("http://localhost:8265/api/ray/version").status_code == 200,
        timeout=15,
    )
    try:
        yield
    finally:
        # Change the directory back to the original one.
        os.chdir(original_working_dir)
        print(f"\nChanged working directory back to {original_working_dir}\n")

        subprocess.check_output(["ray", "stop", "--force"])
        wait_for_condition(
            check_ray_stop,
            timeout=15,
        )


@pytest.fixture
def ray_instance(
    request: pytest.FixtureRequest,
) -> Generator[Dict[str, Any], None, None]:
    """Starts and stops a Ray instance for this test.

    Args:
        request: request.param should contain a dictionary of env vars and
            their values. The Ray instance will be started with these env vars.

    Yields:
        Dict[str, Any]: The dict returned by ``ray.init`` for the started cluster.
    """

    original_env_vars = os.environ.copy()

    try:
        requested_env_vars = request.param
    except AttributeError:
        requested_env_vars = {}

    os.environ.update(requested_env_vars)
    yield ray.init(
        address="local",
        _metrics_export_port=9999,
        _system_config={
            "metrics_report_interval_ms": 1000,
            "task_retry_delay_ms": 50,
        },
    )

    serve.shutdown()
    ray.shutdown()

    os.environ.clear()
    os.environ.update(original_env_vars)


@pytest.fixture
def manage_ray_with_telemetry(monkeypatch):
    with monkeypatch.context() as m:
        m.setenv("RAY_USAGE_STATS_ENABLED", "1")
        m.setenv(
            "RAY_USAGE_STATS_REPORT_URL",
            f"http://127.0.0.1:8000{TELEMETRY_ROUTE_PREFIX}",
        )
        m.setenv("RAY_USAGE_STATS_REPORT_INTERVAL_S", "1")
        subprocess.check_output(["ray", "stop", "--force"])
        wait_for_condition(check_ray_stopped, timeout=5)

        subprocess.check_output(["ray", "start", "--head"])
        wait_for_condition(check_ray_started, timeout=5)

        storage = start_telemetry_app()
        wait_for_condition(
            lambda: ray.get(storage.get_reports_received.remote()) > 1, timeout=15
        )

        yield storage

        # Call Python API shutdown() methods to clear global variable state
        serve.shutdown()
        ray.shutdown()

        # Reset global state (any keys that may have been set and cached while the
        # workload was running).
        usage_lib.reset_global_state()

        # Shut down Ray cluster with CLI
        subprocess.check_output(["ray", "stop", "--force"])
        wait_for_condition(check_ray_stopped, timeout=5)


def wait_for_metrics_port_free(port=TEST_METRICS_EXPORT_PORT, timeout=30):
    """
    Ensures the metrics export port is freed.
    """

    def port_free():
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            s.bind(("", port))
            return True
        except OSError:
            return False
        finally:
            s.close()

    wait_for_condition(port_free, timeout=timeout, retry_interval_ms=200)


def wait_for_metrics_endpoint(session_name, port=TEST_METRICS_EXPORT_PORT, timeout=30):
    """
    Ensures the current dashboard agent is serving the metrics endpoint. A
    timeout indicates another agent is still running and holding the port.
    """

    def ready():
        try:
            resp = httpx.get(f"http://localhost:{port}/metrics", timeout=1.0)
        except Exception:
            return False
        return resp.status_code == 200 and f'SessionName="{session_name}"' in resp.text

    wait_for_condition(ready, timeout=timeout, retry_interval_ms=500)


@pytest.fixture
def metrics_start_shutdown(request):
    param = request.param if hasattr(request, "param") else None
    request_timeout_s = param if param else None
    """Fixture provides a fresh Ray cluster to prevent metrics state sharing."""
    wait_for_metrics_port_free()
    ray.init(
        address="local",
        _metrics_export_port=TEST_METRICS_EXPORT_PORT,
        _system_config={
            "metrics_report_interval_ms": 100,
            "task_retry_delay_ms": 50,
        },
    )

    try:
        session_name = ray._private.worker._global_node.session_name
        wait_for_metrics_endpoint(session_name)

        grpc_port = 9000
        grpc_servicer_functions = [
            "ray.serve.generated.serve_pb2_grpc.add_UserDefinedServiceServicer_to_server",
            "ray.serve.generated.serve_pb2_grpc.add_FruitServiceServicer_to_server",
        ]
        yield serve.start(
            grpc_options=gRPCOptions(
                port=grpc_port,
                grpc_servicer_functions=grpc_servicer_functions,
                request_timeout_s=request_timeout_s,
            ),
            http_options=HTTPOptions(
                host="0.0.0.0",
                request_timeout_s=request_timeout_s,
            ),
        )
    finally:
        serve.shutdown()
        ray.shutdown()
        reset_ray_address()


# Helper function to return the node ID of a remote worker.
@ray.remote(num_cpus=0)
def _get_node_id():
    return ray.get_runtime_context().get_node_id()


# Test fixture to start a Serve instance in a RayCluster with two labeled nodes
@pytest.fixture(scope="module")
def serve_instance_with_labeled_nodes():
    cluster = Cluster()

    # Unlabeled default node.
    cluster.add_node(num_cpus=3, resources={"worker0": 1})

    # Node 1 - labeled A100 node in us-west.
    cluster.add_node(
        num_cpus=3,
        resources={"worker1": 1},
        labels={"region": "us-west", "gpu-type": "A100"},
    )

    # Node 2 - labeled H100 node in us-east.
    cluster.add_node(
        num_cpus=3,
        resources={"worker2": 1},
        labels={"region": "us-east", "gpu-type": "H100"},
    )

    cluster.wait_for_nodes()

    if ray.is_initialized():
        ray.shutdown()

    ray.init(address=cluster.address)

    node_1_id = ray.get(_get_node_id.options(resources={"worker1": 1}).remote())
    node_2_id = ray.get(_get_node_id.options(resources={"worker2": 1}).remote())

    serve.start()

    yield _get_global_client(), node_1_id, node_2_id, cluster

    serve.shutdown()
    ray.shutdown()
    cluster.shutdown()