import os
import platform
import random
import subprocess
import sys
import tempfile

import numpy as np
import pytest

import ray
from ray._common.test_utils import (
    run_string_as_driver,
    wait_for_condition,
)
from ray._private.external_storage import (
    FileSystemStorage,
)
from ray.tests.test_object_spilling import is_dir_empty

# Note: Disk write speed can be as low as 6 MiB/s in AWS Mac instances, so we have to
# increase the timeout.
pytestmark = [pytest.mark.timeout(900 if platform.system() == "Darwin" else 180)]
condition_wait_timeout = 20 if os.getenv("RAY_DEBUG_MODE") == "1" else 10


def test_delete_objects(object_spilling_config, shutdown_only):
    # Limit our object store to 75 MiB of memory.
    object_spilling_config, temp_folder = object_spilling_config

    ray_context = ray.init(
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 1,
            "min_spilling_size": 0,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
        },
    )
    arr = np.random.rand(1024 * 1024)  # 8 MB data
    replay_buffer = []

    for _ in range(80):
        ref = None
        while ref is None:
            ref = ray.put(arr)
            replay_buffer.append(ref)

    print("-----------------------------------")

    del replay_buffer
    del ref
    wait_for_condition(
        lambda: is_dir_empty(temp_folder, ray_context["node_id"]),
        timeout=condition_wait_timeout,
    )


def test_delete_objects_delete_while_creating(object_spilling_config, shutdown_only):
    # Limit our object store to 75 MiB of memory.
    object_spilling_config, temp_folder = object_spilling_config

    ray_context = ray.init(
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 4,
            "min_spilling_size": 0,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
        },
    )
    arr = np.random.rand(1024 * 1024)  # 8 MB data
    replay_buffer = []

    for _ in range(80):
        ref = None
        while ref is None:
            ref = ray.put(arr)
            replay_buffer.append(ref)
        # Remove the replay buffer with 60% probability.
        if random.randint(0, 9) < 6:
            replay_buffer.pop()

    # Do random sampling.
    for _ in range(200):
        ref = random.choice(replay_buffer)
        sample = ray.get(ref, timeout=None)
        assert np.array_equal(sample, arr)

    # After all, make sure all objects are killed without race condition.
    del replay_buffer
    del ref
    wait_for_condition(
        lambda: is_dir_empty(temp_folder, ray_context["node_id"]),
        timeout=condition_wait_timeout,
    )


@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
def test_delete_objects_on_worker_failure(object_spilling_config, shutdown_only):
    # Limit our object store to 75 MiB of memory.
    object_spilling_config, temp_folder = object_spilling_config

    ray_context = ray.init(
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 4,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
            "min_spilling_size": 0,
            # ↓↓↓ make cleanup fast/consistent in CI
            "object_timeout_milliseconds": 200,
            "local_gc_min_interval_s": 1,
        },
    )

    arr = np.random.rand(1024 * 1024)  # 8 MB data

    @ray.remote
    class Actor:
        def __init__(self):
            self.replay_buffer = []

        def get_pid(self):
            return os.getpid()

        def create_objects(self):
            for _ in range(80):
                ref = None
                while ref is None:
                    ref = ray.put(arr)
                    self.replay_buffer.append(ref)
                # Remove the replay buffer with 60% probability.
                if random.randint(0, 9) < 6:
                    self.replay_buffer.pop()

            # Do random sampling.
            for _ in range(200):
                ref = random.choice(self.replay_buffer)
                sample = ray.get(ref, timeout=None)
                assert np.array_equal(sample, arr)

    a = Actor.remote()
    actor_pid = ray.get(a.get_pid.remote())
    ray.get(a.create_objects.remote())
    os.kill(actor_pid, 9)

    def wait_until_actor_dead():
        try:
            ray.get(a.get_pid.remote())
        except ray.exceptions.RayActorError:
            return True
        return False

    wait_for_condition(wait_until_actor_dead, timeout=condition_wait_timeout)

    # After all, make sure all objects are deleted upon worker failures.
    wait_for_condition(
        lambda: is_dir_empty(temp_folder, ray_context["node_id"]),
        timeout=condition_wait_timeout,
    )


