import sys
import time
from math import ceil

import pytest

import ray
from ray._common.test_utils import (
    PrometheusTimeseries,
    wait_for_condition,
)
from ray._common.utils import get_system_memory
from ray._private import (
    ray_constants,
)
from ray._private.grpc_utils import init_grpc_channel
from ray._private.state_api_test_utils import verify_failed_task
from ray._private.test_utils import raw_metric_timeseries
from ray._private.utils import get_used_memory
from ray.util.state.state_manager import StateDataSourceClient

memory_usage_threshold = 0.5
task_oom_retries = 1
memory_monitor_refresh_ms = 100
expected_worker_eviction_message = (
    "worker(s) were killed due to the node running low on memory"
)


def get_local_state_client():
    gcs_channel = init_grpc_channel(
        ray.worker._global_node.gcs_address,
        ray_constants.GLOBAL_GRPC_OPTIONS,
        asynchronous=True,
    )

    gcs_client = ray._private.worker.global_worker.gcs_client
    return StateDataSourceClient(gcs_channel, gcs_client)


@pytest.fixture
def ray_with_memory_monitor(shutdown_only):
    with ray.init(
        num_cpus=1,
        object_store_memory=100 * 1024 * 1024,
        _system_config={
            "memory_usage_threshold": memory_usage_threshold,
            "memory_monitor_refresh_ms": memory_monitor_refresh_ms,
            "metrics_report_interval_ms": 100,
            "task_failure_entry_ttl_ms": 2 * 60 * 1000,
            "task_oom_retries": task_oom_retries,
            "min_memory_free_bytes": -1,
            "task_oom_retry_delay_base_ms": 0,
        },
    ) as addr:
        yield addr


@ray.remote
def allocate_memory(
    allocate_bytes: int,
    num_chunks: int = 10,
    allocate_interval_s: float = 0,
    post_allocate_sleep_s: float = 0,
):
    start = time.time()
    chunks = []
    # divide by 8 as each element in the array occupies 8 bytes
    bytes_per_chunk = allocate_bytes / 8 / num_chunks
    for _ in range(num_chunks):
        chunks.append([0] * ceil(bytes_per_chunk))
        time.sleep(allocate_interval_s)
    end = time.time()
    time.sleep(post_allocate_sleep_s)
    return end - start


@ray.remote
class Leaker:
    def __init__(self):
        self.leaks = []

    def allocate(self, allocate_bytes: int, sleep_time_ms: int = 0):
        # divide by 8 as each element in the array occupies 8 bytes
        new_list = [0] * ceil(allocate_bytes / 8)
        self.leaks.append(new_list)

        time.sleep(sleep_time_ms / 1000)

    def get_worker_id(self):
        return ray._private.worker.global_worker.core_worker.get_worker_id().hex()

    def get_actor_id(self):
        return ray._private.worker.global_worker.core_worker.get_actor_id().hex()


def get_additional_bytes_to_reach_memory_usage_pct(pct: float) -> int:
    used = get_used_memory()
    total = get_system_memory()
    bytes_needed = int(total * pct) - used
    assert bytes_needed > 0, "node has less memory than what is requested"
    return bytes_needed


def has_metric_tagged_with_value(
    addr, tag, value, timeseries: PrometheusTimeseries
) -> bool:
    metrics = raw_metric_timeseries(addr, timeseries)
    for name, samples in metrics.items():
        for sample in samples:
            if tag in set(sample.labels.values()) and sample.value == value:
                return True
    return False


@pytest.mark.skipif(
    sys.platform != "linux" and sys.platform != "linux2",
    reason="memory monitor is currently only supported on Linux",
)
@pytest.mark.parametrize("restartable", [False, True])
def test_restartable_actor_throws_oom_error(ray_with_memory_monitor, restartable: bool):
    addr = ray_with_memory_monitor
    if restartable:
        leaker = Leaker.options(max_restarts=1, max_task_retries=1).remote()
    else:
        leaker = Leaker.options(max_restarts=0, max_task_retries=0).remote()

    bytes_to_alloc = get_additional_bytes_to_reach_memory_usage_pct(
        memory_usage_threshold + 0.1
    )
    with pytest.raises(ray.exceptions.OutOfMemoryError):
        ray.get(leaker.allocate.remote(bytes_to_alloc, memory_monitor_refresh_ms * 3))

    timeseries = PrometheusTimeseries()
    wait_for_condition(
        has_metric_tagged_with_value,
        timeout=10,
        retry_interval_ms=100,
        addr=addr,
        tag="MemoryManager.ActorEviction.Total",
        value=2.0 if restartable else 1.0,
        timeseries=timeseries,
    )

    wait_for_condition(
        has_metric_tagged_with_value,
        timeout=10,
        retry_interval_ms=100,
        addr=addr,
        tag="Leaker.__init__",
        value=2.0 if restartable else 1.0,
        timeseries=timeseries,
    )


