import json
import platform
import random
import re
import shutil
import sys
import time
import zlib
from collections import defaultdict

import numpy as np
import pytest

import ray
from ray._common.test_utils import wait_for_condition
from ray.cluster_utils import Cluster, cluster_not_supported
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)]


@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
def test_multiple_directories(tmp_path, shutdown_only):
    num_dirs = 3
    temp_dirs = []
    for i in range(num_dirs):
        temp_folder = tmp_path / f"spill_{i}"
        temp_folder.mkdir()
        temp_dirs.append(temp_folder)

    # Limit our object store to 75 MiB of memory.
    min_spilling_size = 0
    object_spilling_config = json.dumps(
        {
            "type": "filesystem",
            "params": {"directory_path": [str(directory) for directory in temp_dirs]},
        }
    )
    ray_context = ray.init(
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 5,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
            "min_spilling_size": min_spilling_size,
        },
    )

    arr = np.ones(74 * 1024 * 1024, dtype=np.uint8)  # 74MB.
    object_refs = []
    # Now the storage is full.
    object_refs.append(ray.put(arr))

    num_object_spilled = 20
    for _ in range(num_object_spilled):
        object_refs.append(ray.put(arr))

    num_files = defaultdict(int)
    for temp_dir in temp_dirs:
        # Under temp_dir there are spilled_objects_dir(s) with name pattern
        # "ray_spilled_objects[_<node_id>]", each containing spilled object files.
        for spilled_objects_dir in temp_folder.iterdir():
            for path in spilled_objects_dir.iterdir():
                num_files[str(temp_folder)] += 1

    for ref in object_refs:
        assert np.array_equal(ray.get(ref), arr)

    print("Check distribution...")
    min_count = 5
    is_distributed = [n_files >= min_count for n_files in num_files.values()]
    assert all(is_distributed)

    print("Check deletion...")
    # Empty object refs.
    object_refs = []
    # Add a new small object so that the last entry is evicted and we don't
    # exceed the spill threshold.
    ref = ray.put(np.ones(5 * 1024 * 1024, dtype=np.uint8))
    for temp_dir in temp_dirs:
        temp_folder = temp_dir
        wait_for_condition(lambda: is_dir_empty(temp_folder, ray_context["node_id"]))

    # Now kill ray and see all directories are deleted.
    print("Check directories are deleted...")
    ray.shutdown()
    for temp_dir in temp_dirs:
        wait_for_condition(lambda: is_dir_empty(temp_dir, ray_context["node_id"]))


def _check_spilled(num_objects_spilled=0):
    def ok():
        s = ray._private.internal_api.memory_summary(stats_only=True)
        if num_objects_spilled == 0:
            return "Spilled " not in s

        m = re.search(r"Spilled (\d+) MiB, (\d+) objects", s)
        if m is not None:
            actual_num_objects = int(m.group(2))
            return actual_num_objects >= num_objects_spilled

        return False

    wait_for_condition(ok, timeout=90, retry_interval_ms=5000)


def _test_object_spilling_threshold(thres, num_objects, num_objects_spilled):
    try:
        ray.init(
            object_store_memory=2_200_000_000,
            _system_config={"object_spilling_threshold": thres} if thres else {},
        )
        objs = []
        for _ in range(num_objects):
            objs.append(ray.put(np.empty(200_000_000, dtype=np.uint8)))
        if num_objects_spilled == 0:
            time.sleep(10)  # Wait for spilling to happen
        _check_spilled(num_objects_spilled)
    finally:
        ray.shutdown()


@pytest.mark.skipif(platform.system() != "Linux", reason="Failing on Windows/macOS.")
def test_object_spilling_threshold_default():
    _test_object_spilling_threshold(None, 10, 5)


@pytest.mark.skipif(platform.system() != "Linux", reason="Failing on Windows/macOS.")
def test_object_spilling_threshold_1_0():
    _test_object_spilling_threshold(1.0, 10, 0)


