"""
This file defines the common pytest fixtures used in current directory.
"""

import copy
import json
import logging
import os
import platform
import shutil
import socket
import subprocess
import tempfile
import time
from contextlib import contextmanager
from pathlib import Path
from tempfile import gettempdir
from typing import List, Optional
from unittest import mock

import pytest

import ray
import ray._private.ray_constants as ray_constants
from ray._common.network_utils import build_address, find_free_port
from ray._common.test_utils import wait_for_condition
from ray._private.authentication_test_utils import (
    authentication_env_guard,
    clear_auth_token_sources,
    reset_auth_token_state,
    set_auth_mode,
    set_env_auth_token,
)
from ray._private.conftest_utils import set_override_dashboard_url  # noqa: F401
from ray._private.runtime_env import virtualenv_utils
from ray._private.test_utils import (
    RayletKiller,
    external_redis_test_enabled,
    get_and_run_resource_killer,
    get_redis_cli,
    init_error_pubsub,
    init_log_pubsub,
    kill_processes,
    redis_replicas,
    redis_sentinel_replicas,
    reset_autoscaler_v2_enabled_cache,
    rocksdb_gcs_test_enabled,
    setup_tls,
    start_redis_instance,
    start_redis_sentinel_instance,
    teardown_tls,
)
from ray.cluster_utils import AutoscalingCluster, Cluster, cluster_not_supported

import psutil

# TODO (mengjin) Improve the logging in the conftest files so that the logger can log
# information in stdout as well as stderr and replace the print statements in the test
# files
logger = logging.getLogger(__name__)

START_REDIS_WAIT_RETRIES = int(os.environ.get("RAY_START_REDIS_WAIT_RETRIES", "60"))


@pytest.fixture(autouse=True)
def pre_envs(monkeypatch):
    # To make test run faster
    monkeypatch.setenv("RAY_NUM_REDIS_GET_RETRIES", "2")
    ray_constants.NUM_REDIS_GET_RETRIES = 2
    yield


def wait_for_redis_to_start(
    redis_ip_address: str,
    redis_port: int,
    password: Optional[str] = None,
    username: Optional[str] = None,
):
    """Wait for a Redis server to be available.

    This is accomplished by creating a Redis client and sending a random
    command to the server until the command gets through.

    Args:
        redis_ip_address: The IP address of the redis server.
        redis_port: The port of the redis server.
        password: The password of the Redis server.
        username: The username of the Redis server.

    Raises:
        Exception: An exception is raised if we could not connect with Redis.
    """
    import redis

    redis_client = redis.StrictRedis(
        host=redis_ip_address, port=redis_port, username=username, password=password
    )
    # Wait for the Redis server to start.
    num_retries = START_REDIS_WAIT_RETRIES

    delay = 0.001
    for i in range(num_retries):
        try:
            # Run some random command and see if it worked.
            logger.debug(
                "Waiting for redis server at {} to respond...".format(
                    build_address(redis_ip_address, redis_port)
                )
            )
            redis_client.client_list()
        # If the Redis service is delayed getting set up for any reason, we may
        # get a redis.ConnectionError: Error 111 connecting to host:port.
        # Connection refused.
        # Unfortunately, redis.ConnectionError is also the base class of
        # redis.AuthenticationError. We *don't* want to obscure a
        # redis.AuthenticationError, because that indicates the user provided a
        # bad password. Thus a double except clause to ensure a
        # redis.AuthenticationError isn't trapped here.
        except redis.AuthenticationError as authEx:
            raise RuntimeError(
                f"Unable to connect to Redis at {build_address(redis_ip_address, redis_port)}."
            ) from authEx
        except redis.ConnectionError as connEx:
            if i >= num_retries - 1:
                raise RuntimeError(
                    f"Unable to connect to Redis at {build_address(redis_ip_address, redis_port)} "
                    f"after {num_retries} retries. Check that "
                    f"{build_address(redis_ip_address, redis_port)} is reachable from this "
                    "machine. If it is not, your firewall may be blocking "
                    "this port. If the problem is a flaky connection, try "
                    "setting the environment variable "
                    "`RAY_START_REDIS_WAIT_RETRIES` to increase the number of"
                    " attempts to ping the Redis server."
                ) from connEx
            # Wait a little bit.
            time.sleep(delay)
            # Make sure the retry interval doesn't increase too large, which will
            # affect the delivery time of the Ray cluster.
            delay = min(1, delay * 2)
        else:
            break
    else:
        raise RuntimeError(
            f"Unable to connect to Redis (after {num_retries} retries). "
            "If the Redis instance is on a different machine, check that "
            "your firewall and relevant Ray ports are configured properly. "
            "You can also set the environment variable "
            "`RAY_START_REDIS_WAIT_RETRIES` to increase the number of "
            "attempts to ping the Redis server."
        )


def get_default_fixure_system_config():
    system_config = {
        "object_timeout_milliseconds": 200,
        "health_check_initial_delay_ms": 0,
        "health_check_failure_threshold": 10,
        "object_store_full_delay_ms": 100,
        "local_gc_min_interval_s": 1,
    }
    return system_config


def get_default_fixture_ray_kwargs():
    system_config = get_default_fixure_system_config()
    ray_kwargs = {
        "num_cpus": 1,
        "object_store_memory": 150 * 1024 * 1024,
        "dashboard_port": None,
        "namespace": "default_test_namespace",
        "_system_config": system_config,
    }
    return ray_kwargs


def is_process_listen_to_port(pid, port):
    retry_num = 10
    interval_time = 0.5
    for _ in range(retry_num):
        try:
            proc = psutil.Process(pid)
            for conn in proc.connections():
                if conn.status == "LISTEN" and conn.laddr.port == port:
                    return True
        except Exception:
            pass
        finally:
            time.sleep(interval_time)
    print(
        f"Process({pid}) has not listened to port {port} "
        + f"for more than {retry_num * interval_time}s."
    )
    return False


