import os
import sys
import time

import numpy as np
import pytest

import ray
from ray._common.test_utils import Semaphore, wait_for_condition
from ray._private.internal_api import memory_summary
from ray.cluster_utils import Cluster, cluster_not_supported

# RayConfig to enable recording call sites during ObjectRej creations.
ray_config = {"record_ref_creation_sites": True}

# Unique strings.
DRIVER_PID = "Driver"
WORKER_PID = "Worker"
UNKNOWN_SIZE = " ? "

# Reference status.
PINNED_IN_MEMORY = "PINNED_IN_MEMORY"
LOCAL_REF = "LOCAL_REFERENCE"
USED_BY_PENDING_TASK = "USED_BY_PENDING_TASK"
CAPTURED_IN_OBJECT = "CAPTURED_IN_OBJECT"
ACTOR_HANDLE = "ACTOR_HANDLE"

# Call sites.
PUT_OBJ = "(put object)"
TASK_CALL_OBJ = "(task call)"
ACTOR_TASK_CALL_OBJ = "(actor call)"
DESER_TASK_ARG = "(deserialize task arg)"
# Only 22 characters can be matched because longer strings are wrapped around.
DESER_ACTOR_TASK_ARG = "(deserialize actor tas"

# Group by and sort by parameters.
NODE_ADDRESS = "node address"
STACK_TRACE = "stack trace"
PID = "pid"
OBJECT_SIZE = "object size"
REFERENCE_TYPE = "reference type"

# Task status.
WAITING_FOR_DEPENDENCIES = "PENDING_ARGS_AVAIL"
SCHEDULED = "PENDING_NODE_ASSIGNMENT"
FINISHED = "FINISHED"
WAITING_FOR_EXECUTION = "SUBMITTED_TO_WORKER"


def data_lines(memory_str):
    for line in memory_str.split("\n"):
        if (
            PINNED_IN_MEMORY in line
            or LOCAL_REF in line
            or USED_BY_PENDING_TASK in line
            or CAPTURED_IN_OBJECT in line
            or ACTOR_HANDLE in line
        ):
            yield line
        else:
            continue


def num_objects(memory_str):
    n = 0
    for line in data_lines(memory_str):
        n += 1
    return n


def count(memory_str, substr):
    substr = substr[:42]
    n = 0
    for line in memory_str.split("\n"):
        if substr in line:
            n += 1
    return n


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_driver_put_ref(ray_start_regular):
    address = ray_start_regular["address"]
    info = memory_summary(address)
    assert num_objects(info) == 0, info
    x_id = ray.put("HI")
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, DRIVER_PID) == 1, info
    assert count(info, WORKER_PID) == 0, info
    del x_id
    info = memory_summary(address)
    assert num_objects(info) == 0, info


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_worker_task_refs(ray_start_regular):
    address = ray_start_regular["address"]

    @ray.remote
    def f(y):
        from ray._private.internal_api import memory_summary

        x_id = ray.put("HI")
        info = memory_summary(address)
        del x_id
        return info

    x_id = f.remote(np.zeros(100000))
    info = ray.get(x_id)
    print(info)
    assert num_objects(info) == 4, info
    # Task argument plus task return ids.
    assert count(info, TASK_CALL_OBJ) == 2, info
    assert count(info, DRIVER_PID) == 2, info
    assert count(info, WORKER_PID) == 2, info
    assert count(info, LOCAL_REF) == 2, info
    assert count(info, PINNED_IN_MEMORY) == 1, info
    assert count(info, PUT_OBJ) == 1, info
    assert count(info, DESER_TASK_ARG) == 1, info
    assert count(info, UNKNOWN_SIZE) == 1, info

    print(ray_start_regular)
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, DRIVER_PID) == 1, info
    assert count(info, TASK_CALL_OBJ) == 1, info
    assert count(info, UNKNOWN_SIZE) == 0, info
    assert count(info, x_id.hex()) == 1, info

    del x_id
    info = memory_summary(address)
    assert num_objects(info) == 0, info


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_actor_task_refs(ray_start_regular):
    address = ray_start_regular["address"]

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

        def f(self, x):
            from ray._private.internal_api import memory_summary

            self.refs.append(x)
            return memory_summary(address)

    def make_actor():
        return Actor.remote()

    actor = make_actor()
    x_id = actor.f.remote(np.zeros(100000))
    info = ray.get(x_id)
    print(info)
    assert num_objects(info) == 4, info
    # Actor handle, task argument id, task return id.
    assert count(info, ACTOR_TASK_CALL_OBJ) == 3, info
    assert count(info, DRIVER_PID) == 3, info
    assert count(info, WORKER_PID) == 1, info
    assert count(info, LOCAL_REF) == 1, info
    assert count(info, PINNED_IN_MEMORY) == 1, info
    assert count(info, USED_BY_PENDING_TASK) == 1, info
    assert count(info, ACTOR_HANDLE) == 1, info
    assert count(info, DESER_ACTOR_TASK_ARG) == 1, info
    del x_id

    # These should accumulate in the actor.
    for _ in range(5):
        ray.get(actor.f.remote([ray.put(np.zeros(100000))]))
    info = memory_summary(address)
    print(info)
    assert count(info, DESER_ACTOR_TASK_ARG) == 5, info
    assert count(info, ACTOR_TASK_CALL_OBJ) == 1, info

    # Cleanup.
    del actor
    time.sleep(1)
    info = memory_summary(address)
    assert num_objects(info) == 0, info


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_nested_object_refs(ray_start_regular):
    address = ray_start_regular["address"]
    x_id = ray.put(np.zeros(100000))
    y_id = ray.put([x_id])
    z_id = ray.put([y_id])
    del x_id, y_id
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 3, info
    assert count(info, LOCAL_REF) == 1, info
    assert count(info, CAPTURED_IN_OBJECT) == 2, info
    del z_id


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_pinned_object_call_site(ray_start_regular):
    address = ray_start_regular["address"]
    # Local ref only.
    x_id = ray.put(np.zeros(100000))
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, LOCAL_REF) == 1, info
    assert count(info, PINNED_IN_MEMORY) == 0, info

    # Local ref + pinned buffer.
    buf = ray.get(x_id)
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, LOCAL_REF) == 0, info
    assert count(info, PINNED_IN_MEMORY) == 1, info

    # Just pinned buffer.
    del x_id
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 1, info
    assert count(info, LOCAL_REF) == 0, info
    assert count(info, PINNED_IN_MEMORY) == 1, info

    # Nothing.
    del buf
    info = memory_summary(address)
    print(info)
    assert num_objects(info) == 0, info