@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
def test_delete_file_non_exists(shutdown_only, tmp_path):
    ray_context = ray.init()

    def create_spilled_files(num_files):
        spilled_files = []
        uris = []
        for _ in range(3):
            fd, path = tempfile.mkstemp()
            with os.fdopen(fd, "w") as tmp:
                tmp.write("stuff")
            spilled_files.append(path)
            uris.append((path + "?offset=0&size=10").encode("ascii"))
        return spilled_files, uris

    for storage in [
        FileSystemStorage(ray_context["node_id"], "/tmp"),
    ]:
        spilled_files, uris = create_spilled_files(3)
        storage.delete_spilled_objects(uris)
        for file in spilled_files:
            assert not os.path.exists(file)

        # delete should succeed even if some files doesn't exist.
        spilled_files1, uris1 = create_spilled_files(3)
        spilled_files += spilled_files1
        uris += uris1
        storage.delete_spilled_objects(uris)
        for file in spilled_files:
            assert not os.path.exists(file)


@pytest.mark.skipif(
    platform.system() in ["Windows"], reason="Failing on Windows and MacOS."
)
def test_delete_objects_multi_node(
    multi_node_object_spilling_config, ray_start_cluster
):
    # Limit our object store to 75 MiB of memory.
    object_spilling_config, temp_folder = multi_node_object_spilling_config

    cluster = ray_start_cluster
    # Head node.
    cluster.add_node(
        num_cpus=1,
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 2,
            "min_spilling_size": 20 * 1024 * 1024,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
        },
    )
    ray.init(address=cluster.address)
    # Add 2 worker nodes.
    worker_node1 = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    worker_node2 = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
    cluster.wait_for_nodes()

    arr = np.random.rand(1024 * 1024)  # 8 MB data

    @ray.remote(num_cpus=1)
    class Actor:
        def __init__(self):
            self.replay_buffer = []

        def ping(self):
            return

        def create_objects(self):
            for _ in range(80):
                ref = None
                while ref is None:
                    ref = ray.put(arr)
                    self.replay_buffer.append(ref)
                # Remove the replay buffer with 60% probability.
                if random.randint(0, 9) < 6:
                    self.replay_buffer.pop()

            # Do random sampling.
            for _ in range(50):
                ref = random.choice(self.replay_buffer)
                sample = ray.get(ref, timeout=10)
                assert np.array_equal(sample, arr)

    actors = [Actor.remote() for _ in range(3)]
    ray.get([actor.create_objects.remote() for actor in actors])

    def wait_until_actor_dead(actor):
        try:
            ray.get(actor.ping.remote())
        except ray.exceptions.RayActorError:
            return True
        return False

    # Kill actors to remove all references.
    for actor in actors:
        ray.kill(actor)
        wait_for_condition(
            lambda: wait_until_actor_dead(actor), timeout=condition_wait_timeout
        )
    # The multi node deletion should work.
    wait_for_condition(
        lambda: is_dir_empty(temp_folder, worker_node1.node_id),
        timeout=condition_wait_timeout,
    )
    wait_for_condition(
        lambda: is_dir_empty(temp_folder, worker_node2.node_id),
        timeout=condition_wait_timeout,
    )


def test_fusion_objects(fs_only_object_spilling_config, shutdown_only):
    # Limit our object store to 75 MiB of memory.
    object_spilling_config, temp_folder = fs_only_object_spilling_config
    min_spilling_size = 10 * 1024 * 1024
    ray.init(
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 3,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
            "min_spilling_size": min_spilling_size,
        },
    )
    replay_buffer = []
    solution_buffer = []
    buffer_length = 100

    # Create objects of more than 800 MiB.
    for _ in range(buffer_length):
        ref = None
        while ref is None:
            multiplier = random.choice([1, 2, 3])
            arr = np.random.rand(multiplier * 1024 * 1024)
            ref = ray.put(arr)
            replay_buffer.append(ref)
            solution_buffer.append(arr)

    print("-----------------------------------")
    # randomly sample objects
    for _ in range(1000):
        index = random.choice(list(range(buffer_length)))
        ref = replay_buffer[index]
        solution = solution_buffer[index]
        sample = ray.get(ref, timeout=None)
        assert np.array_equal(sample, solution)

    is_test_passing = False
    for path in temp_folder.iterdir():
        # Under the temp_folder, there should be a folder called
        # "ray_spilled_objects[_<node_id>]", which contains the
        # spilled objects.
        for spilled_objects_path in path.iterdir():
            file_size = spilled_objects_path.stat().st_size
            # Make sure there are at least one
            # file_size that exceeds the min_spilling_size.
            # If we don't fusion correctly, this cannot happen.
            if file_size >= min_spilling_size:
                is_test_passing = True
    assert is_test_passing