def find_user_process_by_port_and_status(
    port: int, statuses_to_check: Optional[list[str]]
):
    """
    Test helper function to find the processes that have a connection to a provided
    port and with statuses in the provided list.

    Args:
        port: The port to check.
        statuses_to_check: The list of statuses to check. If None, the function will not
            check the status of the connection.

    Returns:
        The first process that have a connection
    """
    # Here the function finds all the processes and checks if each of them is with the
    # provided port and status. This is inefficient comparing to leveraging
    # psutil.net_connections to directly filter the processes by port. However, the
    # method is chosen because the net_connections method will need root access to run
    # on macOS: https://psutil.readthedocs.io/en/latest/#psutil.net_connections.
    # Therefore, the current solution is chosen to make the function work for all
    processes = []
    for pid in psutil.pids():
        process = psutil.Process(pid)
        try:
            conns = process.connections()
            for conn in conns:
                if conn.laddr.port == port:
                    if statuses_to_check is None or conn.status in statuses_to_check:
                        processes.append(process)
        except (psutil.AccessDenied, psutil.ZombieProcess, psutil.NoSuchProcess):
            continue

    if not processes:
        print(
            f"Failed to find processes that have connections to the port {port} and "
            f"with connection status in {statuses_to_check}.  It could "
            "be because the process needs higher privilege to access its "
            "information or the port is not listened by any processes."
        )
    return processes


def redis_alive(port, enable_tls):
    try:
        # If there is no redis libs installed, skip the check.
        # This could happen In minimal test, where we don't have
        # redis.
        import redis
    except Exception:
        return True

    params = {}
    if enable_tls:
        from ray._raylet import Config

        params = {"ssl": True, "ssl_cert_reqs": "required"}
        if Config.REDIS_CA_CERT():
            params["ssl_ca_certs"] = Config.REDIS_CA_CERT()
        if Config.REDIS_CLIENT_CERT():
            params["ssl_certfile"] = Config.REDIS_CLIENT_CERT()
        if Config.REDIS_CLIENT_KEY():
            params["ssl_keyfile"] = Config.REDIS_CLIENT_KEY()

    cli = redis.Redis("localhost", port, **params)

    try:
        return cli.ping()
    except Exception:
        pass
    return False


def _find_available_ports(start: int, end: int, *, num: int = 1) -> List[int]:
    ports = []
    for _ in range(num):
        random_port = 0
        with socket.socket() as s:
            s.bind(("", 0))
            random_port = s.getsockname()[1]
        if random_port >= start and random_port <= end and random_port not in ports:
            ports.append(random_port)
            continue

        for port in range(start, end + 1):
            if port in ports:
                continue
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind(("", port))
                ports.append(port)
                break
            except OSError:
                pass

    if len(ports) != num:
        raise RuntimeError(f"Can't find {num} available port from {start} to {end}.")

    return ports


def start_redis_with_sentinel(db_dir):
    temp_dir = ray._common.utils.get_default_ray_temp_dir()

    redis_ports = _find_available_ports(49159, 55535, num=redis_sentinel_replicas() + 1)
    sentinel_port = redis_ports[0]
    master_port = redis_ports[1]
    redis_processes = [
        start_redis_instance(temp_dir, p, listen_to_localhost_only=True, db_dir=db_dir)[
            1
        ]
        for p in redis_ports[1:]
    ]

    # ensure all redis servers are up
    for port in redis_ports[1:]:
        wait_for_condition(redis_alive, 3, 100, port=port, enable_tls=False)

    # setup replicas of the master
    for port in redis_ports[2:]:
        redis_cli = get_redis_cli(port, False)
        redis_cli.replicaof("127.0.0.1", master_port)
        sentinel_process = start_redis_sentinel_instance(
            temp_dir, sentinel_port, master_port
        )
        address_str = f"127.0.0.1:{sentinel_port}"
        return address_str, redis_processes + [sentinel_process]


def start_redis(db_dir):
    retry_num = 0
    while True:
        is_need_restart = False
        processes = []
        enable_tls = "RAY_REDIS_CA_CERT" in os.environ
        leader_port = None
        leader_id = None
        redis_ports = []
        while len(redis_ports) != redis_replicas():
            temp_dir = ray._common.utils.get_default_ray_temp_dir()
            port, free_port = _find_available_ports(49159, 55535, num=2)
            try:
                node_id = None
                proc = None
                node_id, proc = start_redis_instance(
                    temp_dir,
                    port,
                    enable_tls=enable_tls,
                    replica_of=leader_port,
                    leader_id=leader_id,
                    db_dir=db_dir,
                    free_port=free_port,
                )
                wait_for_condition(
                    redis_alive, 3, 100, port=port, enable_tls=enable_tls
                )
            except Exception as e:
                print(f"Fails to start redis on port {port} with exception {e}")
                if (
                    proc is not None
                    and proc.process is not None
                    and proc.process.poll() is None
                ):
                    proc.process.kill()

                # TODO (mengjin) Here we added more debug logs here to help further
                # troubleshoot the potential race condition where the available port
                # we found above is taken by another process and the Redis server
                # cannot be started. Here we won't fail the test but we can check the
                # output log of the test to further investigate the issue if needed.
                if "Redis process exited unexpectedly" in str(e):
                    # Output the process that listens to the port
                    processes = find_user_process_by_port_and_status(
                        port, [psutil.CONN_LISTEN]
                    )

                    for process in processes:
                        print(
                            f"Another process({process.pid}) with command"
                            f"\"{' '.join(process.args)}\" is listening on the port"
                            f"{port}"
                        )

                continue

            redis_ports.append(port)
            if leader_port is None:
                leader_port = port
                leader_id = node_id
            processes.append(proc)
            # Check if th redis has started successfully and is listening on the port.
            if not is_process_listen_to_port(proc.process.pid, port):
                is_need_restart = True
                break

        if is_need_restart:
            retry_num += 1
            for proc in processes:
                proc.process.kill()

            if retry_num > 5:
                raise RuntimeError(
                    f"Failed to start Redis on port {port} after {retry_num} attempts."
                )
            print(
                "Retry to start redis because the process failed to "
                + f"listen to the port({port}), retry num:{retry_num}."
            )
            continue

        if redis_replicas() > 1:

            redis_cli = get_redis_cli(str(leader_port), enable_tls)
            while redis_cli.cluster("info")["cluster_state"] != "ok":
                pass

        scheme = "rediss://" if enable_tls else ""
        address_str = f"{scheme}127.0.0.1:{redis_ports[-1]}"
        return address_str, processes


