import os
import sys
import time

import pytest

import ray
import ray.cluster_utils
from ray._common.test_utils import wait_for_condition
from ray._private.runtime_env.context import RuntimeEnvContext
from ray._private.runtime_env.plugin import RuntimeEnvPlugin
from ray._private.test_utils import (
    get_other_nodes,
    is_placement_group_removed,
    placement_group_assert_no_leak,
)
from ray._raylet import PlacementGroupID
from ray.util.placement_group import PlacementGroup
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

MOCK_WORKER_STARTUP_SLOWLY_PLUGIN_CLASS_PATH = (
    "ray.tests.test_placement_group_4.MockWorkerStartupSlowlyPlugin"  # noqa
)
MOCK_WORKER_STARTUP_SLOWLY_PLUGIN_NAME = "MockWorkerStartupSlowlyPlugin"


class MockWorkerStartupSlowlyPlugin(RuntimeEnvPlugin):

    name = MOCK_WORKER_STARTUP_SLOWLY_PLUGIN_NAME

    def validate(runtime_env_dict: dict) -> str:
        return "success"

    @staticmethod
    def create(uri: str, runtime_env_dict: dict, ctx: RuntimeEnvContext) -> float:
        time.sleep(60)
        return 0


def test_remove_placement_group(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=4)
    ray.init(address=cluster.address)

    @ray.remote
    def warmup():
        pass

    # warm up the cluster.
    ray.get([warmup.remote() for _ in range(4)])

    # First try to remove a placement group that doesn't
    # exist. This should not do anything.
    random_group_id = PlacementGroupID.from_random()
    random_placement_group = PlacementGroup(random_group_id)
    for _ in range(3):
        ray.util.remove_placement_group(random_placement_group)

    # Creating a placement group as soon as it is
    # created should work.
    placement_group = ray.util.placement_group([{"CPU": 2}, {"CPU": 2}])
    assert placement_group.wait(10)

    ray.util.remove_placement_group(placement_group)

    wait_for_condition(lambda: is_placement_group_removed(placement_group))

    # # Now let's create a placement group.
    placement_group = ray.util.placement_group([{"CPU": 2}, {"CPU": 2}])
    assert placement_group.wait(10)

    # Create an actor that occupies resources.
    @ray.remote(num_cpus=2)
    class A:
        def f(self):
            return 3

    # Currently, there's no way to prevent
    # tasks to be retried for removed placement group.
    # Set max_retries=0 for testing.
    # TODO(sang): Handle this edge case.
    @ray.remote(num_cpus=2, max_retries=0)
    def long_running_task():
        print(os.getpid())
        time.sleep(50)

    # Schedule a long running task and actor.
    task_ref = long_running_task.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group
        )
    ).remote()
    a = A.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group
        )
    ).remote()
    assert ray.get(a.f.remote()) == 3

    ray.util.remove_placement_group(placement_group)
    # Subsequent remove request shouldn't do anything.
    for _ in range(3):
        ray.util.remove_placement_group(placement_group)

    # Make sure placement group resources are
    # released and we can schedule this task.
    @ray.remote(num_cpus=4)
    def f():
        return 3

    assert ray.get(f.remote()) == 3
    # Since the placement group is removed,
    # the actor should've been killed.
    # That means this request should fail.
    with pytest.raises(ray.exceptions.RayActorError, match="actor died"):
        ray.get(a.f.remote(), timeout=3.0)
    with pytest.raises(ray.exceptions.WorkerCrashedError):
        ray.get(task_ref)


@pytest.mark.parametrize(
    "set_runtime_env_plugins",
    [
        '[{"class":"' + MOCK_WORKER_STARTUP_SLOWLY_PLUGIN_CLASS_PATH + '"}]',
    ],
    indirect=True,
)
def test_remove_placement_group_worker_startup_slowly(
    set_runtime_env_plugins, ray_start_cluster
):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=4)
    ray.init(address=cluster.address)

    placement_group = ray.util.placement_group([{"CPU": 2}, {"CPU": 2}])
    assert placement_group.wait(10)

    @ray.remote(num_cpus=2)
    class A:
        def ready(self):
            return "ok"

        def hang(self):
            time.sleep(60)

    @ray.remote(num_cpus=2, max_retries=0)
    def long_running_task():
        time.sleep(60)

    # Schedule a long-running task that uses
    # runtime env to mock worker start up slowly.
    task_ref = long_running_task.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group
        ),
        runtime_env={MOCK_WORKER_STARTUP_SLOWLY_PLUGIN_NAME: {}},
    ).remote()
    a = A.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group
        )
    ).remote()
    assert ray.get(a.ready.remote()) == "ok"

    # Remove the PG, check that the actor and task are failed.
    ray.util.remove_placement_group(placement_group)

    with pytest.raises(ray.exceptions.RayActorError, match="actor died"):
        ray.get(a.hang.remote(), timeout=10)

    # The long-running task should still be in the state
    # of leasing-worker bacause of the worker startup delay.
    with pytest.raises(ray.exceptions.TaskPlacementGroupRemoved):
        ray.get(task_ref)