# https://github.com/ray-project/ray/issues/12912
def test_release_resource(object_spilling_config, shutdown_only):
    object_spilling_config, temp_folder = object_spilling_config
    ray.init(
        num_cpus=1,
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 1,
            "automatic_object_spilling_enabled": True,
            "object_spilling_config": object_spilling_config,
        },
    )
    plasma_obj = ray.put(np.ones(50 * 1024 * 1024, dtype=np.uint8))
    for _ in range(5):
        ray.put(np.ones(50 * 1024 * 1024, dtype=np.uint8))  # Force spilling

    @ray.remote
    def sneaky_task_tries_to_steal_released_resources():
        print("resources were released!")

    @ray.remote
    def f(dep):
        while True:
            try:
                ray.get(dep[0], timeout=0.001)
            except ray.exceptions.GetTimeoutError:
                pass

    done = f.remote([plasma_obj])  # noqa
    canary = sneaky_task_tries_to_steal_released_resources.remote()
    ready, _ = ray.wait([canary], timeout=2)
    assert not ready


def test_spill_objects_on_object_transfer(
    object_spilling_config, ray_start_cluster_enabled
):
    object_spilling_config, _ = object_spilling_config
    # This test checks that objects get spilled to make room for transferred
    # objects.
    cluster = ray_start_cluster_enabled
    object_size = int(1e7)
    num_objects = 10
    num_tasks = 10
    # Head node can fit all of the objects at once.
    cluster.add_node(
        num_cpus=0,
        object_store_memory=2 * num_tasks * num_objects * object_size,
        _system_config={
            "max_io_workers": 1,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
            "min_spilling_size": 0,
        },
    )
    cluster.wait_for_nodes()
    ray.init(address=cluster.address)

    # Worker node can fit 1 tasks at a time.
    cluster.add_node(num_cpus=1, object_store_memory=1.5 * num_objects * object_size)
    cluster.wait_for_nodes()

    @ray.remote
    def foo(*args):
        return

    @ray.remote
    def allocate(*args):
        return np.zeros(object_size, dtype=np.uint8)

    # Allocate some objects that must be spilled to make room for foo's
    # arguments.
    allocated = [allocate.remote() for _ in range(num_objects)]
    ray.get(allocated)
    print("done allocating")

    args = []
    for _ in range(num_tasks):
        task_args = [
            ray.put(np.zeros(object_size, dtype=np.uint8)) for _ in range(num_objects)
        ]
        args.append(task_args)

    # Check that tasks scheduled to the worker node have enough room after
    # spilling.
    tasks = [foo.remote(*task_args) for task_args in args]
    ray.get(tasks)


@pytest.mark.skipif(
    platform.system() in ["Windows"], reason="Failing on Windows and Mac."
)
def test_file_deleted_when_driver_exits(tmp_path, shutdown_only):
    temp_folder = tmp_path / "spill"
    temp_folder.mkdir()

    driver = """
import json
import os
import signal
import numpy as np
import ray
ray.init(
    object_store_memory=75 * 1024 * 1024,
    _system_config={{
        "max_io_workers": 2,
        "min_spilling_size": 0,
        "automatic_object_spilling_enabled": True,
        "object_store_full_delay_ms": 100,
        "object_spilling_config": json.dumps({{
            "type": "filesystem",
            "params": {{
                "directory_path": "{temp_dir}"
            }}
        }}),
    }})
arr = np.random.rand(1024 * 1024)  # 8 MB data
replay_buffer = []
# Spill lots of objects
for _ in range(30):
    ref = None
    while ref is None:
        ref = ray.put(arr)
        replay_buffer.append(ref)
# Send sigterm to itself.
signum = {signum}
sig = None
if signum == 2:
    sig = signal.SIGINT
elif signum == 15:
    sig = signal.SIGTERM
os.kill(os.getpid(), sig)
"""

    # Run a driver with sigint.
    print("Sending sigint...")
    with pytest.raises(subprocess.CalledProcessError):
        print(run_string_as_driver(driver.format(temp_dir=str(temp_folder), signum=2)))
    # node_id is not actually used in the following check, so we pass in a dummy one
    wait_for_condition(
        lambda: is_dir_empty(temp_folder, "dummy_node_id", append_path=False),
        timeout=condition_wait_timeout,
    )


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