def kill_all_redis_server():
    """
    Find all redis server processes running on this host via cmdline
    and kill them.
    Note: killed redis process will raise ResourceWarning
          when the python Subprocess tracking the
          underlying process is garbage collected.
    """
    # Find Redis server processes
    redis_procs = []
    for proc in psutil.process_iter(["name", "cmdline"]):
        try:
            if proc.name() == "redis-server":
                redis_procs.append(proc)
        except psutil.NoSuchProcess:
            pass

    # Kill Redis server processes
    for proc in redis_procs:
        try:
            proc.kill()
        except psutil.NoSuchProcess:
            pass


@contextmanager
def _setup_redis(request, with_sentinel=False):
    with tempfile.TemporaryDirectory() as tmpdirname:
        kill_all_redis_server()
        address_str, processes = (
            start_redis_with_sentinel(tmpdirname)
            if with_sentinel
            else start_redis(tmpdirname)
        )
        old_addr = os.environ.get("RAY_REDIS_ADDRESS")
        os.environ["RAY_REDIS_ADDRESS"] = address_str
        import uuid

        ns = str(uuid.uuid4())
        old_ns = os.environ.get("RAY_external_storage_namespace")
        os.environ["RAY_external_storage_namespace"] = ns

        yield
        if old_addr is not None:
            os.environ["RAY_REDIS_ADDRESS"] = old_addr
        else:
            del os.environ["RAY_REDIS_ADDRESS"]

        if old_ns is not None:
            os.environ["RAY_external_storage_namespace"] = old_ns
        else:
            del os.environ["RAY_external_storage_namespace"]

        kill_processes(processes)


@contextmanager
def _setup_rocksdb_gcs(request):
    """Configure the env so a Ray cluster started inside the fixture uses
    the RocksDB GCS backend (REP-64). The DB lives in a tempdir scoped
    to the fixture; nothing persists beyond the test.
    """
    with tempfile.TemporaryDirectory() as tmpdirname:
        old_storage = os.environ.get("RAY_gcs_storage")
        old_path = os.environ.get("RAY_gcs_storage_path")
        os.environ["RAY_gcs_storage"] = "rocksdb"
        os.environ["RAY_gcs_storage_path"] = tmpdirname
        try:
            yield
        finally:
            if old_storage is not None:
                os.environ["RAY_gcs_storage"] = old_storage
            else:
                del os.environ["RAY_gcs_storage"]
            if old_path is not None:
                os.environ["RAY_gcs_storage_path"] = old_path
            else:
                del os.environ["RAY_gcs_storage_path"]


@pytest.fixture
def maybe_setup_external_redis(request):
    # Dispatches the configured GCS storage backend based on CI env
    # vars. Despite the historical name, this fixture also handles the
    # RocksDB GCS backend (REP-64) so existing cluster fixtures pick up
    # rocksdb-mode behavior automatically when TEST_GCS_ROCKSDB=1.
    if external_redis_test_enabled():
        with _setup_redis(request):
            yield
    elif rocksdb_gcs_test_enabled():
        with _setup_rocksdb_gcs(request):
            yield
    else:
        yield


@pytest.fixture(scope="module")
def maybe_setup_external_redis_shared(request):
    if external_redis_test_enabled():
        with _setup_redis(request):
            yield
    elif rocksdb_gcs_test_enabled():
        with _setup_rocksdb_gcs(request):
            yield
    else:
        yield


@pytest.fixture
def external_redis(request):
    with _setup_redis(request):
        yield


@pytest.fixture
def external_redis_with_sentinel(request):
    with _setup_redis(request, True):
        yield


@pytest.fixture
def local_autoscaling_cluster(request, enable_v2):
    reset_autoscaler_v2_enabled_cache()

    # Start a mock multi-node autoscaling cluster.
    head_resources, worker_node_types, system_config = copy.deepcopy(request.param)
    cluster = AutoscalingCluster(
        head_resources=head_resources,
        worker_node_types=worker_node_types,
        autoscaler_v2=enable_v2,
    )
    cluster.start(_system_config=system_config)

    yield None

    # Shutdown the cluster
    cluster.shutdown()


@pytest.fixture
def shutdown_only(maybe_setup_external_redis):
    yield None
    # The code after the yield will run as teardown code.
    ray.shutdown()
    # Delete the cluster address just in case.
    ray._common.utils.reset_ray_address()


@pytest.fixture
def propagate_logs():
    # Ensure that logs are propagated to ancestor handles. This is required if using the
    # caplog or capsys fixtures with Ray's logging.
    # NOTE: This only enables log propagation in the driver process, not the workers!
    logging.getLogger("ray").propagate = True
    logging.getLogger("ray.data").propagate = True
    yield
    logging.getLogger("ray").propagate = False
    logging.getLogger("ray.data").propagate = False


# Provide a shared Ray instance for a test class
@pytest.fixture(scope="class")
def class_ray_instance():
    yield ray.init()
    ray.shutdown()
    # Delete the cluster address just in case.
    ray._common.utils.reset_ray_address()


@contextmanager
def _ray_start(**kwargs):
    init_kwargs = get_default_fixture_ray_kwargs()
    init_kwargs.update(kwargs)
    # Start the Ray processes.
    address_info = ray.init("local", **init_kwargs)

    yield address_info
    # The code after the yield will run as teardown code.
    ray.shutdown()
    # Delete the cluster address just in case.
    ray._common.utils.reset_ray_address()