@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
def test_multi_node_stats(shutdown_only):
    # NOTE(mwtian): using env var only enables the feature on workers, while
    # using head_node_args={"_system_config": ray_config} only enables the
    # feature on the driver.
    os.environ["RAY_record_ref_creation_sites"] = "1"
    cluster = Cluster()
    for _ in range(2):
        cluster.add_node(num_cpus=1)

    ray.init(address=cluster.address)

    @ray.remote(num_cpus=1)
    class Actor:
        def __init__(self):
            self.ref = ray.put(np.zeros(100000))

        def ping(self):
            pass

    # Each actor will be on a different node.
    a = Actor.remote()
    b = Actor.remote()
    ray.get(a.ping.remote())
    ray.get(b.ping.remote())

    # Verify we have collected stats across the nodes.
    info = memory_summary(cluster.address)
    print(info)
    assert count(info, PUT_OBJ) == 2, info


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_group_by_sort_by(ray_start_regular):
    address = ray_start_regular["address"]

    @ray.remote
    def f(y):
        x_id = ray.put("HI")
        info_a = memory_summary(
            address, group_by="STACK_TRACE", sort_by="REFERENCE_TYPE"
        )
        info_b = memory_summary(address, group_by="NODE_ADDRESS", sort_by="OBJECT_SIZE")
        info_c = memory_summary(address, group_by="NODE_ADDRESS", sort_by="PID")
        del x_id
        return info_a, info_b, info_c

    x_id = f.remote(np.zeros(100000))
    info_a, info_b, info_c = ray.get(x_id)
    print(info_c)
    assert count(info_a, STACK_TRACE) == 7, info_a
    assert count(info_a, REFERENCE_TYPE) == 1, info_a
    assert count(info_b, NODE_ADDRESS) == 3, info_b
    assert count(info_b, OBJECT_SIZE) == 1, info_b
    assert count(info_c, NODE_ADDRESS) == 3, info_c
    assert count(info_c, PID) == 1, info_c


@pytest.mark.parametrize(
    "ray_start_regular", [{"_system_config": ray_config}], indirect=True
)
def test_memory_used_output(ray_start_regular):
    address = ray_start_regular["address"]
    _ = ray.put(np.ones(8 * 1024 * 1024, dtype=np.int8))

    info = memory_summary(address)
    print(info)
    assert count(info, "Plasma memory usage 8 MiB") == 1, info
    assert count(info, "8388862.0 B") == 2, info


def test_task_status(ray_start_regular):
    address = ray_start_regular["address"]

    @ray.remote
    def dep(sema, x=None):
        ray.get(sema.acquire.remote())
        return

    @ray.remote(num_gpus=1)
    def impossible():
        pass

    # Filter out actor handle refs.
    def filtered_summary():
        data = "\n".join(
            [
                line
                for line in memory_summary(address, line_wrap=False).split("\n")
                if "ACTOR_HANDLE" not in line
            ]
        )
        print(data)
        return data

    sema = Semaphore.remote(value=0)
    x = dep.remote(sema)
    y = dep.remote(sema, x=x)
    im = impossible.remote()  # noqa
    # x and its semaphore task are scheduled. im cannot
    # be scheduled, so it is pending forever.
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 2)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 1)
    wait_for_condition(lambda: count(filtered_summary(), SCHEDULED) == 1)

    z = dep.remote(sema, x=x)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 2)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 2)
    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 0)

    sema.release.remote()
    time.sleep(2)
    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 1)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 0)
    # y, z, and two semaphore tasks are scheduled.
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 4)

    sema.release.remote()
    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 2)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 0)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 2)

    sema.release.remote()
    ray.get(y)
    ray.get(z)
    wait_for_condition(lambda: count(filtered_summary(), FINISHED) == 3)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_DEPENDENCIES) == 0)
    wait_for_condition(lambda: count(filtered_summary(), WAITING_FOR_EXECUTION) == 0)
    wait_for_condition(lambda: count(filtered_summary(), SCHEDULED) == 1)


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