@pytest.mark.skipif(
    sys.platform != "linux" and sys.platform != "linux2",
    reason="memory monitor is currently only supported on Linux",
)
def test_non_retryable_task_killed_by_memory_monitor_with_oom_error(
    ray_with_memory_monitor,
):
    addr = ray_with_memory_monitor
    bytes_to_alloc = get_additional_bytes_to_reach_memory_usage_pct(1.1)
    with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
        ray.get(allocate_memory.options(max_retries=0).remote(bytes_to_alloc))

    timeseries = PrometheusTimeseries()
    wait_for_condition(
        has_metric_tagged_with_value,
        timeout=10,
        retry_interval_ms=100,
        addr=addr,
        tag="MemoryManager.TaskEviction.Total",
        value=1.0,
        timeseries=timeseries,
    )
    wait_for_condition(
        has_metric_tagged_with_value,
        timeout=10,
        retry_interval_ms=100,
        addr=addr,
        tag="allocate_memory",
        value=1.0,
        timeseries=timeseries,
    )


@pytest.mark.asyncio
@pytest.mark.skipif(
    sys.platform != "linux" and sys.platform != "linux2",
    reason="memory monitor is currently only supported on Linux",
)
async def test_actor_oom_logs_error(ray_with_memory_monitor):
    first_actor = Leaker.options(name="first_random_actor", max_restarts=0).remote()
    ray.get(first_actor.get_worker_id.remote())

    oom_actor = Leaker.options(name="the_real_oom_actor", max_restarts=0).remote()
    worker_id = ray.get(oom_actor.get_worker_id.remote())
    actor_id = ray.get(oom_actor.get_actor_id.remote())

    bytes_to_alloc = get_additional_bytes_to_reach_memory_usage_pct(1)
    with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
        ray.get(
            oom_actor.allocate.remote(bytes_to_alloc, memory_monitor_refresh_ms * 3)
        )

    state_api_client = get_local_state_client()
    result = await state_api_client.get_all_worker_info(timeout=5, limit=10)
    verified = False
    for worker in result.worker_table_data:
        if worker.worker_address.worker_id.hex() == worker_id:
            assert expected_worker_eviction_message in worker.exit_detail
            verified = True
    assert verified

    result = await state_api_client.get_all_actor_info(timeout=5, limit=10)
    verified = False
    for actor in result.actor_table_data:
        if actor.actor_id.hex() == actor_id:
            assert actor.death_cause
            assert actor.death_cause.oom_context
            assert (
                expected_worker_eviction_message
                in actor.death_cause.oom_context.error_message
            )
            verified = True
    assert verified

    # TODO(clarng): verify log info once state api can dump log info


@pytest.mark.asyncio
@pytest.mark.skipif(
    sys.platform != "linux" and sys.platform != "linux2",
    reason="memory monitor is currently only supported on Linux",
)
async def test_task_oom_logs_error(ray_with_memory_monitor):
    bytes_to_alloc = get_additional_bytes_to_reach_memory_usage_pct(1)
    with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
        ray.get(
            allocate_memory.options(max_retries=0, name="allocate_memory").remote(
                allocate_bytes=bytes_to_alloc,
                allocate_interval_s=0,
                post_allocate_sleep_s=1000,
            )
        )

    state_api_client = get_local_state_client()
    result = await state_api_client.get_all_worker_info(timeout=5, limit=10)
    verified = False
    for worker in result.worker_table_data:
        if worker.exit_detail:
            assert expected_worker_eviction_message in worker.exit_detail
        verified = True
    assert verified

    wait_for_condition(
        verify_failed_task,
        name="allocate_memory",
        error_type="OUT_OF_MEMORY",
        error_message="worker(s) were killed due to the node running low on memory",
    )

    # TODO(clarng): verify log info once state api can dump log info


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