import os
import platform
import shutil
import sys
import tempfile
import time
from contextlib import contextmanager

import numpy as np
import pytest

import ray
from ray._common.test_utils import wait_for_condition
from ray.util.state import list_cluster_events


def calculate_capacity_threshold(disk_capacity_in_bytes):
    usage = shutil.disk_usage("/tmp")
    threshold = min(1, 1.0 - 1.0 * (usage.free - disk_capacity_in_bytes) / usage.total)
    return threshold


def get_current_usage():
    usage = shutil.disk_usage("/tmp")
    print(f"free: {usage.free} ")
    usage_percentage = 1.0 - 1.0 * usage.free / usage.total
    print(f"current usage: {usage_percentage}")
    return usage_percentage


@contextmanager
def create_tmp_file(n):
    tmp_dir = tempfile.mkdtemp(dir="/tmp")
    tmp_path = os.path.join(tmp_dir, "test.txt")
    with open(tmp_path, "wb") as f:
        f.write(os.urandom(n))
    try:
        yield tmp_path
    finally:
        os.remove(tmp_path)


@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
def test_put_out_of_disk(shutdown_only):
    local_fs_capacity_threshold = calculate_capacity_threshold(200 * 1024 * 1024)
    ray.init(
        num_cpus=1,
        object_store_memory=80 * 1024 * 1024,
        _system_config={
            "local_fs_capacity_threshold": local_fs_capacity_threshold,
            "local_fs_monitor_interval_ms": 10,
        },
    )
    assert get_current_usage() < local_fs_capacity_threshold
    ref = ray.put(np.random.rand(20 * 1024 * 1024))
    del ref
    # create a temp file so that the disk size is over the threshold.
    # ray.put doesn't work is that fallback allocation uses mmaped file
    # that doesn't neccssary allocate disk spaces.
    with create_tmp_file(250 * 1024 * 1024):
        wait_for_condition(lambda: get_current_usage() > local_fs_capacity_threshold)
        with pytest.raises(ray.exceptions.OutOfDiskError):
            ray.put(np.random.rand(20 * 1024 * 1024))
        # delete tmp file to reclaim space back.

    # Ideally get_current_usage() should immediately reflect the latest disk usage
    # after the tmp file is deleted, but somehow there is some delays on CI machines
    # so I use wait_for_condition here.
    wait_for_condition(lambda: get_current_usage() < local_fs_capacity_threshold)
    ray.put(np.random.rand(20 * 1024 * 1024))


@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
def test_task_returns(shutdown_only):
    local_fs_capacity_threshold = calculate_capacity_threshold(10 * 1024 * 1024)
    ray.init(
        num_cpus=1,
        object_store_memory=80 * 1024 * 1024,
        _system_config={
            "local_fs_capacity_threshold": local_fs_capacity_threshold,
            "local_fs_monitor_interval_ms": 10,
        },
    )

    # create a temp file so that the disk size is over the threshold.
    # ray.put doesn't work is that fallback allocation uses mmaped file
    # that doesn't neccssary allocate disk spaces.
    with create_tmp_file(250 * 1024 * 1024):
        assert get_current_usage() > local_fs_capacity_threshold
        time.sleep(1)

        @ray.remote
        def foo():
            time.sleep(1)
            return np.random.rand(20 * 1024 * 1024)  # 160 MB data

        try:
            ray.get(foo.remote())
        except ray.exceptions.RayTaskError as e:
            assert isinstance(e.cause, ray.exceptions.OutOfDiskError)


@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
def test_task_put(shutdown_only):
    local_fs_capacity_threshold = calculate_capacity_threshold(1 * 1024 * 1024)
    ray.init(
        num_cpus=1,
        object_store_memory=80 * 1024 * 1024,
        _system_config={
            "local_fs_capacity_threshold": local_fs_capacity_threshold,
            "local_fs_monitor_interval_ms": 10,
        },
    )

    # create a temp file so that the disk size is over the threshold.
    # ray.put doesn't work is that fallback allocation uses mmaped file
    # that doesn't neccssary allocate disk spaces.
    with create_tmp_file(250 * 1024 * 1024):
        assert get_current_usage() > local_fs_capacity_threshold
        time.sleep(1)

        @ray.remote
        def foo():
            ref = ray.put(np.random.rand(20 * 1024 * 1024))  # 160 MB data
            return ref

        try:
            ray.get(foo.remote())
        except ray.exceptions.RayTaskError as e:
            assert isinstance(e.cause, ray.exceptions.OutOfDiskError)