@pytest.mark.skipif(platform.system() != "Linux", reason="Failing on Windows/macOS.")
def test_object_spilling_threshold_0_1():
    _test_object_spilling_threshold(0.1, 10, 5)


def test_partial_retval_allocation(ray_start_cluster_enabled):
    cluster = ray_start_cluster_enabled
    cluster.add_node(object_store_memory=100 * 1024 * 1024)
    ray.init(cluster.address)

    @ray.remote(num_returns=4)
    def f():
        return [np.zeros(50 * 1024 * 1024, dtype=np.uint8) for _ in range(4)]

    ret = f.remote()
    for obj in ret:
        obj = ray.get(obj)
        print(obj.size)


def test_pull_spilled_object(
    ray_start_cluster_enabled, multi_node_object_spilling_config, shutdown_only
):
    cluster = ray_start_cluster_enabled
    object_spilling_config, _ = multi_node_object_spilling_config

    # Head node.
    cluster.add_node(
        num_cpus=1,
        resources={"custom": 0},
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 2,
            "min_spilling_size": 1 * 1024 * 1024,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
        },
    )
    ray.init(cluster.address)

    # add 1 worker node
    cluster.add_node(
        num_cpus=1, resources={"custom": 1}, object_store_memory=75 * 1024 * 1024
    )
    cluster.wait_for_nodes()

    @ray.remote(num_cpus=1, resources={"custom": 1})
    def create_objects():
        results = []
        for size in range(5):
            arr = np.random.rand(size * 1024 * 1024)
            hash_value = zlib.crc32(arr.tobytes())
            results.append([ray.put(arr), hash_value])
        # ensure the objects are spilled
        arr = np.random.rand(5 * 1024 * 1024)
        ray.get(ray.put(arr))
        ray.get(ray.put(arr))
        return results

    @ray.remote(num_cpus=1, resources={"custom": 0})
    def get_object(arr):
        return zlib.crc32(arr.tobytes())

    results = ray.get(create_objects.remote())
    for value_ref, hash_value in results:
        hash_value1 = ray.get(get_object.remote(value_ref))
        assert hash_value == hash_value1


# TODO(chenshen): fix error handling when spilled file
# missing/corrupted
@pytest.mark.skipif(True, reason="Currently hangs.")
def test_pull_spilled_object_failure(object_spilling_config, ray_start_cluster):
    object_spilling_config, temp_folder = object_spilling_config
    cluster = ray_start_cluster

    # Head node.
    cluster.add_node(
        num_cpus=1,
        resources={"custom": 0},
        object_store_memory=75 * 1024 * 1024,
        _system_config={
            "max_io_workers": 2,
            "min_spilling_size": 1 * 1024 * 1024,
            "automatic_object_spilling_enabled": True,
            "object_store_full_delay_ms": 100,
            "object_spilling_config": object_spilling_config,
        },
    )
    ray.init(cluster.address)

    # add 1 worker node
    cluster.add_node(
        num_cpus=1, resources={"custom": 1}, object_store_memory=75 * 1024 * 1024
    )
    cluster.wait_for_nodes()

    @ray.remote(num_cpus=1, resources={"custom": 1})
    def create_objects():
        arr = np.random.rand(5 * 1024 * 1024)
        hash_value = zlib.crc32(arr.tobytes())
        results = [ray.put(arr), hash_value]
        # ensure the objects are spilled
        arr = np.random.rand(5 * 1024 * 1024)
        ray.get(ray.put(arr))
        ray.get(ray.put(arr))
        return results

    @ray.remote(num_cpus=1, resources={"custom": 0})
    def get_object(arr):
        return zlib.crc32(arr.tobytes())

    [ref, hash_value] = ray.get(create_objects.remote())

    # remove spilled file
    shutil.rmtree(temp_folder)

    hash_value1 = ray.get(get_object.remote(ref))
    assert hash_value == hash_value1