def test_remove_pending_placement_group(ray_start_cluster):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=4)
    ray.init(address=cluster.address)

    # Create a placement group that cannot be scheduled now.
    placement_group = ray.util.placement_group([{"GPU": 2}, {"CPU": 2}])
    wait_for_condition(
        lambda: (ray.util.placement_group_table(placement_group) or {}).get("state")
        == "PENDING"
    )
    ray.util.remove_placement_group(placement_group)
    wait_for_condition(lambda: is_placement_group_removed(placement_group))

    @ray.remote(num_cpus=4)
    def f():
        return 3

    # Make sure this task is still schedulable.
    assert ray.get(f.remote()) == 3

    placement_group_assert_no_leak([placement_group])


def test_placement_group_table(ray_start_cluster):
    @ray.remote(num_cpus=2)
    class Actor(object):
        def __init__(self):
            self.n = 0

        def value(self):
            return self.n

    cluster = ray_start_cluster
    num_nodes = 2
    for _ in range(num_nodes):
        cluster.add_node(num_cpus=4)
    ray.init(address=cluster.address)
    pgs_created = []

    # Originally placement group creation should be pending because
    # there are no resources.
    name = "name"
    strategy = "PACK"
    bundles = [{"CPU": 2, "GPU": 1}, {"CPU": 2}]
    placement_group = ray.util.placement_group(
        name=name, strategy=strategy, bundles=bundles
    )
    pgs_created.append(placement_group)
    result = ray.util.placement_group_table(placement_group)
    assert result["name"] == name
    assert result["strategy"] == strategy
    for i in range(len(bundles)):
        assert bundles[i] == result["bundles"][i]
    assert result["state"] == "PENDING"

    # Now the placement group should be scheduled.
    cluster.add_node(num_cpus=5, num_gpus=1)
    cluster.wait_for_nodes()

    actor_1 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=0
        )
    ).remote()
    ray.get(actor_1.value.remote())

    result = ray.util.placement_group_table(placement_group)
    assert result["state"] == "CREATED"

    # Add tow more placement group for placement group table test.
    second_strategy = "SPREAD"
    pgs_created.append(
        ray.util.placement_group(
            name="second_placement_group", strategy=second_strategy, bundles=bundles
        )
    )
    pgs_created.append(
        ray.util.placement_group(
            name="third_placement_group", strategy=second_strategy, bundles=bundles
        )
    )

    placement_group_table = ray.util.placement_group_table()
    assert len(placement_group_table) == 3

    true_name_set = {"name", "second_placement_group", "third_placement_group"}
    get_name_set = set()

    for _, placement_group_data in placement_group_table.items():
        get_name_set.add(placement_group_data["name"])

    assert true_name_set == get_name_set

    placement_group_assert_no_leak(pgs_created)