@pytest.fixture
def ray_start_with_dashboard(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    if param.get("num_cpus") is None:
        param["num_cpus"] = 1
    with _ray_start(include_dashboard=True, **param) as address_info:
        yield address_info


@pytest.fixture
def ray_start_with_dashboard_and_proxy(request, httpserver, maybe_setup_external_redis):
    hsurl = httpserver.url_for("/")

    param = getattr(request, "param", {})
    if param.get("num_cpus") is None:
        param["num_cpus"] = 1
    with _ray_start(include_dashboard=True, proxy_server_url=hsurl, **param) as info:
        yield info


@pytest.fixture
def make_sure_dashboard_http_port_unused():
    """Make sure the dashboard agent http port is unused."""
    for process in psutil.process_iter():
        should_kill = False
        try:
            for conn in process.connections():
                if conn.laddr.port == ray_constants.DEFAULT_DASHBOARD_AGENT_LISTEN_PORT:
                    should_kill = True
                    break
        except Exception:
            continue
        if should_kill:
            try:
                process.kill()
                process.wait()
            except Exception:
                pass
    yield


# The following fixture will start ray with 0 cpu.
@pytest.fixture
def ray_start_no_cpu(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start(num_cpus=0, **param) as res:
        yield res


# The following fixture will start ray with 1 cpu.
@pytest.fixture
def ray_start_regular(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start(**param) as res:
        yield res


# We can compose external_redis and ray_start_regular instead of creating this
# separate fixture, if there is a good way to ensure external_redis runs before
# ray_start_regular.
@pytest.fixture
def ray_start_regular_with_external_redis(request, external_redis):
    param = getattr(request, "param", {})
    with _ray_start(**param) as res:
        yield res


@pytest.fixture(scope="module")
def ray_start_regular_shared(request):
    param = getattr(request, "param", {})
    with _ray_start(**param) as res:
        yield res


@pytest.fixture(scope="module")
def ray_start_regular_shared_2_cpus(request):
    param = getattr(request, "param", {})
    with _ray_start(num_cpus=2, **param) as res:
        yield res


@pytest.fixture
def ray_start_2_cpus(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start(num_cpus=2, **param) as res:
        yield res


@pytest.fixture
def ray_start_10_cpus(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start(num_cpus=10, **param) as res:
        yield res


@contextmanager
def _ray_start_cluster(**kwargs):
    if kwargs.pop("skip_cluster", cluster_not_supported):
        pytest.skip("Cluster not supported")

    init_kwargs = get_default_fixture_ray_kwargs()
    num_nodes = 0
    do_init = False
    # num_nodes & do_init are not arguments for ray.init, so delete them.
    if "num_nodes" in kwargs:
        num_nodes = kwargs["num_nodes"]
        del kwargs["num_nodes"]
    if "do_init" in kwargs:
        do_init = kwargs["do_init"]
        del kwargs["do_init"]
    elif num_nodes > 0:
        do_init = True
    init_kwargs.update(kwargs)
    namespace = init_kwargs.pop("namespace")

    cluster = Cluster()
    remote_nodes = []
    for i in range(num_nodes):
        if i > 0 and "_system_config" in init_kwargs:
            del init_kwargs["_system_config"]
        remote_nodes.append(cluster.add_node(**init_kwargs))
        # We assume driver will connect to the head (first node),
        # so ray init will be invoked if do_init is true
        if len(remote_nodes) == 1 and do_init:
            ray.init(address=cluster.address, namespace=namespace)
    yield cluster

    # The code after the yield will run as teardown code.
    ray.shutdown()
    cluster.shutdown()


# This fixture will start a cluster with empty nodes.
@pytest.fixture
def ray_start_cluster(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start_cluster(**param) as res:
        yield res


@pytest.fixture
def ray_start_cluster_enabled(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    param["skip_cluster"] = False
    with _ray_start_cluster(**param) as res:
        yield res


@pytest.fixture(scope="module")
def ray_start_cluster_enabled_shared(request, maybe_setup_external_redis_shared):
    param = getattr(request, "param", {})
    param["skip_cluster"] = False
    with _ray_start_cluster(**param) as res:
        yield res


@pytest.fixture
def ray_start_cluster_init(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start_cluster(do_init=True, **param) as res:
        yield res


@pytest.fixture
def ray_start_cluster_head(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start_cluster(do_init=True, num_nodes=1, **param) as res:
        yield res


# We can compose external_redis and ray_start_cluster_head instead of creating
# this separate fixture, if there is a good way to ensure external_redis runs
# before ray_start_cluster_head.
@pytest.fixture
def ray_start_cluster_head_with_external_redis(request, external_redis):
    param = getattr(request, "param", {})
    with _ray_start_cluster(do_init=True, num_nodes=1, **param) as res:
        yield res


@pytest.fixture
def ray_start_cluster_head_with_external_redis_sentinel(
    request, external_redis_with_sentinel
):
    param = getattr(request, "param", {})
    with _ray_start_cluster(do_init=True, num_nodes=1, **param) as res:
        yield res


@pytest.fixture
def ray_start_cluster_head_with_env_vars(
    request, maybe_setup_external_redis, monkeypatch
):
    param = getattr(request, "param", {})
    env_vars = param.get("env_vars", {})
    # Create a copy of param without env_vars to pass to _ray_start_cluster
    cluster_param = {k: v for k, v in param.items() if k != "env_vars"}
    for k, v in env_vars.items():
        monkeypatch.setenv(k, v)
    with _ray_start_cluster(do_init=True, num_nodes=1, **cluster_param) as res:
        yield res


@pytest.fixture
def ray_start_cluster_2_nodes(request, maybe_setup_external_redis):
    param = getattr(request, "param", {})
    with _ray_start_cluster(do_init=True, num_nodes=2, **param) as res:
        yield res


@pytest.fixture
def ray_start_object_store_memory(request, maybe_setup_external_redis):
    # Start the Ray processes.
    store_size = request.param
    system_config = get_default_fixure_system_config()
    init_kwargs = {
        "num_cpus": 1,
        "_system_config": system_config,
        "object_store_memory": store_size,
    }
    ray.init("local", **init_kwargs)
    yield store_size
    # The code after the yield will run as teardown code.
    ray.shutdown()


@pytest.fixture
def call_ray_start(request):
    with call_ray_start_context(request) as address:
        yield address


# This fixture will start an httpserver and use it as the proxy-server-url parameters
@pytest.fixture
def call_ray_start_context_with_proxy_server(httpserver):
    hsurl = httpserver.url_for("/")
    cmd = f"ray start --head --num-cpus=1 --proxy-server-url={hsurl} --port 0 --min-worker-port=0 --max-worker-port=0"
    tempObject = type("Temp", (), {"param": cmd})

    with call_ray_start_context(tempObject) as address:
        yield address


@contextmanager
def call_ray_start_context(request):
    default_cmd = (
        "ray start --head --num-cpus=1 --min-worker-port=0 "
        "--max-worker-port=0 --port 0"
    )
    parameter = getattr(request, "param", default_cmd)
    env = None

    if isinstance(parameter, dict):
        if "env" in parameter:
            env = {**os.environ, **parameter.get("env")}

        parameter = parameter.get("cmd", default_cmd)

    command_args = parameter.split(" ")
    try:
        out = ray._common.utils.decode(
            subprocess.check_output(command_args, stderr=subprocess.STDOUT, env=env)
        )
    # If the exit code is non-zero subprocess.check_output raises a CalledProcessError
    except subprocess.CalledProcessError as e:
        print("Ray start cmd failed!")
        print(f"Command: {' '.join(e.cmd)}")
        print(f"Exit code: {e.returncode}")
        if e.output:
            print(f"Output:\n{e.output.decode()}")
        raise
    except Exception as e:
        print(type(e), e)
        raise

    # Get the redis address from the output.
    redis_substring_prefix = "--address='"
    idx = out.find(redis_substring_prefix)
    if idx >= 0:
        address_location = idx + len(redis_substring_prefix)
        address = out[address_location:]
        address = address.split("'")[0]
    else:
        address = None

    yield address

    # Disconnect from the Ray cluster.
    ray.shutdown()
    # Kill the Ray cluster.
    subprocess.check_call(["ray", "stop"], env=env)
    # Delete the cluster address just in case.
    ray._common.utils.reset_ray_address()


@pytest.fixture
def init_and_serve():
    import ray.util.client.server.server as ray_client_server

    server_handle, _ = ray_client_server.init_and_serve("localhost", 50051)
    yield server_handle
    ray_client_server.shutdown_with_server(server_handle.grpc_server)
    time.sleep(2)


@pytest.fixture
def call_ray_stop_only():
    yield
    subprocess.check_call(["ray", "stop"])
    # Delete the cluster address just in case.
    ray._common.utils.reset_ray_address()


def _start_cluster(cluster, request):
    cluster.add_node(num_cpus=4, dashboard_agent_listen_port=find_free_port())
    return cluster, cluster.address


# Used to enforce that `start_cluster` and `start_cluster_shared` fixtures aren't mixed.
_START_CLUSTER_SHARED_USED = False


@pytest.fixture
def start_cluster(ray_start_cluster_enabled, request):
    if _START_CLUSTER_SHARED_USED:
        pytest.fail(
            "Cannot mix `start_cluster` and `start_cluster_shared` "
            "fixtures in the same session."
        )
    yield _start_cluster(ray_start_cluster_enabled, request)


@pytest.fixture(scope="module")
def _start_cluster_shared(ray_start_cluster_enabled_shared, request):
    global _START_CLUSTER_SHARED_USED
    _START_CLUSTER_SHARED_USED = True

    yield _start_cluster(ray_start_cluster_enabled_shared, request)


# Same as `start_cluster` but module-scoped (shared across all tests in a file).
# Runs ray.shutdown after each test.
@pytest.fixture
def start_cluster_shared(_start_cluster_shared, request):
    yield _start_cluster_shared
    ray.shutdown()


@pytest.fixture(scope="function")
def tmp_working_dir():
    with tempfile.TemporaryDirectory() as tmp_dir:
        path = Path(tmp_dir)

        hello_file = path / "hello"
        with hello_file.open(mode="w") as f:
            f.write("world")

        test_file_module = path / "file_module.py"
        with test_file_module.open(mode="w") as f:
            f.write("def hello():\n")
            f.write("    return 'hello'\n")

        module_path = path / "test_module"
        module_path.mkdir(parents=True)

        test_file = module_path / "test.py"
        with test_file.open(mode="w") as f:
            f.write("def one():\n")
            f.write("    return 1\n")

        init_file = module_path / "__init__.py"
        with init_file.open(mode="w") as f:
            f.write("from test_module.test import one\n")

        yield tmp_dir


@pytest.fixture
def enable_pickle_debug():
    os.environ["RAY_PICKLE_VERBOSE_DEBUG"] = "1"
    yield
    del os.environ["RAY_PICKLE_VERBOSE_DEBUG"]


@pytest.fixture
def set_enable_auto_connect(enable_auto_connect: bool = False):
    from ray._private import auto_init_hook

    try:
        old_value = auto_init_hook.enable_auto_connect
        auto_init_hook.enable_auto_connect = enable_auto_connect
        yield enable_auto_connect
    finally:
        auto_init_hook.enable_auto_connect = old_value


@pytest.fixture
def enable_mac_large_object_store():
    os.environ["RAY_ENABLE_MAC_LARGE_OBJECT_STORE"] = "1"
    yield
    del os.environ["RAY_ENABLE_MAC_LARGE_OBJECT_STORE"]


@pytest.fixture()
def two_node_cluster():
    system_config = {
        "object_timeout_milliseconds": 200,
    }
    if cluster_not_supported:
        pytest.skip("Cluster not supported")
    cluster = ray.cluster_utils.Cluster(
        head_node_args={"_system_config": system_config}
    )
    for _ in range(2):
        remote_node = cluster.add_node(num_cpus=1)
    ray.init(address=cluster.address)
    yield cluster, remote_node

    # The code after the yield will run as teardown code.
    ray.shutdown()
    cluster.shutdown()


@pytest.fixture()
def error_pubsub():
    p = init_error_pubsub()
    yield p
    p.close()


@pytest.fixture()
def log_pubsub():
    p = init_log_pubsub()
    yield p
    p.close()


@pytest.fixture
def use_tls(request):
    if request.param:
        key_filepath, cert_filepath, temp_dir = setup_tls()
    yield request.param
    if request.param:
        teardown_tls(key_filepath, cert_filepath, temp_dir)


"""
Object spilling test fixture
"""
# -- Smart open param --
bucket_name = "object-spilling-test"

# -- File system param --
spill_local_path = "/tmp/spill"

# -- Spilling configs --
file_system_object_spilling_config = {
    "type": "filesystem",
    "params": {"directory_path": spill_local_path},
}

buffer_object_spilling_config = {
    "type": "filesystem",
    "params": {"directory_path": spill_local_path, "buffer_size": 1_000_000},
}

# Since we have differet protocol for a local external storage (e.g., fs)
# and distributed external storage (e.g., S3), we need to test both cases.
# This mocks the distributed fs with cluster utils.
mock_distributed_fs_object_spilling_config = {
    "type": "mock_distributed_fs",
    "params": {"directory_path": spill_local_path},
}
smart_open_object_spilling_config = {
    "type": "smart_open",
    "params": {"uri": f"s3://{bucket_name}/"},
}
buffer_open_object_spilling_config = {
    "type": "smart_open",
    "params": {"uri": f"s3://{bucket_name}/", "buffer_size": 1000},
}
multi_smart_open_object_spilling_config = {
    "type": "smart_open",
    "params": {"uri": [f"s3://{bucket_name}/{i}" for i in range(3)]},
}

unstable_object_spilling_config = {
    "type": "unstable_fs",
    "params": {
        "directory_path": spill_local_path,
    },
}
slow_object_spilling_config = {
    "type": "slow_fs",
    "params": {
        "directory_path": spill_local_path,
    },
}


def create_object_spilling_config(request, tmp_path):
    temp_folder = tmp_path / "spill"
    temp_folder.mkdir()
    if (
        request.param["type"] == "filesystem"
        or request.param["type"] == "mock_distributed_fs"
    ):
        request.param["params"]["directory_path"] = str(temp_folder)
    return json.dumps(request.param), temp_folder


@pytest.fixture(
    scope="function",
    params=[
        file_system_object_spilling_config,
    ],
)
def fs_only_object_spilling_config(request, tmp_path):
    yield create_object_spilling_config(request, tmp_path)


@pytest.fixture(
    scope="function",
    params=[
        file_system_object_spilling_config,
    ],
)
def object_spilling_config(request, tmp_path):
    yield create_object_spilling_config(request, tmp_path)


@pytest.fixture(
    scope="function",
    params=[
        file_system_object_spilling_config,
        mock_distributed_fs_object_spilling_config,
    ],
)
def multi_node_object_spilling_config(request, tmp_path):
    yield create_object_spilling_config(request, tmp_path)


@pytest.fixture(
    scope="function",
    params=[
        unstable_object_spilling_config,
    ],
)
def unstable_spilling_config(request, tmp_path):
    yield create_object_spilling_config(request, tmp_path)


@pytest.fixture(
    scope="function",
    params=[
        slow_object_spilling_config,
    ],
)
def slow_spilling_config(request, tmp_path):
    yield create_object_spilling_config(request, tmp_path)


def _ray_start_chaos_cluster(request):
    param = getattr(request, "param", {})
    kill_interval = param.pop("kill_interval", None)
    config = param.pop("_system_config", {})
    config.update(
        {
            "task_retry_delay_ms": 100,
        }
    )
    # Config of workers that are re-started.
    head_resources = param.pop("head_resources")
    worker_node_types = param.pop("worker_node_types")
    cluster = AutoscalingCluster(
        head_resources,
        worker_node_types,
        idle_timeout_minutes=10,  # Don't take down nodes.
        **param,
    )
    cluster.start(_system_config=config)
    ray.init("auto")
    nodes = ray.nodes()
    assert len(nodes) == 1

    if kill_interval is not None:
        node_killer = get_and_run_resource_killer(RayletKiller, kill_interval)

    yield cluster

    if kill_interval is not None:
        ray.get(node_killer.stop_run.remote())
        killed = {
            node_id for node_id, _, _ in ray.get(node_killer.get_killed_nodes.remote())
        }
        assert len(killed) > 0
        died = {node["NodeID"] for node in ray.nodes() if not node["Alive"]}
        assert died.issubset(
            killed
        ), f"Raylets {died - killed} that we did not kill crashed"

    ray.shutdown()
    cluster.shutdown()


@pytest.fixture
def ray_start_chaos_cluster(request):
    """Returns the cluster and chaos thread."""
    for x in _ray_start_chaos_cluster(request):
        yield x


# Set scope to "class" to force this to run before start_cluster, whose scope
# is "function".  We need these env vars to be set before Ray is started.
@pytest.fixture(scope="class")
def runtime_env_disable_URI_cache():
    with mock.patch.dict(
        os.environ,
        {
            "RAY_RUNTIME_ENV_CONDA_CACHE_SIZE_GB": "0",
            "RAY_RUNTIME_ENV_PIP_CACHE_SIZE_GB": "0",
            "RAY_RUNTIME_ENV_WORKING_DIR_CACHE_SIZE_GB": "0",
            "RAY_RUNTIME_ENV_PY_MODULES_CACHE_SIZE_GB": "0",
        },
    ):
        print(
            "URI caching disabled (conda, pip, working_dir, py_modules cache "
            "size set to 0)."
        )
        yield


# Use to create virtualenv that clone from current python env.
# The difference between this fixture and `pytest_virtual.virtual` is that
# `pytest_virtual.virtual` will not inherit current python env's site-package.
# Note: Can't use in virtualenv, this must be noted when testing locally.
@pytest.fixture(scope="function")
def cloned_virtualenv():
    # Lazy import pytest_virtualenv,
    # aviod import `pytest_virtualenv` in test case `Minimal install`
    from pytest_virtualenv import VirtualEnv

    if virtualenv_utils.is_in_virtualenv():
        raise RuntimeError("Forbid the use of this fixture in virtualenv")

    venv = VirtualEnv(
        args=[
            "--system-site-packages",
            "--reset-app-data",
            "--no-periodic-update",
            "--no-download",
        ],
    )
    yield venv
    venv.teardown()


@pytest.fixture
def set_runtime_env_retry_times(request):
    runtime_env_retry_times = getattr(request, "param", "0")
    try:
        os.environ["RUNTIME_ENV_RETRY_TIMES"] = runtime_env_retry_times
        yield runtime_env_retry_times
    finally:
        del os.environ["RUNTIME_ENV_RETRY_TIMES"]


@pytest.fixture
def listen_port(request):
    port = getattr(request, "param", 0)
    try:
        sock = socket.socket()
        if hasattr(socket, "SO_REUSEPORT"):
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 0)

        # Try up to 10 times.
        MAX_RETRY = 10
        for i in range(MAX_RETRY):
            try:
                sock.bind(("127.0.0.1", port))
                break
            except OSError as e:
                if i == MAX_RETRY - 1:
                    raise e
                else:
                    print(f"failed to bind on a port {port}. {e}")
                    time.sleep(1)
        yield port
    finally:
        sock.close()


@pytest.fixture
def set_bad_runtime_env_cache_ttl_seconds(request):
    ttl = getattr(request, "param", "0")
    os.environ["BAD_RUNTIME_ENV_CACHE_TTL_SECONDS"] = ttl
    yield ttl
    del os.environ["BAD_RUNTIME_ENV_CACHE_TTL_SECONDS"]


@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    # execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    try:
        append_short_test_summary(rep)
    except Exception as e:
        print(f"+++ Error creating PyTest summary\n{e}")
    try:
        create_ray_logs_for_failed_test(rep)
    except Exception as e:
        print(f"+++ Error saving Ray logs for failing test\n{e}")


def append_short_test_summary(rep):
    """Writes a short summary txt for failed tests to be printed later."""
    if rep.when != "call":
        return

    summary_dir = os.environ.get("RAY_TEST_SUMMARY_DIR")

    if platform.system() == "Darwin":
        summary_dir = os.environ.get("RAY_TEST_SUMMARY_DIR_HOST")

    if not summary_dir:
        return

    if not os.path.exists(summary_dir):
        os.makedirs(summary_dir, exist_ok=True)

    test_name = rep.nodeid.replace(os.sep, "::")

    if platform.system() == "Windows":
        # ":" is not legal in filenames in windows
        test_name = test_name.replace(":", "$")

    summary_file = os.path.join(summary_dir, test_name + ".txt")

    if rep.passed and os.path.exists(summary_file):
        # The test succeeded after failing, thus it is flaky.
        # We do not want to annotate flaky tests just now, so remove report.
        os.remove(summary_file)
        return

    # Only consider failed tests from now on
    if not rep.failed:
        return

    # No failing test information
    if rep.longrepr is None:
        return

    # No failing test information
    if not hasattr(rep.longrepr, "chain"):
        return

    # Use `wt` here to overwrite so we only have one result per test (exclude retries)
    with open(summary_file, "wt") as fp:
        fp.write(_get_markdown_annotation(rep))


def _get_markdown_annotation(rep) -> str:
    # Main traceback is the last in the chain (where the last error is raised)
    main_tb, main_loc, _ = rep.longrepr.chain[-1]
    markdown = ""

    # Only keep last line of the message
    short_message = list(filter(None, main_loc.message.split("\n")))[-1]

    # Header: Main error message
    markdown += f"#### {rep.nodeid}\n\n"
    markdown += "<details>\n"
    markdown += f"<summary>{short_message}</summary>\n\n"

    # Add location to the test definition
    test_file, test_lineno, _test_node = rep.location
    test_path = os.path.abspath(test_file)
    markdown += f"Test location: {test_path}:{test_lineno}\n\n"

    # Print main traceback
    markdown += "##### Traceback\n\n"
    markdown += "```\n"
    markdown += str(main_tb)
    markdown += "\n```\n\n"

    # Print test definition location
    markdown += f"{main_loc.path}:{main_loc.lineno}\n\n"

    # If this is a longer exception chain, users can expand the full traceback
    if len(rep.longrepr.chain) > 1:
        markdown += "<details><summary>Full traceback</summary>\n\n"

        # Here we just print each traceback and the respective lines.
        for tb, loc, _ in rep.longrepr.chain:
            markdown += "```\n"
            markdown += str(tb)
            markdown += "\n```\n\n"
            if loc:
                markdown += f"{loc.path}:{loc.lineno}\n\n"

        markdown += "</details>\n"

    markdown += "<details><summary>PIP packages</summary>\n\n"
    markdown += "```\n"
    markdown += "\n".join(_get_pip_packages())
    markdown += "\n```\n\n"
    markdown += "</details>\n"

    markdown += "</details>\n\n"
    return markdown


def _get_pip_packages() -> List[str]:
    try:
        from pip._internal.operations import freeze

        return list(freeze.freeze())
    except Exception:
        return ["invalid"]


def create_ray_logs_for_failed_test(rep):
    """Creates artifact zip of /tmp/ray/session_latest/logs for failed tests"""

    # We temporarily restrict to Linux until we have artifact dirs
    # for Windows and Mac
    if platform.system() != "Linux" and platform.system() != "Windows":
        return

    # Only archive failed tests after the "call" phase of the test
    if rep.when != "call" or not rep.failed:
        return

    # Get dir to write zipped logs to
    archive_dir = os.environ.get("RAY_TEST_FAILURE_LOGS_ARCHIVE_DIR")

    if not archive_dir:
        return

    if not os.path.exists(archive_dir):
        os.makedirs(archive_dir)

    # Get logs dir from the latest ray session
    tmp_dir = gettempdir()
    logs_dir = os.path.join(tmp_dir, "ray", "session_latest", "logs")

    if not os.path.exists(logs_dir):
        return

    # Write zipped logs to logs archive dir
    test_name = rep.nodeid.replace(os.sep, "::")
    if platform.system() == "Windows":
        # ":" is not legal in filenames in windows
        test_name = test_name.replace(":", "$")
    output_file = os.path.join(archive_dir, f"{test_name}_{time.time():.4f}")
    shutil.make_archive(output_file, "zip", logs_dir)


@pytest.fixture(params=[True, False])
def start_http_proxy(request):
    env = {}

    proxy = None
    try:
        if request.param:
            # the `proxy` command is from the proxy.py package.
            proxy = subprocess.Popen(
                ["proxy", "--port", "8899", "--log-level", "ERROR"]
            )
            env["RAY_grpc_enable_http_proxy"] = "1"
            proxy_url = "http://localhost:8899"
        else:
            proxy_url = "http://example.com"
        env["http_proxy"] = proxy_url
        env["https_proxy"] = proxy_url
        yield env
    finally:
        if proxy:
            proxy.terminate()
            proxy.wait()


@pytest.fixture
def set_runtime_env_plugins(request):
    runtime_env_plugins = getattr(request, "param", "0")
    try:
        os.environ["RAY_RUNTIME_ENV_PLUGINS"] = runtime_env_plugins
        yield runtime_env_plugins
    finally:
        del os.environ["RAY_RUNTIME_ENV_PLUGINS"]


@pytest.fixture(scope="function")
def temp_file(request):
    with tempfile.NamedTemporaryFile("r+b") as fp:
        yield fp


@pytest.fixture(scope="function")
def temp_dir(request):
    with tempfile.TemporaryDirectory("r+b") as d:
        yield d


@pytest.fixture(scope="module")
def random_ascii_file(request):
    import random
    import string

    file_size = getattr(request, "param", 1 << 10)

    with tempfile.NamedTemporaryFile(mode="r+b") as fp:
        fp.write("".join(random.choices(string.ascii_letters, k=file_size)).encode())
        fp.flush()

        yield fp


# Clean up Ray address file before the test run starts, since sometimes bazel test times out
# and kill the test process, without cleaning up the Ray address file.
def pytest_sessionstart(session):
    """Called after the Session object has been created and before performing collection and entering the run test loop."""

    # Delete the cluster address file just in case.
    ray._common.utils.reset_ray_address()


"""
pytest httpserver related test fixtures
"""


@pytest.fixture(scope="module")
def make_httpserver(httpserver_listen_address, httpserver_ssl_context):
    """
    Module-scoped override of pytest-httpserver's make_httpserver fixture.
    Copies the implementation the make_httpserver fixture.
    """
    # Lazy import pytest_httpserver to avoid import errors in library tests that doesn't
    # have pytest_httpserver installed.
    from pytest_httpserver.httpserver import HTTPServer

    host, port = httpserver_listen_address
    if not host:
        host = HTTPServer.DEFAULT_LISTEN_HOST
    if not port:
        port = HTTPServer.DEFAULT_LISTEN_PORT

    server = HTTPServer(host=host, port=port, ssl_context=httpserver_ssl_context)
    server.start()
    yield server
    server.clear()
    if server.is_running():
        server.stop()


@pytest.fixture(scope="function")
def event_routing_config(request, monkeypatch):
    """
    fixture to toggle event routing modes.
    Modes:
      - "default": Uses the existing core_worker to gcs code path.
      - "aggregator": Enable publishing events to GCS through the Aggregator agent.
    """
    mode = getattr(request, "param", "default")
    # clear envs to ensure default behavior
    monkeypatch.delenv(
        "RAY_DASHBOARD_AGGREGATOR_AGENT_PUBLISH_EVENTS_TO_GCS", raising=False
    )
    monkeypatch.delenv("RAY_enable_core_worker_ray_event_to_aggregator", raising=False)

    if mode == "aggregator":
        print("using aggregator mode")
        # Enable aggregator path in core worker
        monkeypatch.setenv("RAY_enable_core_worker_ray_event_to_aggregator", "1")
        # Explicitly disable core worker to GCS so that all events are only sent to GCS once (through the aggregator pathway)
        monkeypatch.setenv("RAY_enable_core_worker_task_event_to_gcs", "0")
        # Ensure aggregator agent publishes to GCS
        monkeypatch.setenv(
            "RAY_DASHBOARD_AGGREGATOR_AGENT_PUBLISH_EVENTS_TO_GCS", "True"
        )
    yield


@pytest.fixture
def cleanup_auth_token_env():
    """Reset authentication environment variables, files, and caches."""

    with authentication_env_guard():
        clear_auth_token_sources(remove_default=True)
        reset_auth_token_state()
        yield
        reset_auth_token_state()


@pytest.fixture(autouse=False)
def clean_token_sources(cleanup_auth_token_env):
    """Ensure authentication-related state is clean around each test."""
    clear_auth_token_sources(remove_default=True)
    reset_auth_token_state()

    yield

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

    subprocess.run(
        ["ray", "stop", "--force"],
        capture_output=True,
        timeout=60,
        check=False,
    )

    reset_auth_token_state()


@pytest.fixture
def setup_cluster_with_token_auth(cleanup_auth_token_env):
    """Spin up a Ray cluster with token authentication enabled."""

    test_token = "test_token_12345678901234567890123456789012"
    set_auth_mode("token")
    set_env_auth_token(test_token)
    reset_auth_token_state()

    cluster = Cluster()
    # Use dynamic port to avoid port conflicts on Windows where sockets
    # linger in TIME_WAIT state between tests
    cluster.add_node(dashboard_agent_listen_port=find_free_port())

    try:
        context = ray.init(address=cluster.address)
        dashboard_url = context.address_info["webui_url"]
        yield {
            "cluster": cluster,
            "dashboard_url": f"http://{dashboard_url}",
            "token": test_token,
        }
    finally:
        ray.shutdown()
        cluster.shutdown()


@pytest.fixture
def setup_cluster_without_token_auth(cleanup_auth_token_env):
    """Spin up a Ray cluster with authentication disabled."""

    set_auth_mode("disabled")
    clear_auth_token_sources(remove_default=True)
    reset_auth_token_state()

    cluster = Cluster()
    # Use dynamic port to avoid port conflicts on Windows where sockets
    # linger in TIME_WAIT state between tests
    cluster.add_node(dashboard_agent_listen_port=find_free_port())

    try:
        context = ray.init(address=cluster.address)
        dashboard_url = context.address_info["webui_url"]
        yield {
            "cluster": cluster,
            "dashboard_url": f"http://{dashboard_url}",
        }
    finally:
        ray.shutdown()
        cluster.shutdown()


@pytest.fixture
def ray_start_cluster_with_zero_copy_tensors(monkeypatch):
    """Start a Ray cluster with zero-copy PyTorch tensors enabled."""
    with monkeypatch.context() as m:
        # Enable zero-copy sharing of PyTorch tensors in Ray
        m.setenv("RAY_ENABLE_ZERO_COPY_TORCH_TENSORS", "1")

        # Initialize Ray with the required environment variable.
        ray.init(runtime_env={"env_vars": {"RAY_ENABLE_ZERO_COPY_TORCH_TENSORS": "1"}})

        # Yield control to the test session
        yield

        # Shutdown Ray after tests complete
        ray.shutdown()