@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
def test_task_args(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(
        num_cpus=1,
        object_store_memory=80 * 1024 * 1024,
        _system_config={
            "local_fs_capacity_threshold": 0,
        },
        resources={"out_of_memory": 1},
    )
    cluster.add_node(
        num_cpus=1,
        object_store_memory=200 * 1024 * 1024,
        resources={"sufficient_memory": 1},
    )
    cluster.wait_for_nodes()
    ray.init(address=cluster.address)

    @ray.remote
    def foo():
        return np.random.rand(20 * 1024 * 1024)  # 160 MB data

    @ray.remote
    def bar(obj):
        print(obj)

    ref = foo.options(resources={"sufficient_memory": 1}).remote()
    try:
        ray.get(bar.options(resources={"out_of_memory": 1}).remote(ref))
    except ray.exceptions.RayTaskError as e:
        assert isinstance(e.cause, ray.exceptions.OutOfDiskError)


@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
def test_actor(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(
        num_cpus=1,
        object_store_memory=80 * 1024 * 1024,
        _system_config={
            "local_fs_capacity_threshold": 0,
        },
        resources={"out_of_memory": 1},
    )
    cluster.add_node(
        num_cpus=1,
        object_store_memory=200 * 1024 * 1024,
        resources={"sufficient_memory": 1},
    )
    cluster.wait_for_nodes()
    ray.init(address=cluster.address)

    @ray.remote
    def foo():
        return np.random.rand(20 * 1024 * 1024)  # 160 MB data

    @ray.remote
    class Actor:
        def __init__(self, obj):
            self._obj = obj

        def foo(self):
            print(self._obj)

        def args_ood(self, obj):
            print(obj)

        def return_ood(self):
            return np.random.rand(20 * 1024 * 1024)

    ref = foo.options(resources={"sufficient_memory": 1}).remote()
    with pytest.raises(ray.exceptions.RayActorError):
        a = Actor.options(resources={"out_of_memory": 0.001}).remote(ref)
        ray.get(a.foo.remote())

    a = Actor.options(resources={"out_of_memory": 1}).remote(1)
    ray.get(a.foo.remote())
    try:
        ray.get(a.args_ood.remote(ref))
    except ray.exceptions.RayTaskError as e:
        assert isinstance(e.cause, ray.exceptions.OutOfDiskError)

    ray.get(a.foo.remote())
    try:
        ray.get(a.return_ood.remote())
    except ray.exceptions.RayTaskError as e:
        assert isinstance(e.cause, ray.exceptions.OutOfDiskError)


@pytest.mark.skipif(platform.system() == "Windows", reason="Not targeting Windows")
def test_ood_events(shutdown_only):
    local_fs_capacity_threshold = calculate_capacity_threshold(200 * 1024 * 1024)
    ray.init(
        num_cpus=1,
        object_store_memory=80 * 1024 * 1024,
        _system_config={
            "local_fs_capacity_threshold": local_fs_capacity_threshold,
            "local_fs_monitor_interval_ms": 10,
        },
    )

    # create a temp file so that the disk size is over the threshold.
    # ray.put doesn't work is that fallback allocation uses mmaped file
    # that doesn't neccssary allocate disk spaces.
    with create_tmp_file(250 * 1024 * 1024):
        assert get_current_usage() > local_fs_capacity_threshold
        time.sleep(1)

        @ray.remote
        def foo():
            ref = ray.put(np.random.rand(20 * 1024 * 1024))  # 160 MB data
            return ref

        try:
            ray.get(foo.remote())
        except ray.exceptions.RayTaskError as e:
            assert isinstance(e.cause, ray.exceptions.OutOfDiskError)

    # Give it some time for events to appear.
    # TODO(core-team): provide some way to wait for events to be flushed.
    time.sleep(2)

    events = list_cluster_events()
    print(events)
    # There could be more than 1 event depending on the test timing.
    assert len(events) >= 1
    assert (
        "Object creation will fail if spilling is required" in events[0]["message"]
    )  # noqa
    assert events[0]["severity"] == "ERROR"


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