def test_placement_group_stats(ray_start_cluster):
    cluster = ray_start_cluster
    num_nodes = 1
    for _ in range(num_nodes):
        cluster.add_node(num_cpus=4, num_gpus=1)
    ray.init(address=cluster.address)

    # Test createable pgs.
    pg = ray.util.placement_group(bundles=[{"CPU": 4, "GPU": 1}])
    ray.get(pg.ready())
    stats = ray.util.placement_group_table(pg)["stats"]
    assert stats["scheduling_attempt"] == 1
    assert stats["scheduling_state"] == "FINISHED"
    assert stats["end_to_end_creation_latency_ms"] != 0

    # Create a pending pg.
    pg2 = ray.util.placement_group(bundles=[{"CPU": 4, "GPU": 1}])

    def assert_scheduling_state():
        stats = ray.util.placement_group_table(pg2)["stats"]
        if stats["scheduling_attempt"] != 1:
            return False
        if stats["scheduling_state"] != "NO_RESOURCES":
            return False
        if stats["end_to_end_creation_latency_ms"] != 0:
            return False
        return True

    wait_for_condition(assert_scheduling_state)

    # Remove the first pg, and the second
    # pg should be schedulable now.
    ray.util.remove_placement_group(pg)

    def assert_scheduling_state():
        stats = ray.util.placement_group_table(pg2)["stats"]
        if stats["scheduling_state"] != "FINISHED":
            return False
        if stats["end_to_end_creation_latency_ms"] == 0:
            return False
        return True

    wait_for_condition(assert_scheduling_state)

    # Infeasible pg.
    pg3 = ray.util.placement_group(bundles=[{"CPU": 4, "a": 1}])
    # TODO This is supposed to be infeasible, but it is printed
    # as NO_RESOURCES. Fix the issue.
    # def assert_scheduling_state():
    #     stats = ray.util.placement_group_table(pg3)["stats"]
    #     print(stats)
    #     if stats["scheduling_state"] != "INFEASIBLE":
    #         return False
    #     return True
    # wait_for_condition(assert_scheduling_state)

    ray.util.remove_placement_group(pg3)

    def assert_scheduling_state():
        stats = ray.util.placement_group_table(pg3)["stats"]
        if stats["scheduling_state"] != "REMOVED":
            return False
        return True

    wait_for_condition(assert_scheduling_state)

    placement_group_assert_no_leak([pg2])


def test_cuda_visible_devices(ray_start_cluster):
    @ray.remote(num_gpus=1)
    def f():
        return os.environ["CUDA_VISIBLE_DEVICES"]

    cluster = ray_start_cluster
    num_nodes = 1
    for _ in range(num_nodes):
        cluster.add_node(num_gpus=1)
    ray.init(address=cluster.address)

    g1 = ray.util.placement_group([{"CPU": 1, "GPU": 1}])
    o1 = f.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=g1)
    ).remote()

    devices = ray.get(o1)
    assert devices == "0", devices
    placement_group_assert_no_leak([g1])


def test_placement_group_reschedule_when_node_dead(ray_start_cluster):
    @ray.remote(num_cpus=1)
    class Actor(object):
        def __init__(self):
            self.n = 0

        def value(self):
            return self.n

    cluster = ray_start_cluster
    cluster.add_node(num_cpus=4)
    cluster.add_node(num_cpus=4)
    cluster.add_node(num_cpus=4)
    cluster.wait_for_nodes()
    ray.init(address=cluster.address, namespace="default_test_namespace")

    # Make sure both head and worker node are alive.
    nodes = ray.nodes()
    assert len(nodes) == 3
    assert nodes[0]["alive"] and nodes[1]["alive"] and nodes[2]["alive"]

    placement_group = ray.util.placement_group(
        name="name", strategy="SPREAD", bundles=[{"CPU": 2}, {"CPU": 2}, {"CPU": 2}]
    )
    actor_1 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=0
        ),
        lifetime="detached",
    ).remote()
    actor_2 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=1
        ),
        lifetime="detached",
    ).remote()
    actor_3 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=2
        ),
        lifetime="detached",
    ).remote()
    ray.get(actor_1.value.remote())
    ray.get(actor_2.value.remote())
    ray.get(actor_3.value.remote())

    cluster.remove_node(get_other_nodes(cluster, exclude_head=True)[-1])
    cluster.wait_for_nodes()

    actor_4 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=0
        ),
        lifetime="detached",
    ).remote()
    actor_5 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=1
        ),
        lifetime="detached",
    ).remote()
    actor_6 = Actor.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=placement_group, placement_group_bundle_index=2
        ),
        lifetime="detached",
    ).remote()
    ray.get(actor_4.value.remote())
    ray.get(actor_5.value.remote())
    ray.get(actor_6.value.remote())
    placement_group_assert_no_leak([placement_group])


def test_infeasible_pg(ray_start_cluster):
    """Test infeasible pgs are scheduled after new nodes are added."""
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=2)
    ray.init("auto")

    bundle = {"CPU": 4, "GPU": 1}
    pg = ray.util.placement_group([bundle], name="worker_1", strategy="STRICT_PACK")

    # Placement group is infeasible.
    with pytest.raises(ray.exceptions.GetTimeoutError):
        ray.get(pg.ready(), timeout=3)

    state = ray.util.placement_group_table()[pg.id.hex()]["stats"]["scheduling_state"]
    assert state == "INFEASIBLE"

    # Add a new node. PG can now be scheduled.
    cluster.add_node(num_cpus=4, num_gpus=1)
    assert ray.get(pg.ready(), timeout=10)


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