@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
def test_spill_dir_cleanup_on_node_removal(fs_only_object_spilling_config):
    object_spilling_config, temp_folder = fs_only_object_spilling_config
    cluster = Cluster()
    cluster.add_node(
        num_cpus=0,
        object_store_memory=75 * 1024 * 1024,
        _system_config={"object_spilling_config": object_spilling_config},
    )
    ray.init(address=cluster.address)
    node2 = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)

    # This task will run on node 2 because node 1 has no CPU resource
    @ray.remote(num_cpus=1)
    def run_workload():
        ids = []
        for _ in range(2):
            arr = np.random.rand(5 * 1024 * 1024)  # 40 MB
            ids.append(ray.put(arr))
        return ids

    ids = ray.get(run_workload.remote())
    node2_id = node2.node_id
    assert not is_dir_empty(temp_folder, node2_id)

    # Kill node 2
    cluster.remove_node(node2)

    # Verify that the spill folder is cleaned up upon node removal
    assert is_dir_empty(temp_folder, node2_id)

    # We hold the object refs to prevent them from being deleted
    # due to out of scope.
    del ids
    ray.shutdown()
    cluster.shutdown()


def test_spill_deadlock(object_spilling_config, shutdown_only):
    object_spilling_config, _ = object_spilling_config
    # Limit our object store to 75 MiB of memory.
    ray.init(
        object_store_memory=75 * 1024 * 1024,
        _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,
        },
    )
    arr = np.random.rand(1024 * 1024)  # 8 MB data
    replay_buffer = []

    # Create objects of more than 400 MiB.
    for _ in range(50):
        ref = None
        while ref is None:
            ref = ray.put(arr)
            replay_buffer.append(ref)
        # This is doing random sampling with 50% prob.
        if random.randint(0, 9) < 5:
            for _ in range(5):
                ref = random.choice(replay_buffer)
                sample = ray.get(ref, timeout=None)
                assert np.array_equal(sample, arr)


def test_spill_reconstruction_errors(ray_start_cluster, object_spilling_config):
    config = {
        "health_check_failure_threshold": 10,
        "health_check_period_ms": 100,
        "health_check_initial_delay_ms": 0,
        "max_direct_call_object_size": 100,
        "task_retry_delay_ms": 100,
        "object_timeout_milliseconds": 200,
    }
    cluster = ray_start_cluster
    # Head node with no resources.
    cluster.add_node(num_cpus=0, _system_config=config, object_store_memory=10**8)
    ray.init(address=cluster.address)
    # Node to place the initial object.
    node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)

    @ray.remote
    def put():
        return np.zeros(10**5, dtype=np.uint8)

    @ray.remote
    def check(x):
        return

    ref = put.remote()
    for _ in range(4):
        ray.get(check.remote(ref))
        cluster.remove_node(node_to_kill, allow_graceful=False)
        node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)

    # All reconstruction attempts used up. The object's value should now be an
    # error in the local store.
    # Force object spilling and check that it can complete.
    xs = []
    for _ in range(20):
        xs.append(ray.put(np.zeros(10**7, dtype=np.uint8)))
    for x in xs:
        ray.get(x, timeout=10)

    with pytest.raises(ray.exceptions.ObjectLostError):
        ray.get(ref)


def test_evict_secondary_copies_before_spill(ray_start_cluster, object_spilling_config):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=1, object_store_memory=10**8)
    ray.init(address=cluster.address)
    for _ in range(3):
        cluster.add_node(num_cpus=1, object_store_memory=10**8)
    wait_for_condition(lambda: ray.cluster_resources()["CPU"] >= 4)

    # Spread data onto all nodes.
    @ray.remote
    def gen():
        time.sleep(0.5)
        return np.ones(10 * 1024 * 1024, dtype=np.uint8)

    refs = [
        gen.options(scheduling_strategy="SPREAD").remote() for _ in range(16)
    ]  # 160MiB

    # Iterate over the data on the worker nodes from the head node.
    for i in range(10):
        for j, r in enumerate(refs):
            print("Iteration", i, j)
            ray.get(r)

    _check_spilled()


if __name__ == "__main__":

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